Fix gcp support

gcp.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bytes"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -12,49 +11,37 @@ import (
 
 	"github.com/taskcluster/generic-worker/gwconfig"
 	"github.com/taskcluster/httpbackoff"
+	"github.com/taskcluster/taskcluster-client-go/tcworkermanager"
 )
 
 var (
 	// not a const, because in testing we swap this out
-	GCPMetadataBaseURL = "http://metadata.google.internal/computeMetadata/v1/"
+	GCPMetadataBaseURL = "http://metadata.google.internal/computeMetadata/v1"
 )
 
 type GCPUserData struct {
-	WorkerType                string `json:"workerType"`
-	WorkerGroup               string `json:"workerGroup"`
-	ProvisionerID             string `json:"provisionerId"`
-	CredentialURL             string `json:"credentialURL"`
-	Audience                  string `json:"audience"`
-	Ed25519SigningKeyLocation string `json:"ed25519SigningKeyLocation"`
-	RootURL                   string `json:"rootURL"`
+	WorkerPoolID string                       `json:"workerPoolId"`
+	ProviderID   string                       `json:"providerId"`
+	WorkerGroup  string                       `json:"workerGroup"`
+	RootURL      string                       `json:"rootURL"`
+	WorkerConfig WorkerTypeDefinitionUserData `json:"workerConfig"`
 }
 
-type CredentialRequestData struct {
-	Token string `json:"token"`
-}
-
-type TaskclusterCreds struct {
-	AccessToken string `json:"accessToken"`
-	ClientID    string `json:"clientId"`
-	Certificate string `json:"certificate"`
-}
-
-func queryGCPMetaData(client *http.Client, path string) (string, error) {
+func queryGCPMetaData(client *http.Client, path string) ([]byte, error) {
 	req, err := http.NewRequest("GET", GCPMetadataBaseURL+path, nil)
 
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	req.Header.Add("Metadata-Flavor", "Google")
 
 	resp, _, err := httpbackoff.ClientDo(client, req)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	defer resp.Body.Close()
-	content, err := ioutil.ReadAll(resp.Body)
-	return string(content), err
+	return ioutil.ReadAll(resp.Body)
 }
 
 func updateConfigWithGCPSettings(c *gwconfig.Config) error {
@@ -64,77 +51,80 @@ func updateConfigWithGCPSettings(c *gwconfig.Config) error {
 	c.ShutdownMachineOnIdle = true
 
 	client := &http.Client{}
-	userDataString, err := queryGCPMetaData(client, "instance/attributes/config")
+
+	workerID, err := queryGCPMetaData(client, "/instance/id")
 	if err != nil {
-		return err
+		return fmt.Errorf("Could not query instance ID: %v", err)
+	}
+	c.WorkerID = string(workerID)
+
+	taskclusterConfig, err := queryGCPMetaData(client, "/instance/attributes/taskcluster")
+	if err != nil {
+		return fmt.Errorf("Could not query taskcluster configuration: %v", err)
 	}
 
 	var userData GCPUserData
-	err = json.Unmarshal([]byte(userDataString), &userData)
+	err = json.Unmarshal(taskclusterConfig, &userData)
 	if err != nil {
 		return err
 	}
 
-	c.ProvisionerID = userData.ProvisionerID
-	c.WorkerType = userData.WorkerType
-	c.WorkerGroup = userData.WorkerGroup
+	wp := strings.SplitN(userData.WorkerPoolID, "/", -1)
+	if len(wp) != 2 {
+		return fmt.Errorf("Was expecting WorkerPoolID to have syntax <provisionerId>/<workerType> but was %q", userData.WorkerPoolID)
+	}
 
+	c.ProvisionerID = wp[0]
+	c.WorkerType = wp[1]
+	c.WorkerGroup = userData.WorkerGroup
 	c.RootURL = userData.RootURL
-	c.Ed25519SigningKeyLocation = userData.Ed25519SigningKeyLocation
 
-	// Now we get taskcluster credentials via instance identity
-	// TODO: Disable getting instance identity after first run
-	audience := userData.Audience
-	instanceIDPath := fmt.Sprintf("instance/service-accounts/default/identity?audience=%s&format=full", audience)
-	instanceIDToken, err := queryGCPMetaData(client, instanceIDPath)
-	if err != nil {
-		return err
-	}
+	// We need a worker manager client for fetching taskcluster credentials.
+	// Ensure auth is disabled in client, since we don't have credentials yet.
+	wm := c.WorkerManager()
+	wm.Authenticate = false
+	wm.Credentials = nil
 
-	data := CredentialRequestData{Token: instanceIDToken}
-	reqData, err := json.Marshal(data)
+	identity, err := queryGCPMetaData(client, "/instance/service-accounts/default/identity?audience="+userData.RootURL+"&format=full")
 	if err != nil {
-		return err
+		return fmt.Errorf("Could not query google indentity token: %v", err)
+	}
+	providerType := tcworkermanager.GoogleProviderType{
+		Token: string(identity),
 	}
-	dataBuffer := bytes.NewBuffer(reqData)
 
-	credentialURL := userData.CredentialURL
-	req, err := http.NewRequest("POST", credentialURL, dataBuffer)
+	workerIdentityProof, err := json.Marshal(providerType)
 	if err != nil {
-		return err
+		return fmt.Errorf("Could not marshal google provider type %#v: %v", providerType, err)
 	}
 
-	req.Header.Add("Content-Type", "application/json")
+	reg, err := wm.RegisterWorker(&tcworkermanager.RegisterWorkerRequest{
+		WorkerPoolID:        userData.WorkerPoolID,
+		ProviderID:          userData.ProviderID,
+		WorkerGroup:         userData.WorkerGroup,
+		WorkerID:            c.WorkerID,
+		WorkerIdentityProof: json.RawMessage(workerIdentityProof),
+	})
 
-	resp, _, err := httpbackoff.ClientDo(client, req)
 	if err != nil {
-		return err
-	}
-	defer resp.Body.Close()
-	content, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return err
+		return fmt.Errorf("Could not register worker: %v", err)
 	}
 
-	var creds TaskclusterCreds
-	err = json.Unmarshal([]byte(content), &creds)
-	if err != nil {
-		return err
-	}
+	c.AccessToken = reg.Credentials.AccessToken
+	c.Certificate = reg.Credentials.Certificate
+	c.ClientID = reg.Credentials.ClientID
 
-	c.AccessToken = creds.AccessToken
-	c.ClientID = creds.ClientID
-	c.Certificate = creds.Certificate
+	// TODO: process reg.Expires
 
 	gcpMetadata := map[string]interface{}{}
 	for _, path := range []string{
-		"instance/image",
-		"instance/id",
-		"instance/machine-type",
-		"instance/network-interfaces/0/access-configs/0/external-ip",
-		"instance/zone",
-		"instance/hostname",
-		"instance/network-interfaces/0/ip",
+		"/instance/image",
+		"/instance/id",
+		"/instance/machine-type",
+		"/instance/network-interfaces/0/access-configs/0/external-ip",
+		"/instance/zone",
+		"/instance/hostname",
+		"/instance/network-interfaces/0/ip",
 	} {
 		key := path[strings.LastIndex(path, "/")+1:]
 		value, err := queryGCPMetaData(client, path)
@@ -151,8 +141,31 @@ func updateConfigWithGCPSettings(c *gwconfig.Config) error {
 	c.InstanceType = gcpMetadata["machine-type"].(string)
 	c.AvailabilityZone = gcpMetadata["zone"].(string)
 
-	// TODO: Fetch these from secrets
-	c.LiveLogSecret = "foobar"
+	// Parse the config before applying it, to ensure that no disallowed fields
+	// are included.
+	_, err = userData.WorkerConfig.PublicHostSetup()
+	if err != nil {
+		return fmt.Errorf("Error retrieving/interpreting host setup from GCP metadata: %v", err)
+	}
+
+	// Host setup per worker type "userData" section.
+	//
+	// Note, we first update configuration from public host setup, before
+	// calling tc-secrets to get private host setup, in case secretsBaseURL is
+	// configured in userdata.
+	err = c.MergeInJSON(userData.WorkerConfig.GenericWorker, func(a map[string]interface{}) map[string]interface{} {
+		if config, exists := a["config"]; exists {
+			return config.(map[string]interface{})
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("Error applying /data/genericWorker/config from AWS userdata to config: %v", err)
+	}
+
+	if c.IdleTimeoutSecs == 0 {
+		c.IdleTimeoutSecs = 3600
+	}
 
 	return nil
 }

gwconfig/config.go

@@ -13,6 +13,7 @@ import (
 	"github.com/taskcluster/taskcluster-client-go/tcpurgecache"
 	"github.com/taskcluster/taskcluster-client-go/tcqueue"
 	"github.com/taskcluster/taskcluster-client-go/tcsecrets"
+	"github.com/taskcluster/taskcluster-client-go/tcworkermanager"
 )
 
 type (
@@ -63,6 +64,7 @@ type (
 		TasksDir                       string                 `json:"tasksDir"`
 		WorkerGroup                    string                 `json:"workerGroup"`
 		WorkerID                       string                 `json:"workerId"`
+		WorkerManagerBaseURL           string                 `json:"workerManagerBaseURL"`
 		WorkerType                     string                 `json:"workerType"`
 		WorkerTypeMetadata             map[string]interface{} `json:"workerTypeMetadata"`
 		WSTAudience                    string                 `json:"wstAudience"`
@@ -207,3 +209,12 @@ func (c *Config) Secrets() *tcsecrets.Secrets {
 	}
 	return secrets
 }
+
+func (c *Config) WorkerManager() *tcworkermanager.WorkerManager {
+	workerManager := tcworkermanager.New(c.Credentials(), c.RootURL)
+	// If workerManagerBaseURL provided, it should take precedence over rootURL
+	if c.WorkerManagerBaseURL != "" {
+		workerManager.BaseURL = c.WorkerManagerBaseURL
+	}
+	return workerManager
+}

main.go

@@ -220,6 +220,7 @@ func loadConfig(filename string, queryAWSUserData bool, queryGCPMetaData bool) (
 			TaskclusterProxyPort:           80,
 			TasksDir:                       defaultTasksDir(),
 			WorkerGroup:                    "test-worker-group",
+			WorkerManagerBaseURL:           "",
 			WorkerTypeMetadata:             map[string]interface{}{},
 		},
 	}

usage.go

@@ -254,6 +254,11 @@ and reports back results to the queue.
                                             identifier to uniquely identify which pool of
                                             workers this worker logically belongs to.
                                             [default: test-worker-group]
+          workerManagerBaseURL              The base URL for taskcluster worker-manager API calls.
+                                            If not provided, the base URL for API calls is
+                                            instead derived from rootURL setting as follows:
+                                              * https://worker-manager.taskcluster.net/v1 for rootURL https://taskcluster.net
+                                              * <rootURL>/api/worker-manager/v1 for all other rootURLs
           workerTypeMetaData                This arbitrary json blob will be included at the
                                             top of each task log. Providing information here,
                                             such as a URL to the code/config used to set up the