src/pkg/cos/artifacts.go

package cos

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/storage"
	"cos.googlesource.com/cos/tools.git/src/pkg/gcs"
	"github.com/golang/glog"
	"github.com/pkg/errors"
)

const (
	cosToolsGCS      = "cos-tools"
	cosToolsGCSAsia  = "cos-tools-asia"
	cosToolsGCSEU    = "cos-tools-eu"
	kernelInfo       = "kernel_info"
	kernelSrcArchive = "kernel-src.tar.gz"
	kernelHeaders    = "kernel-headers.tgz"
	toolchainURL     = "toolchain_url"
	toolchainArchive = "toolchain.tar.xz"
	toolchainEnv     = "toolchain_env"
	crosKernelRepo   = "https://chromium.googlesource.com/chromiumos/third_party/kernel"
)

// Map VM zone prefix to specific cos-tools bucket for geo-redundancy.
var cosToolsPrefixMap = map[string]string{
	"us":           cosToolsGCS,
	"northamerica": cosToolsGCS,
	"southamerica": cosToolsGCS,
	"europe":       cosToolsGCSEU,
	"asia":         cosToolsGCSAsia,
	"australia":    cosToolsGCSAsia,
}

// ArtifactsDownloader defines the interface to download COS artifacts.
type ArtifactsDownloader interface {
	DownloadKernelSrc(ctx context.Context, destDir string) error
	DownloadToolchainEnv(ctx context.Context, destDir string) error
	DownloadToolchain(ctx context.Context, destDir string) error
	DownloadKernelHeaders(dctx context.Context, estDir string) error
	DownloadArtifact(ctx context.Context, destDir, artifact string) error
	GetArtifact(ctx context.Context, artifact string) ([]byte, error)
	ArtifactExists(ctx context.Context, artifact string) (bool, error)
	ListArtifacts(ctx context.Context, prefix string) ([]string, error)
}

// GCSDownloader is the struct downloading COS artifacts from GCS bucket.
type GCSDownloader struct {
	gcsClient         *storage.Client
	envReader         *EnvReader
	gcsDownloadBucket string
	gcsDownloadPrefix string
	mutex             sync.Mutex
}

// NewGCSDownloader creates a GCSDownloader instance.
func NewGCSDownloader(gcsClient *storage.Client, e *EnvReader, bucket, prefix string) *GCSDownloader {
	// If bucket is not set, use cos-tools, cos-tools-asia or cos-tools-eu
	// according to the zone the VM is running in for geo-redundancy.
	// If cannot fetch zone from metadata or get an unknown zone prefix,
	// use cos-tools as the default GCS bucket.
	if bucket == "" {
		zone, err := metadata.Zone()
		if err != nil {
			glog.Warningf("failed to get zone from metadata, will use 'gs://cos-tools' as artifact bucket, err: %v", err)
			bucket = cosToolsGCS
		} else {
			zonePrefix := strings.Split(zone, "-")[0]
			if geoBucket, found := cosToolsPrefixMap[zonePrefix]; found {
				bucket = geoBucket
			} else {
				bucket = cosToolsGCS
			}
		}
	}
	// Use {build number}/{board} as the default GCS download prefix.
	if prefix == "" {
		prefix = path.Join(e.BuildNumber(), e.Board())
	}
	return &GCSDownloader{gcsClient: gcsClient, envReader: e, gcsDownloadBucket: bucket, gcsDownloadPrefix: prefix}
}

// DownloadKernelSrc downloads COS kernel sources to destination directory.
func (d *GCSDownloader) DownloadKernelSrc(ctx context.Context, destDir string) error {
	return d.DownloadArtifact(ctx, destDir, kernelSrcArchive)
}

// DownloadToolchainEnv downloads toolchain compilation environment variables to destination directory.
func (d *GCSDownloader) DownloadToolchainEnv(ctx context.Context, destDir string) error {
	return d.DownloadArtifact(ctx, destDir, toolchainEnv)
}

// DownloadToolchain downloads toolchain package to destination directory.
func (d *GCSDownloader) DownloadToolchain(ctx context.Context, destDir string) error {
	return d.DownloadArtifact(ctx, destDir, toolchainArchive)
}

// DownloadKernelHeaders downloads COS kernel headers to destination directory.
func (d *GCSDownloader) DownloadKernelHeaders(ctx context.Context, destDir string) error {
	return d.DownloadArtifact(ctx, destDir, kernelHeaders)
}

// GetArtifact gets an artifact from GCS buckets and returns its content.
func (d *GCSDownloader) GetArtifact(ctx context.Context, artifactPath string) ([]byte, error) {
	tmpDir, err := ioutil.TempDir("", "tmp")
	if err != nil {
		return nil, errors.Wrap(err, "failed to create temp dir")
	}
	defer os.RemoveAll(tmpDir)

	if err = d.DownloadArtifact(ctx, tmpDir, artifactPath); err != nil {
		return nil, errors.Wrapf(err, "failed to download artifact %s", artifactPath)
	}

	content, err := ioutil.ReadFile(filepath.Join(tmpDir, filepath.Base(artifactPath)))
	if err != nil {
		return nil, errors.Wrapf(err, "failed to read file %s", filepath.Join(tmpDir, artifactPath))
	}

	return content, nil
}

// DownloadArtifact downloads an artifact from the GCS prefix configured in GCSDownloader.
func (d *GCSDownloader) DownloadArtifact(ctx context.Context, destDir, artifactPath string) error {
	gcsPath := path.Join(d.gcsDownloadPrefix, artifactPath)
	filename := filepath.Base(gcsPath)
	outputPath := filepath.Join(destDir, filename)
	gcsClient, err := d.getGCSClient(ctx)
	if err != nil {
		return err
	}
	glog.Infof("Start to download %s artifact from bucket: %s, object: %s to %s.", filename, d.gcsDownloadBucket, gcsPath, outputPath)
	if err := gcs.DownloadGCSObject(ctx, gcsClient, d.gcsDownloadBucket, gcsPath, outputPath); err != nil {
		glog.Errorf("Failed to download %s artifact from bucket: %s, object: %s to %s with error: %v.", filename, d.gcsDownloadBucket, gcsPath, outputPath, err)
		return fmt.Errorf("failed to download %s artifact from bucket: %s, object: %s to %s with error: %w", filename, d.gcsDownloadBucket, gcsPath, outputPath, err)
	}
	glog.Infof("Sucessfully downloaded %s artifact from bucket: %s, object: %s to %s.", filename, d.gcsDownloadBucket, gcsPath, outputPath)
	return nil

}

// ArtifactExists check whether the artifactpath exists.
func (d *GCSDownloader) ArtifactExists(ctx context.Context, artifactPath string) (bool, error) {
	gcsPath := path.Join(d.gcsDownloadPrefix, artifactPath)
	glog.Infof("Start to check whether artifact: %s exists in bucket: %s", gcsPath, d.gcsDownloadBucket)
	gcsClient, err := d.getGCSClient(ctx)
	if err != nil {
		return false, err
	}
	exists, error := gcs.GCSObjectExists(ctx, gcsClient, d.gcsDownloadBucket, gcsPath)
	if error != nil {
		glog.Errorf("Error appears when using gcs to check whether artifact: %s exists in bucket: %s - %v", gcsPath, d.gcsDownloadBucket, error)
		return false, fmt.Errorf("error appears when using gcs to check whether artifact: %s exists in bucket: %s - %w", gcsPath, d.gcsDownloadBucket, error)
	}
	return exists, nil
}

// ListArtifacts: this function lists the artifacts with `gcsDownloadPrefix/prefix` prefix.
func (d *GCSDownloader) ListArtifacts(ctx context.Context, prefix string) ([]string, error) {
	var objects []string
	var err error
	gcsPath := path.Join(d.gcsDownloadPrefix, prefix)
	gcsClient, err := d.getGCSClient(ctx)
	if err != nil {
		return nil, err
	}
	if objects, err = gcs.ListGCSBucket(ctx, gcsClient, d.gcsDownloadBucket, gcsPath); err != nil {
		glog.Errorf("Error appears when listing artifacts in bucket: %s with prefix: %s - %v", d.gcsDownloadBucket, gcsPath, err)
		return nil, fmt.Errorf("error appears when listing artifacts in bucket: %s with prefix: %s - %w", d.gcsDownloadBucket, gcsPath, err)
	}
	return objects, nil
}

// DownloadArtifactFromURL will download the artifact from url and save it to the destinationPath.
func (d *GCSDownloader) DownloadArtifactFromURL(ctx context.Context, url string, destinationPath string) error {
	gcsClient, err := d.getGCSClient(ctx)
	if err != nil {
		return err
	}
	return gcs.DownloadGCSObjectFromURL(ctx, gcsClient, url, destinationPath)
}

func (d *GCSDownloader) getGCSClient(ctx context.Context) (*storage.Client, error) {
	glog.Infof("Start to fetch the GCSClient.")
	d.mutex.Lock()
	defer d.mutex.Unlock()
	// Initialize the GCSClient only if it hasn't been created before.
	if d.gcsClient == nil {
		glog.Infof("GCSClient hasn't been created, start to create a new GCSClient.")
		var err error
		d.gcsClient, err = storage.NewClient(ctx)
		if err != nil {
			return nil, fmt.Errorf("can not create the GCSClient in GCSDownloader due to: %w", err)
		}
	}
	return d.gcsClient, nil
}