Skip to content

Commit

Permalink
cmd/bootstrapswarm: add bootstrapswarm used to bootstrap swarming bot
Browse files Browse the repository at this point in the history
This change adds a bootstrapswarm which bootstraps the swarming bot
on two different environments (on GCE and not on GCE). It can be extended
in the future to start the swarming client on other clouds as needed.

Updates golang/go#60468
Updates golang/go#60640

Change-Id: Iead5f980d27441d3bc6d8161d8baf695a5b55d56
Reviewed-on: https://go-review.googlesource.com/c/build/+/504821
Run-TryBot: Dmitri Shuralyov <[email protected]>
Auto-Submit: Dmitri Shuralyov <[email protected]>
Reviewed-by: Dmitri Shuralyov <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: Heschi Kreinick <[email protected]>
Reviewed-by: Carlos Amedee <[email protected]>
  • Loading branch information
cagedmantis authored and gopherbot committed Jun 26, 2023
1 parent 4de57f2 commit 14925fa
Showing 1 changed file with 183 additions and 0 deletions.
183 changes: 183 additions & 0 deletions cmd/bootstrapswarm/bootstrapswarm.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// bootstapswarm will bootstrap the swarming bot depending
// on the environment that it is run on.
//
// On GCE: bootstrapswarm will retrieve authentication credentials
// from the GCE metadata service and use those credentials to download
// the swarming bot. It will then start the swarming bot in a directory
// within the user's home directory.
//
// Requirements:
// - Python3 installed and in the calling user's PATH.
//
// Not on GCE: bootstrapswarm will read the token file and retrieve the
// the luci machine token. It will use that token to authenticate and
// download the swarming bot. It will then start the swarming bot in a
// directory within the user's home directory.
//
// Requirements:
// - Python3 installed and in the calling user's PATH.
// - luci_machine_tokend running as root in a cron job.
// https://chromium.googlesource.com/infra/luci/luci-go/+/refs/heads/main/tokenserver
// Further instructions can be found at https://github.com/golang/go/wiki/DashboardBuilders
// The default locations for the token files should be used if possible:
// Most OS: /var/lib/luci_machine_tokend/token.json
// Windows: C:\luci_machine_tokend\token.json
// - bootstrapswarm should not be run as a privileged user.
package main

import (
"context"
"encoding/json"
"flag"
"fmt"
"io"
"log"
"net/http"
"os"
"os/exec"
"path/filepath"
"runtime"

"cloud.google.com/go/compute/metadata"
)

var (
tokenFilePath = flag.String("token-file-path", defaultTokenLocation(), "Path to the token file (used when not on GCE)")
hostname = flag.String("hostname", os.Getenv("HOSTNAME"), "Hostname of machine to bootstrap (required)")
)

func main() {
flag.Usage = func() {
fmt.Fprintln(os.Stderr, "Usage: bootstrapswarm")
flag.PrintDefaults()
}
flag.Parse()
if *hostname == "" {
flag.Usage()
os.Exit(2)
}
ctx := context.Background()
if err := bootstrap(ctx, *hostname, *tokenFilePath); err != nil {
log.Fatal(err)
}
}

var httpClient = http.DefaultClient

func bootstrap(ctx context.Context, hostname, tokenPath string) error {
httpHeaders := map[string]string{"X-Luci-Swarming-Bot-ID": hostname}
if metadata.OnGCE() {
log.Println("Bootstrapping the swarming bot with GCE authentication")
log.Println("retrieving the GCE VM token")
token, err := retrieveGCEVMToken(ctx)
if err != nil {
return fmt.Errorf("unable to retrieve GCE Machine Token: %w", err)
}
httpHeaders["X-Luci-Gce-Vm-Token"] = token
} else {
log.Println("Bootstrapping the swarming bot with certificate authentication")
log.Println("retrieving the luci-machine-token from the token file")
tokBytes, err := os.ReadFile(tokenPath)
if err != nil {
return fmt.Errorf("unable to read file %q: %w", tokenPath, err)
}
type token struct {
LuciMachineToken string `json:"luci_machine_token"`
}
var tok token
if err := json.Unmarshal(tokBytes, &tok); err != nil {
return fmt.Errorf("unable to unmarshal token %s: %w", tokenPath, err)
}
if tok.LuciMachineToken == "" {
return fmt.Errorf("unable to retrieve machine token from token file %s", tokenPath)
}
httpHeaders["X-Luci-Machine-Token"] = tok.LuciMachineToken
}
log.Println("Downloading the swarming bot")
req, err := http.NewRequestWithContext(ctx, http.MethodGet, `https://chromium-swarm.appspot.com/bot_code`, nil)
if err != nil {
return fmt.Errorf("http.NewRequest: %w", err)
}
for k, v := range httpHeaders {
req.Header.Set(k, v)
}
resp, err := httpClient.Do(req)
if err != nil {
return fmt.Errorf("client.Do: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
return fmt.Errorf("status code %d", resp.StatusCode)
}
botBytes, err := io.ReadAll(resp.Body)
if err != nil {
return fmt.Errorf("io.ReadAll: %w", err)
}
botPath, err := writeToWorkDirectory(botBytes, "swarming_bot.zip")
if err != nil {
return fmt.Errorf("unable to save swarming bot to disk: %w", err)
}
log.Printf("Starting the swarming bot %s", botPath)
cmd := exec.CommandContext(ctx, "python3", botPath, "start_bot")
// swarming client checks the SWARMING_BOT_ID environment variable for hostname overrides.
cmd.Env = append(os.Environ(), fmt.Sprintf("SWARMING_BOT_ID=%s", hostname))
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
return fmt.Errorf("command execution %s: %s", cmd, err)
}
return nil
}

// writeToWorkDirectory writes a file to the swarming working directory and returns the path
// to where the file was written.
func writeToWorkDirectory(b []byte, filename string) (string, error) {
homeDir, err := os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("os.UserHomeDir: %w", err)
}
workDir := filepath.Join(homeDir, ".swarming")
if err := os.Mkdir(workDir, 0755); err != nil && !os.IsExist(err) {
return "", fmt.Errorf("os.Mkdir(%s): %w", workDir, err)
}
path := filepath.Join(workDir, filename)
if err = os.WriteFile(path, b, 0644); err != nil {
return "", fmt.Errorf("os.WriteFile(%s): %w", path, err)
}
return path, nil
}

// retrieveGCEVMToken retrieves a GCE VM token from the GCP metadata service.
func retrieveGCEVMToken(ctx context.Context) (string, error) {
const url = `http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?audience=https://chromium-swarm.appspot.com&format=full`
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return "", fmt.Errorf("http.NewRequest: %w", err)
}
req.Header.Set("Metadata-Flavor", "Google")
resp, err := httpClient.Do(req)
if err != nil {
return "", fmt.Errorf("client.Do: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
return "", fmt.Errorf("status code %d", resp.StatusCode)
}
b, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("io.ReadAll: %w", err)
}
return string(b), nil
}

func defaultTokenLocation() string {
out := "/var/lib/luci_machine_tokend/token.json"
if runtime.GOOS == "windows" {
return `C:\luci_machine_tokend\token.json`
}
return out
}

0 comments on commit 14925fa

Please sign in to comment.