vm/gce: accept local image

Currently gce accepts precreated GCE image name as image config param,
while all other VM types accept local file path as image.
This makes it impossible to write generic code that works with all VM types,
i.e. after building a new image it's unclear if it needs to be uploaded
to GCE or not, and what needs to be passed as image in config.

Eliminate this difference by making gce accept local image file as well.

3 files changed, 43 insertions, 38 deletions

diff --git a/pkg/gcs/gcs.go b/pkg/gcs/gcs.go
index ac95ec420..0ac9fc455 100644
--- a/pkg/gcs/gcs.go
+++ b/pkg/gcs/gcs.go
@@ -49,6 +49,10 @@ func NewClient() (*Client, error) {
 	return client, nil
 }
 
+func (client *Client) Close() {
+	client.client.Close()
+}
+
 func (client *Client) Read(gcsFile string) (*File, error) {
 	bucket, filename, err := split(gcsFile)
 	if err != nil {
diff --git a/syz-gce/syz-gce.go b/syz-gce/syz-gce.go
index 0fc39f650..e979385ba 100644
--- a/syz-gce/syz-gce.go
+++ b/syz-gce/syz-gce.go
@@ -62,7 +62,6 @@ type Config struct {
 	Hub_Key               string
 	Image_Archive         string
 	Image_Path            string
-	Image_Name            string
 	Http_Port             int
 	Machine_Type          string
 	Machine_Count         int
@@ -138,14 +137,10 @@ func main() {
 			Branch:       cfg.Linux_Branch,
 			Compiler:     cfg.Linux_Compiler,
 			UserspaceDir: abs(wd, cfg.Linux_Userspace),
-			ImagePath:    cfg.Image_Path,
-			ImageName:    cfg.Image_Name,
 		})
 	} else {
 		actions = append(actions, &GCSImageAction{
 			ImageArchive: cfg.Image_Archive,
-			ImagePath:    cfg.Image_Path,
-			ImageName:    cfg.Image_Name,
 		})
 	}
 	currHashes := make(map[string]string)
@@ -332,8 +327,6 @@ type LocalBuildAction struct {
 	Branch       string
 	Compiler     string
 	UserspaceDir string
-	ImagePath    string
-	ImageName    string
 }
 
 func (a *LocalBuildAction) Name() string {
@@ -389,8 +382,8 @@ func (a *LocalBuildAction) Build() error {
 	if err := os.Rename(vmlinux, "image/obj/vmlinux"); err != nil {
 		return fmt.Errorf("failed to rename vmlinux file: %v", err)
 	}
-	if err := createImage(filepath.Join(a.Dir, "disk.tar.gz"), a.ImagePath, a.ImageName); err != nil {
-		return err
+	if err := os.Rename(filepath.Join(a.Dir, "disk.tar.gz"), "image/disk.tar.gz"); err != nil {
+		return fmt.Errorf("failed to rename vmlinux file: %v", err)
 	}
 	return nil
 }
@@ -425,8 +418,6 @@ func (a *LocalBuildAction) apply(p dashboard.Patch) error {
 
 type GCSImageAction struct {
 	ImageArchive string
-	ImagePath    string
-	ImageName    string
 
 	file *gcs.File
 }
@@ -452,9 +443,6 @@ func (a *GCSImageAction) Build() error {
 	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 {
-		return err
-	}
 	return nil
 }
 
@@ -483,7 +471,7 @@ func writeManagerConfig(cfg *Config, httpPort int, file string) error {
 		Tag:              string(tag),
 		Syzkaller:        "gopath/src/github.com/google/syzkaller",
 		Type:             "gce",
-		Image:            cfg.Image_Name,
+		Image:            "image/disk.tar.gz",
 		Sandbox:          cfg.Sandbox,
 		Procs:            cfg.Procs,
 		Enable_Syscalls:  cfg.Enable_Syscalls,
@@ -560,21 +548,6 @@ func downloadAndExtract(f *gcs.File, dir string) error {
 	return nil
 }
 
-func createImage(localFile, gcsFile, imageName string) error {
-	Logf(0, "uploading image...")
-	if err := GCS.UploadFile(localFile, gcsFile); err != nil {
-		return fmt.Errorf("failed to upload image: %v", err)
-	}
-	Logf(0, "creating gce image...")
-	if err := GCE.DeleteImage(imageName); err != nil {
-		return fmt.Errorf("failed to delete GCE image: %v", err)
-	}
-	if err := GCE.CreateImage(imageName, gcsFile); err != nil {
-		return fmt.Errorf("failed to create GCE image: %v", err)
-	}
-	return nil
-}
-
 func buildKernel(dir, ccompiler string) error {
 	os.Remove(filepath.Join(dir, ".config"))
 	if _, err := runCmd(dir, "make", "defconfig"); err != nil {
diff --git a/vm/gce/gce.go b/vm/gce/gce.go
index afb393fd4..d0a181250 100644
--- a/vm/gce/gce.go
+++ b/vm/gce/gce.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/google/syzkaller/pkg/config"
 	"github.com/google/syzkaller/pkg/gce"
+	"github.com/google/syzkaller/pkg/gcs"
 	. "github.com/google/syzkaller/pkg/log"
 	"github.com/google/syzkaller/pkg/osutil"
 	"github.com/google/syzkaller/vm/vmimpl"
@@ -32,13 +33,15 @@ func init() {
 type Config struct {
 	Count        int    // number of VMs to use
 	Machine_Type string // GCE machine type (e.g. "n1-highcpu-2")
+	GCS_Path     string // GCS path to upload image
 	Sshkey       string // root ssh key for the image
 }
 
 type Pool struct {
-	env *vmimpl.Env
-	cfg *Config
-	GCE *gce.Context
+	env      *vmimpl.Env
+	cfg      *Config
+	GCE      *gce.Context
+	gceImage string
 }
 
 type instance struct {
@@ -58,6 +61,9 @@ func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
 	if env.Name == "" {
 		return nil, fmt.Errorf("config param name is empty (required for GCE)")
 	}
+	if env.Image == "" {
+		return nil, fmt.Errorf("config param image is empty (required for GCE)")
+	}
 	cfg := &Config{
 		Count: 1,
 	}
@@ -73,18 +79,40 @@ func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
 	if cfg.Machine_Type == "" {
 		return nil, fmt.Errorf("machine_type parameter is empty")
 	}
+	if cfg.GCS_Path == "" {
+		return nil, fmt.Errorf("gcs_path parameter is empty")
+	}
 	cfg.Sshkey = osutil.Abs(cfg.Sshkey)
 
 	GCE, err := gce.NewContext()
 	if err != nil {
 		return nil, fmt.Errorf("failed to init gce: %v", err)
 	}
-	Logf(0, "gce initialized: running on %v, internal IP %v, project %v, zone %v",
+	Logf(0, "GCE initialized: running on %v, internal IP %v, project %v, zone %v",
 		GCE.Instance, GCE.InternalIP, GCE.ProjectID, GCE.ZoneID)
+	GCS, err := gcs.NewClient()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create GCS client: %v", err)
+	}
+	defer GCS.Close()
+	gcsImage := filepath.Join(cfg.GCS_Path, env.Name+"-image.tar.gz")
+	gceImage := env.Name
+	Logf(0, "uploading image to %v...", gcsImage)
+	if err := GCS.UploadFile(env.Image, gcsImage); err != nil {
+		return nil, fmt.Errorf("failed to upload image: %v", err)
+	}
+	Logf(0, "creating GCE image...")
+	if err := GCE.DeleteImage(gceImage); err != nil {
+		return nil, fmt.Errorf("failed to delete GCE image: %v", err)
+	}
+	if err := GCE.CreateImage(gceImage, gcsImage); err != nil {
+		return nil, fmt.Errorf("failed to create GCE image: %v", err)
+	}
 	pool := &Pool{
-		cfg: cfg,
-		env: env,
-		GCE: GCE,
+		cfg:      cfg,
+		env:      env,
+		GCE:      GCE,
+		gceImage: gceImage,
 	}
 	return pool, nil
 }
@@ -111,7 +139,7 @@ func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
 		return nil, err
 	}
 	Logf(0, "creating instance: %v", name)
-	ip, err := pool.GCE.CreateInstance(name, pool.cfg.Machine_Type, pool.env.Image, string(gceKeyPub))
+	ip, err := pool.GCE.CreateInstance(name, pool.cfg.Machine_Type, pool.gceImage, string(gceKeyPub))
 	if err != nil {
 		return nil, err
 	}