diff options
| -rw-r--r-- | pkg/gcs/gcs.go | 107 | ||||
| -rw-r--r-- | syz-gce/syz-gce.go | 61 |
2 files changed, 120 insertions, 48 deletions
diff --git a/pkg/gcs/gcs.go b/pkg/gcs/gcs.go new file mode 100644 index 000000000..ac95ec420 --- /dev/null +++ b/pkg/gcs/gcs.go @@ -0,0 +1,107 @@ +// Copyright 2017 syzkaller project authors. All rights reserved. +// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. + +// Package gcs provides wrappers around Google Cloud Storage (GCS) APIs. +// Package uses Application Default Credentials assuming that the program runs on GCE. +// +// See the following links for details and API reference: +// https://cloud.google.com/go/getting-started/using-cloud-storage +// https://godoc.org/cloud.google.com/go/storage +package gcs + +import ( + "context" + "fmt" + "io" + "os" + "strings" + "time" + + "cloud.google.com/go/storage" +) + +type Client struct { + client *storage.Client + ctx context.Context +} + +type File struct { + Updated time.Time + + ctx context.Context + handle *storage.ObjectHandle +} + +func (file *File) Reader() (io.ReadCloser, error) { + return file.handle.NewReader(file.ctx) +} + +func NewClient() (*Client, error) { + ctx := context.Background() + storageClient, err := storage.NewClient(ctx) + if err != nil { + return nil, err + } + client := &Client{ + client: storageClient, + ctx: ctx, + } + return client, nil +} + +func (client *Client) Read(gcsFile string) (*File, error) { + bucket, filename, err := split(gcsFile) + if err != nil { + return nil, err + } + bkt := client.client.Bucket(bucket) + f := bkt.Object(filename) + attrs, err := f.Attrs(client.ctx) + if err != nil { + return nil, fmt.Errorf("failed to read %v attributes: %v", gcsFile, err) + } + if !attrs.Deleted.IsZero() { + return nil, fmt.Errorf("file %v is deleted", gcsFile) + } + handle := f.If(storage.Conditions{ + GenerationMatch: attrs.Generation, + MetagenerationMatch: attrs.Metageneration, + }) + file := &File{ + Updated: attrs.Updated, + ctx: client.ctx, + handle: handle, + } + return file, nil +} + +func (client *Client) UploadFile(localFile, gcsFile string) error { + bucket, filename, err := split(gcsFile) + if err != nil { + return err + } + local, err := os.Open(localFile) + if err != nil { + return err + } + defer local.Close() + bkt := client.client.Bucket(bucket) + f := bkt.Object(filename) + w := f.NewWriter(client.ctx) + defer w.Close() + if _, err := io.Copy(w, local); err != nil { + return err + } + return nil +} + +func split(file string) (bucket, filename string, err error) { + if strings.HasPrefix(file, "gs://") { + file = file[5:] + } + pos := strings.IndexByte(file, '/') + if pos == -1 { + return "", "", fmt.Errorf("invalid GCS file name: %v", file) + } + return file[:pos], file[pos+1:], nil +} diff --git a/syz-gce/syz-gce.go b/syz-gce/syz-gce.go index ddc410536..e900ea941 100644 --- a/syz-gce/syz-gce.go +++ b/syz-gce/syz-gce.go @@ -32,27 +32,24 @@ import ( "path/filepath" "runtime" "strconv" - "strings" "sync/atomic" "syscall" "time" - "cloud.google.com/go/storage" "github.com/google/syzkaller/dashboard" "github.com/google/syzkaller/gce" . "github.com/google/syzkaller/log" pkgconfig "github.com/google/syzkaller/pkg/config" + "github.com/google/syzkaller/pkg/gcs" "github.com/google/syzkaller/pkg/git" "github.com/google/syzkaller/syz-manager/config" - "golang.org/x/net/context" ) var ( flagConfig = flag.String("config", "", "config file") - cfg *Config - ctx context.Context - storageClient *storage.Client + cfg *Config + GCS *gcs.Client GCE *gce.Context managerHttpPort uint32 patchesHash string @@ -107,9 +104,7 @@ func main() { gopath := abs(wd, "gopath") os.Setenv("GOPATH", gopath) - ctx = context.Background() - storageClient, err = storage.NewClient(ctx) - if err != nil { + if GCS, err = gcs.NewClient(); err != nil { Fatalf("failed to create cloud storage client: %v", err) } @@ -433,7 +428,7 @@ type GCSImageAction struct { ImagePath string ImageName string - handle *storage.ObjectHandle + file *gcs.File } func (a *GCSImageAction) Name() string { @@ -441,24 +436,12 @@ func (a *GCSImageAction) Name() string { } func (a *GCSImageAction) Poll() (string, error) { - pos := strings.IndexByte(a.ImageArchive, '/') - if pos == -1 { - return "", fmt.Errorf("invalid GCS file name: %v", a.ImageArchive) - } - bkt := storageClient.Bucket(a.ImageArchive[:pos]) - f := bkt.Object(a.ImageArchive[pos+1:]) - attrs, err := f.Attrs(ctx) + f, err := GCS.Read(a.ImageArchive) if err != nil { - return "", fmt.Errorf("failed to read %v attributes: %v", a.ImageArchive, err) - } - if !attrs.Deleted.IsZero() { - return "", fmt.Errorf("file %v is deleted", a.ImageArchive) + return "", err } - a.handle = f.If(storage.Conditions{ - GenerationMatch: attrs.Generation, - MetagenerationMatch: attrs.Metageneration, - }) - return attrs.Updated.Format(time.RFC1123Z), nil + a.file = f + return f.Updated.Format(time.RFC1123Z), nil } func (a *GCSImageAction) Build() error { @@ -466,7 +449,7 @@ func (a *GCSImageAction) Build() error { if err := os.RemoveAll("image"); err != nil { return fmt.Errorf("failed to remove image dir: %v", err) } - if err := downloadAndExtract(a.handle, "image"); err != nil { + if err := downloadAndExtract(a.file, "image"); err != nil { return fmt.Errorf("failed to download and extract %v: %v", a.ImageArchive, err) } if err := createImage("image/disk.tar.gz", a.ImagePath, a.ImageName); err != nil { @@ -529,8 +512,8 @@ func chooseUnusedPort() (int, error) { return port, nil } -func downloadAndExtract(f *storage.ObjectHandle, dir string) error { - r, err := f.NewReader(ctx) +func downloadAndExtract(f *gcs.File, dir string) error { + r, err := f.Reader() if err != nil { return err } @@ -578,7 +561,7 @@ func downloadAndExtract(f *storage.ObjectHandle, dir string) error { func createImage(localFile, gcsFile, imageName string) error { Logf(0, "uploading image...") - if err := uploadFile(localFile, gcsFile); err != nil { + if err := GCS.UploadFile(localFile, gcsFile); err != nil { return fmt.Errorf("failed to upload image: %v", err) } Logf(0, "creating gce image...") @@ -591,24 +574,6 @@ func createImage(localFile, gcsFile, imageName string) error { return nil } -func uploadFile(localFile, gcsFile string) error { - local, err := os.Open(localFile) - if err != nil { - return err - } - defer local.Close() - pos := strings.IndexByte(gcsFile, '/') - if pos == -1 { - return fmt.Errorf("invalid GCS file name: %v", gcsFile) - } - bkt := storageClient.Bucket(gcsFile[:pos]) - f := bkt.Object(gcsFile[pos+1:]) - w := f.NewWriter(ctx) - defer w.Close() - io.Copy(w, local) - return nil -} - func buildKernel(dir, ccompiler string) error { os.Remove(filepath.Join(dir, ".config")) if _, err := runCmd(dir, "make", "defconfig"); err != nil { |