From ba1ab8a4b129e26d02a72a131a76f7430bb28a3f Mon Sep 17 00:00:00 2001
From: Josh Kalderimis <josh.kalderimis@gmail.com>
Date: Wed, 8 May 2019 22:02:29 +0200
Subject: [PATCH 1/3] support additional startup script commands

this allows additional startup commands to be passed into cloud-init for GCE instances.
---
 backend/gce.go | 115 +++++++++++++++++++++++++++----------------------
 1 file changed, 63 insertions(+), 52 deletions(-)

diff --git a/backend/gce.go b/backend/gce.go
index ccd3de5f1..06a71a421 100644
--- a/backend/gce.go
+++ b/backend/gce.go
@@ -107,19 +107,20 @@ var (
 		"RATE_LIMIT_DYNAMIC_CONFIG":     "get max-calls and duration dynamically through redis (default false)",
 		"RATE_LIMIT_DYNAMIC_CONFIG_TTL": fmt.Sprintf("time to cache dynamic config for (default %v)", defaultGCERateLimitDynamicConfigTTL),
 
-		"BACKOFF_RETRY_MAX":     "Maximum allowed duration of generic exponential backoff retries (default 1m)",
-		"REGION":                fmt.Sprintf("only takes effect when SUBNETWORK is defined; region in which to deploy (default %v)", defaultGCERegion),
-		"SKIP_STOP_POLL":        "immediately return after issuing first instance deletion request (default false)",
-		"SSH_DIAL_TIMEOUT":      fmt.Sprintf("connection timeout for ssh connections (default %v)", defaultGCESSHDialTimeout),
-		"STOP_POLL_SLEEP":       fmt.Sprintf("sleep interval between polling server for instance stop status (default %v)", defaultGCEStopPollSleep),
-		"STOP_PRE_POLL_SLEEP":   fmt.Sprintf("time to sleep prior to polling server for instance stop status (default %v)", defaultGCEStopPrePollSleep),
-		"SUBNETWORK":            fmt.Sprintf("the subnetwork in which to launch build instances (gce internal default \"%v\")", defaultGCESubnet),
-		"UPLOAD_RETRIES":        fmt.Sprintf("number of times to attempt to upload script before erroring (default %d)", defaultGCEUploadRetries),
-		"UPLOAD_RETRY_SLEEP":    fmt.Sprintf("sleep interval between script upload attempts (default %v)", defaultGCEUploadRetrySleep),
-		"WARMER_URL":            "URL for warmer service",
-		"WARMER_TIMEOUT":        fmt.Sprintf("timeout for requests to warmer service (default %v)", defaultGCEWarmerTimeout),
-		"WARMER_SSH_PASSPHRASE": fmt.Sprintf("The passphrase used to decipher instace SSH keys"),
-		"ZONE":                  fmt.Sprintf("zone name (default %q)", defaultGCEZone),
+		"BACKOFF_RETRY_MAX":       "Maximum allowed duration of generic exponential backoff retries (default 1m)",
+		"REGION":                  fmt.Sprintf("only takes effect when SUBNETWORK is defined; region in which to deploy (default %v)", defaultGCERegion),
+		"SKIP_STOP_POLL":          "immediately return after issuing first instance deletion request (default false)",
+		"SSH_DIAL_TIMEOUT":        fmt.Sprintf("connection timeout for ssh connections (default %v)", defaultGCESSHDialTimeout),
+		"STARTUP_SCRIPT_ADDITION": "Additional commands to run during VM startup",
+		"STOP_POLL_SLEEP":         fmt.Sprintf("sleep interval between polling server for instance stop status (default %v)", defaultGCEStopPollSleep),
+		"STOP_PRE_POLL_SLEEP":     fmt.Sprintf("time to sleep prior to polling server for instance stop status (default %v)", defaultGCEStopPrePollSleep),
+		"SUBNETWORK":              fmt.Sprintf("the subnetwork in which to launch build instances (gce internal default \"%v\")", defaultGCESubnet),
+		"UPLOAD_RETRIES":          fmt.Sprintf("number of times to attempt to upload script before erroring (default %d)", defaultGCEUploadRetries),
+		"UPLOAD_RETRY_SLEEP":      fmt.Sprintf("sleep interval between script upload attempts (default %v)", defaultGCEUploadRetrySleep),
+		"WARMER_URL":              "URL for warmer service",
+		"WARMER_TIMEOUT":          fmt.Sprintf("timeout for requests to warmer service (default %v)", defaultGCEWarmerTimeout),
+		"WARMER_SSH_PASSPHRASE":   fmt.Sprintf("The passphrase used to decipher instace SSH keys"),
+		"ZONE":                    fmt.Sprintf("zone name (default %q)", defaultGCEZone),
 	}
 
 	errGCEMissingIPAddressError   = fmt.Errorf("no IP address found")
@@ -131,6 +132,7 @@ cat > ~travis/.ssh/authorized_keys <<EOF
 {{ .SSHPubKey }}
 EOF
 chown -R travis:travis ~travis/.ssh/
+{{ .StartupScriptAddition }}
 `))
 
 	gceWindowsStartupScript = template.Must(template.New("gce-windows-startup").Parse(`
@@ -146,11 +148,12 @@ Set-LocalUser -Name travis -Password $pw
 )
 
 type gceStartupScriptData struct {
-	AutoImplode        bool
-	HardTimeoutMinutes int64
-	SSHPubKey          string
-	HardTimeoutSeconds int64
-	WindowsPassword    string
+	AutoImplode           bool
+	HardTimeoutMinutes    int64
+	SSHPubKey             string
+	HardTimeoutSeconds 	  int64
+	WindowsPassword       string
+	StartupScriptAddition string
 }
 
 func init() {
@@ -212,23 +215,24 @@ type gceProvider struct {
 }
 
 type gceInstanceConfig struct {
-	MachineType        string
-	PremiumMachineType string
-	Zone               *compute.Zone
-	Network            *compute.Network
-	Subnetwork         *compute.Subnetwork
-	AcceleratorConfig  *compute.AcceleratorConfig
-	DiskSize           int64
-	SSHPubKey          string
-	AutoImplode        bool
-	HardTimeoutMinutes int64
-	StopPollSleep      time.Duration
-	StopPrePollSleep   time.Duration
-	SkipStopPoll       bool
-	Preemptible        bool
-	PublicIP           bool
-	PublicIPConnect    bool
-	Site               string
+	MachineType           string
+	PremiumMachineType    string
+	Zone                  *compute.Zone
+	Network               *compute.Network
+	Subnetwork            *compute.Subnetwork
+	AcceleratorConfig     *compute.AcceleratorConfig
+	DiskSize              int64
+	SSHPubKey             string
+	AutoImplode           bool
+	HardTimeoutMinutes    int64
+	StartupScriptAddition string
+	StopPollSleep         time.Duration
+	StopPrePollSleep      time.Duration
+	SkipStopPoll          bool
+	Preemptible           bool
+	PublicIP              bool
+	PublicIPConnect       bool
+	Site                  string
 }
 
 type gceStartMultistepWrapper struct {
@@ -469,6 +473,11 @@ func newGCEProvider(cfg *config.ProviderConfig) (Provider, error) {
 		autoImplode = ai
 	}
 
+	StartupScriptAddition := ""
+	if cfg.IsSet("STARTUP_SCRIPT") {
+		StartupScriptAddition = cfg.Get("STARTUP_SCRIPT_ADDITION")
+	}
+
 	imageSelectorType := defaultGCEImageSelectorType
 	if cfg.IsSet("IMAGE_SELECTOR_TYPE") {
 		imageSelectorType = cfg.Get("IMAGE_SELECTOR_TYPE")
@@ -605,19 +614,20 @@ func newGCEProvider(cfg *config.ProviderConfig) (Provider, error) {
 		sshDialTimeout:       sshDialTimeout,
 
 		ic: &gceInstanceConfig{
-			Preemptible:        preemptible,
-			PublicIP:           publicIP,
-			PublicIPConnect:    publicIPConnect,
-			DiskSize:           diskSize,
-			SSHPubKey:          string(pubKey),
-			AutoImplode:        autoImplode,
-			StopPollSleep:      stopPollSleep,
-			StopPrePollSleep:   stopPrePollSleep,
-			SkipStopPoll:       skipStopPoll,
-			Site:               site,
-			AcceleratorConfig:  defaultAcceleratorConfig,
-			MachineType:        mtName,
-			PremiumMachineType: premiumMTName,
+			Preemptible:           preemptible,
+			PublicIP:              publicIP,
+			PublicIPConnect:       publicIPConnect,
+			DiskSize:              diskSize,
+			SSHPubKey:             string(pubKey),
+			AutoImplode:           autoImplode,
+			StartupScriptAddition: startupScriptAddition,
+			StopPollSleep:         stopPollSleep,
+			StopPrePollSleep:      stopPrePollSleep,
+			SkipStopPoll:          skipStopPoll,
+			Site:                  site,
+			AcceleratorConfig:     defaultAcceleratorConfig,
+			MachineType:           mtName,
+			PremiumMachineType:    premiumMTName,
 		},
 
 		backoffRetryMax:       backoffRetryMax,
@@ -1009,10 +1019,11 @@ func (p *gceProvider) stepRenderScript(c *gceStartContext) multistep.StepAction
 
 	scriptBuf := bytes.Buffer{}
 	scriptData := gceStartupScriptData{
-		AutoImplode:        p.ic.AutoImplode,
-		HardTimeoutMinutes: int64(c.startAttributes.HardTimeout.Minutes()) + 10,
-		SSHPubKey:          p.ic.SSHPubKey,
-		HardTimeoutSeconds: int64(c.startAttributes.HardTimeout.Seconds()) + 600,
+		AutoImplode:           p.ic.AutoImplode,
+		HardTimeoutMinutes:    int64(c.startAttributes.HardTimeout.Minutes()) + 10,
+		SSHPubKey:             p.ic.SSHPubKey,
+		HardTimeoutSeconds:    int64(c.startAttributes.HardTimeout.Seconds()) + 600,
+		StartupScriptAddition: p.ic.StartupScriptAddition,
 	}
 
 	var err error

From 2d481bee7c70cd6a89f98cbb3aacbd5631214ebf Mon Sep 17 00:00:00 2001
From: Josh Kalderimis <josh.kalderimis@gmail.com>
Date: Wed, 8 May 2019 23:54:51 +0200
Subject: [PATCH 2/3] correct var naming

---
 backend/gce.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/backend/gce.go b/backend/gce.go
index 06a71a421..5be1de503 100644
--- a/backend/gce.go
+++ b/backend/gce.go
@@ -473,9 +473,9 @@ func newGCEProvider(cfg *config.ProviderConfig) (Provider, error) {
 		autoImplode = ai
 	}
 
-	StartupScriptAddition := ""
+	startupScriptAddition := ""
 	if cfg.IsSet("STARTUP_SCRIPT") {
-		StartupScriptAddition = cfg.Get("STARTUP_SCRIPT_ADDITION")
+		startupScriptAddition = cfg.Get("STARTUP_SCRIPT_ADDITION")
 	}
 
 	imageSelectorType := defaultGCEImageSelectorType

From 87c2dc941ae6e1b4a9d6f81018d374bc40ae34a9 Mon Sep 17 00:00:00 2001
From: Josh Kalderimis <josh.kalderimis@gmail.com>
Date: Fri, 10 May 2019 14:19:22 +0200
Subject: [PATCH 3/3] fix an env var check

---
 backend/gce.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/backend/gce.go b/backend/gce.go
index 5be1de503..6d7b57c0c 100644
--- a/backend/gce.go
+++ b/backend/gce.go
@@ -474,7 +474,7 @@ func newGCEProvider(cfg *config.ProviderConfig) (Provider, error) {
 	}
 
 	startupScriptAddition := ""
-	if cfg.IsSet("STARTUP_SCRIPT") {
+	if cfg.IsSet("STARTUP_SCRIPT_ADDITION") {
 		startupScriptAddition = cfg.Get("STARTUP_SCRIPT_ADDITION")
 	}