vm/gce: windows support

Support custom pre-created images.
Support non-root user.
Use dir instead of pwd on windows.
Don't use sudo on windows.

1 files changed, 43 insertions, 31 deletions

diff --git a/vm/gce/gce.go b/vm/gce/gce.go
index 084189522..bc1a7ac3c 100644
--- a/vm/gce/gce.go
+++ b/vm/gce/gce.go
@@ -39,16 +39,17 @@ 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
+	GCE_Image    string // Pre-created GCE image to use
 }
 
 type Pool struct {
-	env      *vmimpl.Env
-	cfg      *Config
-	GCE      *gce.Context
-	gceImage string
+	env *vmimpl.Env
+	cfg *Config
+	GCE *gce.Context
 }
 
 type instance struct {
+	env     *vmimpl.Env
 	cfg     *Config
 	GCE     *gce.Context
 	debug   bool
@@ -84,9 +85,12 @@ 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 == "" {
+	if cfg.GCE_Image == "" && cfg.GCS_Path == "" {
 		return nil, fmt.Errorf("gcs_path parameter is empty")
 	}
+	if cfg.GCE_Image != "" && env.Image != "" {
+		return nil, fmt.Errorf("both image and gce_image are specified")
+	}
 
 	GCE, err := gce.NewContext()
 	if err != nil {
@@ -95,24 +99,25 @@ func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
 	Logf(0, "GCE initialized: running on %v, internal IP %v, project %v, zone %v",
 		GCE.Instance, GCE.InternalIP, GCE.ProjectID, GCE.ZoneID)
 
-	gcsImage := filepath.Join(cfg.GCS_Path, env.Name+"-image.tar.gz")
-	Logf(0, "uploading image to %v...", gcsImage)
-	if err := uploadImageToGCS(env.Image, gcsImage); err != nil {
-		return nil, err
-	}
-	gceImage := env.Name
-	Logf(0, "creating GCE image %v...", gceImage)
-	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)
+	if cfg.GCE_Image == "" {
+		cfg.GCE_Image = env.Name
+		gcsImage := filepath.Join(cfg.GCS_Path, env.Name+"-image.tar.gz")
+		Logf(0, "uploading image to %v...", gcsImage)
+		if err := uploadImageToGCS(env.Image, gcsImage); err != nil {
+			return nil, err
+		}
+		Logf(0, "creating GCE image %v...", cfg.GCE_Image)
+		if err := GCE.DeleteImage(cfg.GCE_Image); err != nil {
+			return nil, fmt.Errorf("failed to delete GCE image: %v", err)
+		}
+		if err := GCE.CreateImage(cfg.GCE_Image, gcsImage); err != nil {
+			return nil, fmt.Errorf("failed to create GCE image: %v", err)
+		}
 	}
 	pool := &Pool{
-		cfg:      cfg,
-		env:      env,
-		GCE:      GCE,
-		gceImage: gceImage,
+		cfg: cfg,
+		env: env,
+		GCE: GCE,
 	}
 	return pool, nil
 }
@@ -139,7 +144,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.gceImage, string(gceKeyPub))
+	ip, err := pool.GCE.CreateInstance(name, pool.cfg.Machine_Type, pool.cfg.GCE_Image, string(gceKeyPub))
 	if err != nil {
 		return nil, err
 	}
@@ -150,19 +155,20 @@ func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
 			pool.GCE.DeleteInstance(name, true)
 		}
 	}()
-	sshKey := pool.env.Sshkey
-	sshUser := "root"
+	sshKey := pool.env.SshKey
+	sshUser := pool.env.SshUser
 	if sshKey == "" {
 		// Assuming image supports GCE ssh fanciness.
 		sshKey = gceKey
 		sshUser = "syzkaller"
 	}
 	Logf(0, "wait instance to boot: %v (%v)", name, ip)
-	if err := waitInstanceBoot(pool.env.Debug, ip, sshKey, sshUser); err != nil {
+	if err := pool.waitInstanceBoot(ip, sshKey, sshUser); err != nil {
 		return nil, err
 	}
 	ok = true
 	inst := &instance{
+		env:     pool.env,
 		cfg:     pool.cfg,
 		debug:   pool.env.Debug,
 		GCE:     pool.GCE,
@@ -272,8 +278,10 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
 		sshRpipe.Close()
 		return nil, nil, err
 	}
-	if inst.sshUser != "root" {
-		command = fmt.Sprintf("sudo bash -c '%v'", command)
+	if inst.env.OS == "linux" {
+		if inst.sshUser != "root" {
+			command = fmt.Sprintf("sudo bash -c '%v'", command)
+		}
 	}
 	args := append(sshArgs(inst.debug, inst.sshKey, "-p", 22), inst.sshUser+"@"+inst.name, command)
 	ssh := exec.Command("ssh", args...)
@@ -334,13 +342,17 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
 	return merger.Output, errc, nil
 }
 
-func waitInstanceBoot(debug bool, ip, sshKey, sshUser string) error {
+func (pool *Pool) waitInstanceBoot(ip, sshKey, sshUser string) error {
+	pwd := "pwd"
+	if pool.env.OS == "windows" {
+		pwd = "dir"
+	}
 	for i := 0; i < 100; i++ {
 		if !vmimpl.SleepInterruptible(5 * time.Second) {
 			return fmt.Errorf("shutdown in progress")
 		}
-		args := append(sshArgs(debug, sshKey, "-p", 22), sshUser+"@"+ip, "pwd")
-		if _, err := runCmd(debug, "ssh", args...); err == nil {
+		args := append(sshArgs(pool.env.Debug, sshKey, "-p", 22), sshUser+"@"+ip, pwd)
+		if _, err := runCmd(pool.env.Debug, "ssh", args...); err == nil {
 			return nil
 		}
 	}
@@ -425,7 +437,7 @@ func sshArgs(debug bool, sshKey, portArg string, port int) []string {
 		"-o", "BatchMode=yes",
 		"-o", "IdentitiesOnly=yes",
 		"-o", "StrictHostKeyChecking=no",
-		"-o", "ConnectTimeout=5",
+		"-o", "ConnectTimeout=10",
 	}
 	if debug {
 		args = append(args, "-v")