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).

1 files changed, 10 insertions, 5 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)
 	}