cmd/catalyst/gcs.go

package main

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"strings"

	"golang.org/x/oauth2/google"
	"google.golang.org/api/option"

	"cloud.google.com/go/storage"
)

var (
	envVarToPathMapping = map[string]string{
		"BOOTLOADER_PATH":       "bootloader.img",
		"CATALYST_ZIRCONA_PATH": "zedboot.zbi",
		"CATALYST_VBMETAA_PATH": "zedboot.vbmeta",
	}
)

func createClient(ctx context.Context, credentialsPath string) (*storage.Client, error) {
	var options []option.ClientOption
	if credentialsPath != "" {
		contents, err := ioutil.ReadFile(credentialsPath)
		if err != nil {
			return nil, err
		}
		creds, err := google.CredentialsFromJSON(ctx, contents, storage.ScopeReadOnly)
		if err != nil {
			return nil, err
		}
		options = append(options, option.WithCredentials(creds))
	}
	client, err := storage.NewClient(ctx, options...)
	if err != nil {
		return nil, err
	}
	return client, nil
}

// DownloadImages is a helper that downloads all images from gcs.
// Specifically, it attempts to download images for the bootloader and
// zedboot. This allows us to put the device in 1 state prior to handing
// off to platform code.
// Returns a list of filenames downloaded on success.
// TODO(rudymathu): Deprecate this in favor of allowing platform to put devices in 1 state
func DownloadImages(ctx context.Context, credentialsPath string) ([]string, error) {
	client, err := createClient(ctx, credentialsPath)
	if err != nil {
		return nil, err
	}
	var downloadedFilenames []string
	for envVarName, localPath := range envVarToPathMapping {
		remotePath := os.Getenv(envVarName)
		if remotePath != "" {
			if err := downloadImage(ctx, client, remotePath, localPath); err != nil {
				return nil, err
			}
			downloadedFilenames = append(downloadedFilenames, localPath)
		}
	}
	return downloadedFilenames, nil
}

// downloadImage is a helper that downloads the specified image from gcs
func downloadImage(ctx context.Context, client *storage.Client, gcsPath, localPath string) error {
	imageURL, err := url.Parse(gcsPath)
	if err != nil {
		return err
	}
	if imageURL.Scheme != "gs" {
		return fmt.Errorf("%s is not a valid gcs URL", gcsPath)
	}
	bucket := imageURL.Host
	bkt := client.Bucket(bucket)

	// Get reader to GCS object.
	r, err := bkt.Object(strings.TrimLeft(imageURL.Path, "/")).NewReader(ctx)
	if err != nil {
		return err
	}
	defer r.Close()

	// Open a local file and download remote data into it.
	file, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err := io.Copy(file, r); err != nil {
		return err
	}
	return nil
}