From dacb3f1c44e9c527b8bb595e3b37a43ca7733cc0 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Mon, 11 Apr 2022 15:00:55 +0000
Subject: vm/gce: make zone_id configurable

At the moment syzkaller can only use the zone where it's running. Make it a
configurable option instead (with the old behavior as a fallback).
---
 pkg/gce/gce.go         | 15 ++++++++++-----
 syz-manager/manager.go |  2 +-
 vm/gce/gce.go          |  3 ++-
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/pkg/gce/gce.go b/pkg/gce/gce.go
index b0100b749..4e797ade6 100644
--- a/pkg/gce/gce.go
+++ b/pkg/gce/gce.go
@@ -43,7 +43,7 @@ type Context struct {
 	apiRateGate <-chan time.Time
 }
 
-func NewContext() (*Context, error) {
+func NewContext(customZoneID string) (*Context, error) {
 	ctx := &Context{
 		apiRateGate: time.NewTicker(time.Second).C,
 	}
@@ -64,18 +64,23 @@ func NewContext() (*Context, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to query gce project-id: %v", err)
 	}
-	ctx.ZoneID, err = ctx.getMeta("instance/zone")
+	myZoneID, err := ctx.getMeta("instance/zone")
 	if err != nil {
 		return nil, fmt.Errorf("failed to query gce zone: %v", err)
 	}
-	if i := strings.LastIndexByte(ctx.ZoneID, '/'); i != -1 {
-		ctx.ZoneID = ctx.ZoneID[i+1:] // the query returns some nonsense prefix
+	if i := strings.LastIndexByte(myZoneID, '/'); i != -1 {
+		myZoneID = myZoneID[i+1:] // the query returns some nonsense prefix
+	}
+	if customZoneID != "" {
+		ctx.ZoneID = customZoneID
+	} else {
+		ctx.ZoneID = myZoneID
 	}
 	ctx.Instance, err = ctx.getMeta("instance/name")
 	if err != nil {
 		return nil, fmt.Errorf("failed to query gce instance name: %v", err)
 	}
-	inst, err := ctx.computeService.Instances.Get(ctx.ProjectID, ctx.ZoneID, ctx.Instance).Do()
+	inst, err := ctx.computeService.Instances.Get(ctx.ProjectID, myZoneID, ctx.Instance).Do()
 	if err != nil {
 		return nil, fmt.Errorf("error getting instance info: %v", err)
 	}
diff --git a/syz-manager/manager.go b/syz-manager/manager.go
index 09762dd64..bb331b542 100644
--- a/syz-manager/manager.go
+++ b/syz-manager/manager.go
@@ -1331,7 +1331,7 @@ func publicWebAddr(addr string) string {
 		if host, err := os.Hostname(); err == nil {
 			addr = net.JoinHostPort(host, port)
 		}
-		if GCE, err := gce.NewContext(); err == nil {
+		if GCE, err := gce.NewContext(""); err == nil {
 			addr = net.JoinHostPort(GCE.ExternalIP, port)
 		}
 	}
diff --git a/vm/gce/gce.go b/vm/gce/gce.go
index 91f76ff68..ef213d12e 100644
--- a/vm/gce/gce.go
+++ b/vm/gce/gce.go
@@ -40,6 +40,7 @@ func init() {
 
 type Config struct {
 	Count         int    `json:"count"`          // number of VMs to use
+	ZoneID        string `json:"zone_id"`        // GCE zone (if it's different from that of syz-manager)
 	MachineType   string `json:"machine_type"`   // GCE machine type (e.g. "n1-highcpu-2")
 	GCSPath       string `json:"gcs_path"`       // GCS path to upload image
 	GCEImage      string `json:"gce_image"`      // pre-created GCE image to use
@@ -99,7 +100,7 @@ func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
 		return nil, fmt.Errorf("both image and gce_image are specified")
 	}
 
-	GCE, err := gce.NewContext()
+	GCE, err := gce.NewContext(cfg.ZoneID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to init gce: %v", err)
 	}
-- 
cgit mrf-deployment