-
Notifications
You must be signed in to change notification settings - Fork 156
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
cmd/bootstrapswarm: add bootstrapswarm used to bootstrap swarming bot
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
1 parent
4de57f2
commit 14925fa
Showing
1 changed file
with
183 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
} |