sweet/cmd/sweet/get.go

// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"archive/tar"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/benchmarks/sweet/cli/bootstrap"
	"golang.org/x/benchmarks/sweet/common"
	"golang.org/x/benchmarks/sweet/common/fileutil"
	"golang.org/x/benchmarks/sweet/common/log"
)

const (
	getUsage = `Retrieves assets for benchmarks from GCS.

Usage: %s get [flags]
`
)

func authOpts(includeNone bool) string {
	i := bootstrap.AuthOption(0)
	if !includeNone {
		i = 1
	}
	s := make([]string, 0, bootstrap.NumAuthOptions)
	for ; i < bootstrap.NumAuthOptions; i++ {
		s = append(s, i.String())
	}
	return strings.Join(s, ", ")
}

type getCmd struct {
	auth           bootstrap.AuthOption
	force          bool
	copyAssets     bool
	cache          string
	bucket         string
	assetsDir      string
	assetsHashFile string
	version        string
}

func (*getCmd) Name() string     { return "get" }
func (*getCmd) Synopsis() string { return "Retrieves assets for benchmarks." }
func (*getCmd) PrintUsage(w io.Writer, base string) {
	fmt.Fprintf(w, getUsage, base)
}

func (c *getCmd) SetFlags(f *flag.FlagSet) {
	f.Var(&c.auth, "auth", fmt.Sprintf("authentication method (options: %s)", authOpts(true)))
	f.BoolVar(&c.force, "force", false, "force download even if assets for this version exist in the cache")
	f.BoolVar(&c.copyAssets, "copy", false, "copy assets to assets-dir instead of symlinking")
	f.StringVar(&c.cache, "cache", bootstrap.CacheDefault(), "cache location for tar'd and compressed assets, if set to \"\" will ignore cache")
	f.StringVar(&c.version, "version", common.Version, "the version to download assets for")
	f.StringVar(&c.bucket, "bucket", "go-sweet-assets", "GCS bucket to download assets from")
	f.StringVar(&c.assetsDir, "assets-dir", "./assets", "location to extract assets into")
	f.StringVar(&c.assetsHashFile, "assets-hash-file", "./assets.hash", "file to check SHA256 hash of the downloaded artifact against")
}

func (c *getCmd) Run(_ []string) error {
	log.SetActivityLog(true)
	if err := bootstrap.ValidateVersion(c.version); err != nil {
		return err
	}
	installAssets := func(todir string, readonly bool) error {
		return downloadAndExtract(todir, c.bucket, c.assetsHashFile, c.version, c.auth, readonly)
	}
	if c.cache == "" {
		log.Printf("Skipping cache...")
		return installAssets(c.assetsDir, false)
	}
	log.Printf("Checking cache: %s", c.cache)
	t, err := bootstrap.CachedAssets(c.cache, c.version)
	if err == bootstrap.ErrNotInCache || (err == nil && c.force) {
		if err := installAssets(t, true); err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	if !c.copyAssets {
		log.Printf("Creating symlink to %s", c.assetsDir)
		if info, err := os.Lstat(c.assetsDir); err == nil {
			if info.Mode()&os.ModeSymlink != 0 {
				// We have a symlink, so just delete it so we can replace it.
				if err := os.Remove(c.assetsDir); err != nil {
					return fmt.Errorf("installing assets: removing %s: %v", c.assetsDir, err)
				}
			} else {
				return fmt.Errorf("installing assets: %s is not a symlink; to install assets here, remove it and re-run this command", c.assetsDir)
			}
		} else if !os.IsNotExist(err) {
			return fmt.Errorf("stat %s: %v", c.assetsDir, err)
		}
		return os.Symlink(t, c.assetsDir)
	}
	if _, err := os.Stat(c.assetsDir); err == nil {
		return fmt.Errorf("installing assets: %s exists; to copy assets here, remove it and re-run this command", c.assetsDir)
	} else if !os.IsNotExist(err) {
		return fmt.Errorf("stat %s: %v", c.assetsDir, err)
	}
	log.Printf("Copying assets %s", c.assetsDir)
	return fileutil.CopyDir(c.assetsDir, t)
}

func downloadAndExtract(todir, bucket, hashfile, version string, auth bootstrap.AuthOption, readonly bool) error {
	log.Printf("Downloading assets archive for version %s to %s", version, todir)

	// Create storage reader for streaming.
	rc, err := bootstrap.NewStorageReader(bucket, version, auth)
	if err != nil {
		return err
	}
	defer rc.Close()

	// Pass everything we read through a hash.
	hash := bootstrap.Hash()
	r := io.TeeReader(rc, hash)

	// Stream and extract the results.
	if err := extractAssets(r, todir, readonly); err != nil {
		return err
	}

	// Check the hash.
	return checkAssetsHash(bootstrap.CanonicalizeHash(hash), hashfile, version)
}

func checkAssetsHash(hash, hashfile, version string) error {
	vals, err := bootstrap.ReadHashesFile(hashfile)
	if err != nil {
		return err
	}
	check, ok := vals.Get(version)
	if !ok {
		return fmt.Errorf("hash for version %s not found", version)
	}
	if hash != check {
		return fmt.Errorf("downloaded artifact has unexpected hash: expected %s, got %s", hash, check)
	}
	return nil
}

func extractAssets(r io.Reader, outdir string, readonly bool) error {
	if err := os.MkdirAll(outdir, os.ModePerm); err != nil {
		return fmt.Errorf("create assets directory: %v", err)
	}
	gr, err := gzip.NewReader(r)
	if err != nil {
		return err
	}
	defer gr.Close()

	tr := tar.NewReader(gr)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		fullpath := filepath.Join(outdir, hdr.Name)
		if err := os.MkdirAll(filepath.Dir(fullpath), os.ModePerm); err != nil {
			return err
		}
		f, err := os.Create(fullpath)
		if err != nil {
			return err
		}
		if _, err := io.Copy(f, tr); err != nil {
			f.Close()
			return err
		}
		fperm := os.FileMode(uint32(hdr.Mode))
		if readonly {
			fperm = 0444 | (fperm & 0555)
		}
		if err := f.Chmod(fperm); err != nil {
			f.Close()
			return err
		}
		f.Close()
	}
	if readonly {
		return filepath.Walk(outdir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				return os.Chmod(path, 0555)
			}
			return nil
		})
	}
	return nil
}