tools/bootserver/image.go - fuchsia - Git at Google

 // Copyright 2019 The Fuchsia 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 bootserver

 import (
 	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net/url"
 	"os"
 	"path/filepath"
 	"strings"

 	"cloud.google.com/go/storage"
 	"go.fuchsia.dev/fuchsia/tools/build"
 	"go.fuchsia.dev/fuchsia/tools/lib/iomisc"
 	"go.fuchsia.dev/fuchsia/tools/lib/logger"
 )

 var (
 	noOpClose = func() error { return nil }
 )

 // Image is a fuchsia image as viewed by bootserver; a simplified version of build.Image.
 type Image struct {
 	// Name is an identifier for this image that usually derives from its target partition.
 	// TODO(fxbug.dev/38517): Remove when BootZedbootShim is deprecated.
 	Name string
 	// Path is the location of the image on disk.
 	// TODO(ihuh): Remove when bootserver/cmd/main.go no longer uses it.
 	Path string
 	// Reader is a reader to the image.
 	Reader io.ReaderAt
 	// Size is the size of the reader in bytes.
 	Size int64
 	// Args correspond to the bootserver args that map to this image type.
 	Args []string
 }

 func getImageArgs(img build.Image, bootMode Mode) []string {
 	switch bootMode {
 	case ModePave:
 		return img.PaveArgs
 	case ModeNetboot:
 		return img.NetbootArgs
 	case ModePaveZedboot:
 		return img.PaveZedbootArgs
 	}
 	return nil
 }

 // ConvertFromBuildImages filters and returns Images corresponding to build Images of a given bootMode.
 func ConvertFromBuildImages(buildImages []build.Image, bootMode Mode, imageDir string) ([]Image, func() error, error) {
 	var imgs []Image
 	closeFunc := noOpClose
 	for _, buildImg := range buildImages {
 		args := getImageArgs(buildImg, bootMode)
 		reader, err := os.Open(filepath.Join(imageDir, buildImg.Path))
 		if err != nil {
 			if os.IsNotExist(err) {
 				// Not all images exist so skip if it doesn't.
 				continue
 			}
 			// Close already opened readers.
 			closeImages(imgs)
 			return nil, closeFunc, err
 		}
 		fi, err := reader.Stat()
 		if err != nil {
 			closeImages(imgs)
 			return nil, closeFunc, err
 		}
 		imgs = append(imgs, Image{
 			Name:   buildImg.Type + "_" + buildImg.Name,
 			Path:   buildImg.Path,
 			Reader: reader,
 			Size:   fi.Size(),
 			Args:   args,
 		})
 	}
 	closeFunc = func() error {
 		return closeImages(imgs)
 	}
 	return imgs, closeFunc, nil
 }

 // ImagesFromLocalFS returns Images of a given bootMode that exist on the local
 // filesystem.
 func ImagesFromLocalFS(manifest string, bootMode Mode) ([]Image, func() error, error) {
 	buildImages, err := build.LoadImages(manifest)
 	if err != nil {
 		return nil, noOpClose, err
 	}
 	return ConvertFromBuildImages(buildImages, bootMode, filepath.Dir(manifest))
 }

 // GCSReader is a wrapper around storage.Reader which implements io.ReaderAt.
 // TODO(https://github.com/googleapis/google-cloud-go/issues/1686): Use the storage.Reader as is once this is fixed.
 type gcsReader struct {
 	obj    *storage.ObjectHandle
 	reader io.ReadCloser
 	index  int64
 }

 func getUncompressedReader(obj *storage.ObjectHandle) (io.ReadCloser, error) {
 	ctx := context.Background()
 	objAttrs, err := obj.Attrs(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get attrs for %q from GCS: %v", obj.ObjectName(), err)
 	}
 	if objAttrs.ContentEncoding != "gzip" {
 		return obj.NewReader(ctx)
 	}
 	r, err := obj.ReadCompressed(true).NewReader(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read %q from GCS: %v", obj.ObjectName(), err)
 	}
 	if r.Attrs.ContentEncoding != "gzip" {
 		return nil, fmt.Errorf("content-encoding expected: gzip, actual: %s", r.Attrs.ContentEncoding)
 	}

 	gr, err := gzip.NewReader(r)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get gzip reader: %v", err)
 	}
 	return struct {
 		io.Reader
 		io.Closer
 	}{gr, iomisc.MultiCloser(gr, r)}, nil
 }

 func (g *gcsReader) ReadAt(buf []byte, offset int64) (int, error) {
 	if g.reader == nil || offset < g.index {
 		g.Close()
 		r, err := getUncompressedReader(g.obj)
 		if err != nil {
 			return 0, err
 		}
 		g.reader = r
 		g.index = 0
 	}
 	// If the offset is greater than the index of the reader, we need to read
 	// up until the requested offset. These bytes will be ignored as they only
 	// need to be read to bring the index up to the offset so that the next
 	// calls to Read will read the correct bytes into buf.
 	if offset > g.index {
 		diff := make([]byte, offset-g.index)
 		n, err := io.ReadAtLeast(g.reader, diff, len(diff))
 		g.index += int64(n)
 		if err != nil && err != io.ErrUnexpectedEOF {
 			return 0, err
 		}
 	}

 	n, err := io.ReadAtLeast(g.reader, buf, len(buf))
 	g.index += int64(n)
 	if err == io.ErrUnexpectedEOF {
 		err = io.EOF
 	}
 	return n, err
 }

 func (g *gcsReader) Close() error {
 	if g.reader != nil {
 		return g.reader.Close()
 	}
 	return nil
 }

 // ImagesFromGCS returns Images of a given bootMode that exist in GCS. The image
 // paths provided in the manifest at the given url are expected to be relative paths
 // to the same directory of the manifest.
 func ImagesFromGCS(ctx context.Context, manifest *url.URL, bootMode Mode) ([]Image, func() error, error) {
 	closeFunc := noOpClose
 	bucket := manifest.Host
 	client, err := storage.NewClient(ctx)
 	if err != nil {
 		return nil, closeFunc, err
 	}
 	bkt := client.Bucket(bucket)
 	manifestGcsPath := strings.TrimLeft(manifest.Path, "/")
 	obj := bkt.Object(manifestGcsPath)
 	r, err := getUncompressedReader(obj)
 	if err != nil {
 		return nil, closeFunc, fmt.Errorf("failed to get image manifest from GCS: %v", err)
 	}
 	defer r.Close()
 	var buildImgs []build.Image
 	if err := json.NewDecoder(r).Decode(&buildImgs); err != nil {
 		return nil, closeFunc, fmt.Errorf("failed to decode image manifest: %v", err)
 	}
 	var imgs []Image
 	namespace := filepath.Dir(manifestGcsPath)
 	for _, buildImg := range buildImgs {
 		args := getImageArgs(buildImg, bootMode)
 		obj := bkt.Object(filepath.Join(namespace, buildImg.Path))
 		objAttrs, err := obj.Attrs(ctx)
 		if err == storage.ErrObjectNotExist {
 			// Not all images may have been uploaded so skip if it doesn't exist.
 			continue
 		} else if err != nil {
 			return nil, closeFunc, fmt.Errorf("failed to get object attributes: %v", err)
 		}

 		imgs = append(imgs, Image{
 			Name:   buildImg.Type + "_" + buildImg.Name,
 			Reader: &gcsReader{obj: obj},
 			Size:   objAttrs.Size,
 			Args:   args,
 		})
 	}

 	closeFunc = func() error {
 		return closeImages(imgs)
 	}
 	return imgs, closeFunc, nil
 }

 // GetImages parses the imageManifest and gets a list of images with readers to
 // each image. It returns the images as well as a func to close the image readers.
 func GetImages(ctx context.Context, imageManifest string, bootMode Mode) ([]Image, func() error, error) {
 	url, err := url.Parse(imageManifest)
 	closeFunc := noOpClose
 	if err != nil {
 		return nil, closeFunc, err
 	}
 	if url == nil {
 		return nil, closeFunc, fmt.Errorf("failed to parse %q", imageManifest)
 	}

 	var imgs []Image
 	if url.Scheme == "gs" {
 		logger.Debugf(ctx, "Fetching images from GCS")
 		imgs, closeFunc, err = ImagesFromGCS(ctx, url, bootMode)
 	} else if url.Scheme == "" && url.Host == "" {
 		logger.Debugf(ctx, "Fetching images from local disk")
 		// Assume that this is a filesystem path.
 		imgs, closeFunc, err = ImagesFromLocalFS(imageManifest, bootMode)
 	} else if url.Scheme != "" && url.Host != "" {
 		// TODO(ihuh|joshuaseaton): handle the case of getting images directly over HTTP
 		err = fmt.Errorf("unimplemented")
 	} else {
 		err = fmt.Errorf("unknown manifest reference: %q", imageManifest)
 	}
 	if err != nil {
 		logger.Errorf(ctx, "Failed to fetch images: %s", err)
 	} else {
 		logger.Errorf(ctx, "Completed fetching images")
 	}
 	return imgs, closeFunc, err
 }

 func closeImages(imgs []Image) error {
 	var errs []error
 	for _, img := range imgs {
 		if img.Reader != nil {
 			if closer, ok := img.Reader.(io.Closer); ok {
 				if err := closer.Close(); err != nil {
 					errs = append(errs, err)
 				}
 			}
 		}
 	}
 	if len(errs) > 0 {
 		return fmt.Errorf("failed to close images: %v", errs)
 	}
 	return nil
 }
	// Copyright 2019 The Fuchsia 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 bootserver

	import (
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"go.fuchsia.dev/fuchsia/tools/build"
	"go.fuchsia.dev/fuchsia/tools/lib/iomisc"
	"go.fuchsia.dev/fuchsia/tools/lib/logger"
	)

	var (
	noOpClose = func() error { return nil }
	)

	// Image is a fuchsia image as viewed by bootserver; a simplified version of build.Image.
	type Image struct {
	// Name is an identifier for this image that usually derives from its target partition.
	// TODO(fxbug.dev/38517): Remove when BootZedbootShim is deprecated.
	Name string
	// Path is the location of the image on disk.
	// TODO(ihuh): Remove when bootserver/cmd/main.go no longer uses it.
	Path string
	// Reader is a reader to the image.
	Reader io.ReaderAt
	// Size is the size of the reader in bytes.
	Size int64
	// Args correspond to the bootserver args that map to this image type.
	Args []string
	}

	func getImageArgs(img build.Image, bootMode Mode) []string {
	switch bootMode {
	case ModePave:
	return img.PaveArgs
	case ModeNetboot:
	return img.NetbootArgs
	case ModePaveZedboot:
	return img.PaveZedbootArgs
	}
	return nil
	}

	// ConvertFromBuildImages filters and returns Images corresponding to build Images of a given bootMode.
	func ConvertFromBuildImages(buildImages []build.Image, bootMode Mode, imageDir string) ([]Image, func() error, error) {
	var imgs []Image
	closeFunc := noOpClose
	for _, buildImg := range buildImages {
	args := getImageArgs(buildImg, bootMode)
	reader, err := os.Open(filepath.Join(imageDir, buildImg.Path))
	if err != nil {
	if os.IsNotExist(err) {
	// Not all images exist so skip if it doesn't.
	continue
	}
	// Close already opened readers.
	closeImages(imgs)
	return nil, closeFunc, err
	}
	fi, err := reader.Stat()
	if err != nil {
	closeImages(imgs)
	return nil, closeFunc, err
	}
	imgs = append(imgs, Image{
	Name: buildImg.Type + "_" + buildImg.Name,
	Path: buildImg.Path,
	Reader: reader,
	Size: fi.Size(),
	Args: args,
	})
	}
	closeFunc = func() error {
	return closeImages(imgs)
	}
	return imgs, closeFunc, nil
	}

	// ImagesFromLocalFS returns Images of a given bootMode that exist on the local
	// filesystem.
	func ImagesFromLocalFS(manifest string, bootMode Mode) ([]Image, func() error, error) {
	buildImages, err := build.LoadImages(manifest)
	if err != nil {
	return nil, noOpClose, err
	}
	return ConvertFromBuildImages(buildImages, bootMode, filepath.Dir(manifest))
	}

	// GCSReader is a wrapper around storage.Reader which implements io.ReaderAt.
	// TODO(https://github.com/googleapis/google-cloud-go/issues/1686): Use the storage.Reader as is once this is fixed.
	type gcsReader struct {
	obj *storage.ObjectHandle
	reader io.ReadCloser
	index int64
	}

	func getUncompressedReader(obj *storage.ObjectHandle) (io.ReadCloser, error) {
	ctx := context.Background()
	objAttrs, err := obj.Attrs(ctx)
	if err != nil {
	return nil, fmt.Errorf("failed to get attrs for %q from GCS: %v", obj.ObjectName(), err)
	}
	if objAttrs.ContentEncoding != "gzip" {
	return obj.NewReader(ctx)
	}
	r, err := obj.ReadCompressed(true).NewReader(ctx)
	if err != nil {
	return nil, fmt.Errorf("failed to read %q from GCS: %v", obj.ObjectName(), err)
	}
	if r.Attrs.ContentEncoding != "gzip" {
	return nil, fmt.Errorf("content-encoding expected: gzip, actual: %s", r.Attrs.ContentEncoding)
	}

	gr, err := gzip.NewReader(r)
	if err != nil {
	return nil, fmt.Errorf("failed to get gzip reader: %v", err)
	}
	return struct {
	io.Reader
	io.Closer
	}{gr, iomisc.MultiCloser(gr, r)}, nil
	}

	func (g *gcsReader) ReadAt(buf []byte, offset int64) (int, error) {
	if g.reader == nil \|\| offset < g.index {
	g.Close()
	r, err := getUncompressedReader(g.obj)
	if err != nil {
	return 0, err
	}
	g.reader = r
	g.index = 0
	}
	// If the offset is greater than the index of the reader, we need to read
	// up until the requested offset. These bytes will be ignored as they only
	// need to be read to bring the index up to the offset so that the next
	// calls to Read will read the correct bytes into buf.
	if offset > g.index {
	diff := make([]byte, offset-g.index)
	n, err := io.ReadAtLeast(g.reader, diff, len(diff))
	g.index += int64(n)
	if err != nil && err != io.ErrUnexpectedEOF {
	return 0, err
	}
	}

	n, err := io.ReadAtLeast(g.reader, buf, len(buf))
	g.index += int64(n)
	if err == io.ErrUnexpectedEOF {
	err = io.EOF
	}
	return n, err
	}

	func (g *gcsReader) Close() error {
	if g.reader != nil {
	return g.reader.Close()
	}
	return nil
	}

	// ImagesFromGCS returns Images of a given bootMode that exist in GCS. The image
	// paths provided in the manifest at the given url are expected to be relative paths
	// to the same directory of the manifest.
	func ImagesFromGCS(ctx context.Context, manifest *url.URL, bootMode Mode) ([]Image, func() error, error) {
	closeFunc := noOpClose
	bucket := manifest.Host
	client, err := storage.NewClient(ctx)
	if err != nil {
	return nil, closeFunc, err
	}
	bkt := client.Bucket(bucket)
	manifestGcsPath := strings.TrimLeft(manifest.Path, "/")
	obj := bkt.Object(manifestGcsPath)
	r, err := getUncompressedReader(obj)
	if err != nil {
	return nil, closeFunc, fmt.Errorf("failed to get image manifest from GCS: %v", err)
	}
	defer r.Close()
	var buildImgs []build.Image
	if err := json.NewDecoder(r).Decode(&buildImgs); err != nil {
	return nil, closeFunc, fmt.Errorf("failed to decode image manifest: %v", err)
	}
	var imgs []Image
	namespace := filepath.Dir(manifestGcsPath)
	for _, buildImg := range buildImgs {
	args := getImageArgs(buildImg, bootMode)
	obj := bkt.Object(filepath.Join(namespace, buildImg.Path))
	objAttrs, err := obj.Attrs(ctx)
	if err == storage.ErrObjectNotExist {
	// Not all images may have been uploaded so skip if it doesn't exist.
	continue
	} else if err != nil {
	return nil, closeFunc, fmt.Errorf("failed to get object attributes: %v", err)
	}

	imgs = append(imgs, Image{
	Name: buildImg.Type + "_" + buildImg.Name,
	Reader: &gcsReader{obj: obj},
	Size: objAttrs.Size,
	Args: args,
	})
	}

	closeFunc = func() error {
	return closeImages(imgs)
	}
	return imgs, closeFunc, nil
	}

	// GetImages parses the imageManifest and gets a list of images with readers to
	// each image. It returns the images as well as a func to close the image readers.
	func GetImages(ctx context.Context, imageManifest string, bootMode Mode) ([]Image, func() error, error) {
	url, err := url.Parse(imageManifest)
	closeFunc := noOpClose
	if err != nil {
	return nil, closeFunc, err
	}
	if url == nil {
	return nil, closeFunc, fmt.Errorf("failed to parse %q", imageManifest)
	}

	var imgs []Image
	if url.Scheme == "gs" {
	logger.Debugf(ctx, "Fetching images from GCS")
	imgs, closeFunc, err = ImagesFromGCS(ctx, url, bootMode)
	} else if url.Scheme == "" && url.Host == "" {
	logger.Debugf(ctx, "Fetching images from local disk")
	// Assume that this is a filesystem path.
	imgs, closeFunc, err = ImagesFromLocalFS(imageManifest, bootMode)
	} else if url.Scheme != "" && url.Host != "" {
	// TODO(ihuh\|joshuaseaton): handle the case of getting images directly over HTTP
	err = fmt.Errorf("unimplemented")
	} else {
	err = fmt.Errorf("unknown manifest reference: %q", imageManifest)
	}
	if err != nil {
	logger.Errorf(ctx, "Failed to fetch images: %s", err)
	} else {
	logger.Errorf(ctx, "Completed fetching images")
	}
	return imgs, closeFunc, err
	}

	func closeImages(imgs []Image) error {
	var errs []error
	for _, img := range imgs {
	if img.Reader != nil {
	if closer, ok := img.Reader.(io.Closer); ok {
	if err := closer.Close(); err != nil {
	errs = append(errs, err)
	}
	}
	}
	}
	if len(errs) > 0 {
	return fmt.Errorf("failed to close images: %v", errs)
	}
	return nil
	}