chromium-setup/main.go

// hetzner-runner provisions an ephemeral Hetzner Cloud server for a
// Chromium build and guarantees it gets torn down afterward.
//
// Usage:
//
//	hetzner-runner up                  # create server, print IP, write state file
//	hetzner-runner down                # delete server + ssh key from state file
//	hetzner-runner run -- <cmd> [args] # up, ssh-exec cmd on the box, down (always)
//	hetzner-runner run -script build.sh -env GITEA_TOKEN -env MIRROR_READ_TOKEN
//	                                    # up, pipe build.sh to `bash -s` on the box
//	                                    # (env vars forwarded from this process's
//	                                    # own environment, not the CLI args), down (always)
//
// HCLOUD_TOKEN must be set in the environment. Never pass secrets as flags
// (shows up in process lists / CI logs) — use -env to forward them instead.
package main

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/hetznercloud/hcloud-go/v2/hcloud"
	"golang.org/x/crypto/ssh"
)

const stateFileName = ".hetzner-runner-state.json"

// selfDestructScript is cloud-init user-data that runs on the box itself.
// It's a failsafe independent of this program's lifetime: if `down` never
// runs (laptop dies, network drops, CI job gets killed), the server deletes
// itself via the Hetzner API after maxLifetime regardless.
const selfDestructScript = `#cloud-config
write_files:
  - path: /root/.hcloud-token
    permissions: '0600'
    content: '%s'
runcmd:
  - [ bash, -c, "echo 'curl -sf -X DELETE -H \"Authorization: Bearer $(cat /root/.hcloud-token)\" https://api.hetzner.cloud/v1/servers/$(curl -sf http://169.254.169.254/hetzner/v1/metadata/instance-id)' | at now + %d hours" ]
`

type state struct {
	ServerID    int64  `json:"server_id"`
	ServerName  string `json:"server_name"`
	SSHKeyID    int64  `json:"ssh_key_id"`
	IPv4        string `json:"ipv4"`
	PrivKeyPath string `json:"priv_key_path"`
}

type config struct {
	serverType  string
	image       string
	location    string
	namePrefix  string
	maxLifetime time.Duration
	sshTimeout  time.Duration
}

func defaultConfig() config {
	return config{
		serverType:  envOr("HETZNER_SERVER_TYPE", "cpx51"),
		image:       envOr("HETZNER_IMAGE", "ubuntu-24.04"),
		location:    envOr("HETZNER_LOCATION", "hil"),
		namePrefix:  envOr("HETZNER_NAME_PREFIX", "chromium-build"),
		maxLifetime: 8 * time.Hour,
		sshTimeout:  3 * time.Minute,
	}
}

func envOr(k, def string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return def
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "usage: hetzner-runner <up|down|run> [flags]")
		os.Exit(2)
	}

	token := os.Getenv("HCLOUD_TOKEN")
	if token == "" {
		fatal("HCLOUD_TOKEN must be set in the environment")
	}
	client := hcloud.NewClient(hcloud.WithToken(token))
	cfg := defaultConfig()

	switch os.Args[1] {
	case "up":
		fs := flag.NewFlagSet("up", flag.ExitOnError)
		bindConfig(fs, &cfg)
		fs.Parse(os.Args[2:])
		st, err := up(context.Background(), client, token, cfg)
		if err != nil {
			fatal("up: %v", err)
		}
		saveState(st)
		fmt.Printf("server ready: id=%d ip=%s\n", st.ServerID, st.IPv4)

	case "down":
		st, err := loadState()
		if err != nil {
			fatal("down: %v (nothing to tear down, or run from the same directory as `up`)", err)
		}
		if err := down(context.Background(), client, st); err != nil {
			fatal("down: %v", err)
		}
		os.Remove(stateFileName)
		fmt.Println("torn down")

	case "run":
		fs := flag.NewFlagSet("run", flag.ExitOnError)
		bindConfig(fs, &cfg)
		var scriptPath string
		var envNames stringSlice
		fs.StringVar(&scriptPath, "script", "", "local script to pipe to `bash -s` on the remote box, instead of a -- command")
		fs.Var(&envNames, "env", "name of an env var to read from this process's environment and forward to the remote script (repeatable)")
		fs.Parse(os.Args[2:])

		var payload remotePayload
		if scriptPath != "" {
			script, err := os.ReadFile(scriptPath)
			if err != nil {
				fatal("run: reading -script %q: %v", scriptPath, err)
			}
			env := map[string]string{}
			for _, name := range envNames {
				v, ok := os.LookupEnv(name)
				if !ok {
					fatal("run: -env %s: not set in this process's environment", name)
				}
				env[name] = v
			}
			payload = remotePayload{script: string(script), env: env}
		} else {
			remainder := fs.Args()
			if len(remainder) == 0 {
				fatal("run: usage: hetzner-runner run [flags] -- <remote command...>  (or -script <path>)")
			}
			payload = remotePayload{command: strings.Join(remainder, " ")}
		}

		os.Exit(runLifecycle(client, token, cfg, payload))

	default:
		fmt.Fprintln(os.Stderr, "usage: hetzner-runner <up|down|run> [flags]")
		os.Exit(2)
	}
}

func bindConfig(fs *flag.FlagSet, cfg *config) {
	fs.StringVar(&cfg.serverType, "type", cfg.serverType, "Hetzner server type")
	fs.StringVar(&cfg.image, "image", cfg.image, "OS image")
	fs.StringVar(&cfg.location, "location", cfg.location, "Hetzner location")
	fs.DurationVar(&cfg.maxLifetime, "max-lifetime", cfg.maxLifetime, "failsafe: server self-deletes after this long even if `down` never runs")
	fs.DurationVar(&cfg.sshTimeout, "ssh-timeout", cfg.sshTimeout, "how long to wait for SSH to come up")
}

// remotePayload is either a plain command (`run -- <cmd>`) or a local
// script piped to `bash -s` on the remote box with forwarded env vars
// (`run -script ...`). Exactly one of command / script is set.
type remotePayload struct {
	command string
	script  string
	env     map[string]string // only used with script; forwarded via `export`, never argv
}

// stringSlice is a repeatable string flag (flag.Value).
type stringSlice []string

func (s *stringSlice) String() string     { return strings.Join(*s, ",") }
func (s *stringSlice) Set(v string) error { *s = append(*s, v); return nil }

// runLifecycle does up -> exec -> down, guaranteeing down runs even if the
// remote command fails or this process is interrupted (Ctrl-C, CI job
// cancellation).
func runLifecycle(client *hcloud.Client, token string, cfg config, payload remotePayload) (exitCode int) {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	st, err := up(ctx, client, token, cfg)
	if err != nil {
		fatal("up: %v", err)
	}
	saveState(st)

	defer func() {
		// Always tear down, regardless of how we got here.
		downCtx, downCancel := context.WithTimeout(context.Background(), 2*time.Minute)
		defer downCancel()
		if err := down(downCtx, client, st); err != nil {
			fmt.Fprintf(os.Stderr, "WARNING: teardown failed, server %d (%s) may still be running and billing: %v\n", st.ServerID, st.IPv4, err)
			fmt.Fprintf(os.Stderr, "  manual cleanup: hcloud server delete %d\n", st.ServerID)
			return
		}
		os.Remove(stateFileName)
		fmt.Println("torn down")
	}()

	var code int
	if payload.script != "" {
		fmt.Printf("server ready: id=%d ip=%s, running script (%d bytes, %d env vars forwarded)\n", st.ServerID, st.IPv4, len(payload.script), len(payload.env))
		code, err = sshExecScript(st, payload.script, payload.env)
	} else {
		fmt.Printf("server ready: id=%d ip=%s, running: %s\n", st.ServerID, st.IPv4, payload.command)
		code, err = sshExec(st, payload.command)
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "remote command failed: %v\n", err)
	}
	return code
}

func up(ctx context.Context, client *hcloud.Client, token string, cfg config) (*state, error) {
	pub, privPath, err := ensureEphemeralKey()
	if err != nil {
		return nil, fmt.Errorf("generating ephemeral ssh key: %w", err)
	}

	name := fmt.Sprintf("%s-%s", cfg.namePrefix, time.Now().Format("20060102-150405"))

	keyResult, _, err := client.SSHKey.Create(ctx, hcloud.SSHKeyCreateOpts{
		Name:      name,
		PublicKey: pub,
	})
	if err != nil {
		return nil, fmt.Errorf("uploading ephemeral ssh key: %w", err)
	}

	serverType, _, err := client.ServerType.GetByName(ctx, cfg.serverType)
	if err != nil || serverType == nil {
		return nil, fmt.Errorf("looking up server type %q: %w", cfg.serverType, err)
	}
	image, _, err := client.Image.GetByNameAndArchitecture(ctx, cfg.image, serverType.Architecture)
	if err != nil || image == nil {
		return nil, fmt.Errorf("looking up image %q: %w", cfg.image, err)
	}
	location, _, err := client.Location.GetByName(ctx, cfg.location)
	if err != nil || location == nil {
		return nil, fmt.Errorf("looking up location %q: %w", cfg.location, err)
	}

	userData := fmt.Sprintf(selfDestructScript, token, int(cfg.maxLifetime.Hours()))

	result, _, err := client.Server.Create(ctx, hcloud.ServerCreateOpts{
		Name:       name,
		ServerType: serverType,
		Image:      image,
		Location:   location,
		SSHKeys:    []*hcloud.SSHKey{keyResult},
		UserData:   userData,
	})
	if err != nil {
		return nil, fmt.Errorf("creating server: %w", err)
	}

	if err := waitForAction(ctx, client, result.Action); err != nil {
		return nil, fmt.Errorf("waiting for server creation: %w", err)
	}

	st := &state{
		ServerID:    result.Server.ID,
		ServerName:  name,
		SSHKeyID:    keyResult.ID,
		IPv4:        result.Server.PublicNet.IPv4.IP.String(),
		PrivKeyPath: privPath,
	}

	if err := waitForSSH(st, cfg.sshTimeout); err != nil {
		// Best-effort cleanup on partial failure so we don't leak a
		// server the caller doesn't know the ID of yet.
		_ = down(ctx, client, st)
		return nil, fmt.Errorf("waiting for ssh: %w", err)
	}

	return st, nil
}

func down(ctx context.Context, client *hcloud.Client, st *state) error {
	var errs []string
	if st.ServerID != 0 {
		if _, _, err := client.Server.DeleteWithResult(ctx, &hcloud.Server{ID: st.ServerID}); err != nil {
			errs = append(errs, fmt.Sprintf("delete server %d: %v", st.ServerID, err))
		}
	}
	if st.SSHKeyID != 0 {
		if _, err := client.SSHKey.Delete(ctx, &hcloud.SSHKey{ID: st.SSHKeyID}); err != nil {
			errs = append(errs, fmt.Sprintf("delete ssh key %d: %v", st.SSHKeyID, err))
		}
	}
	if st.PrivKeyPath != "" {
		os.Remove(st.PrivKeyPath)
	}
	if len(errs) > 0 {
		return fmt.Errorf("%s", strings.Join(errs, "; "))
	}
	return nil
}

func waitForAction(ctx context.Context, client *hcloud.Client, action *hcloud.Action) error {
	progress, errs := client.Action.WatchProgress(ctx, action)
	for {
		select {
		case <-progress:
		case err, ok := <-errs:
			if !ok {
				return nil
			}
			if err != nil {
				return err
			}
		}
	}
}

func waitForSSH(st *state, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	var lastErr error
	for time.Now().Before(deadline) {
		conn, err := net.DialTimeout("tcp", net.JoinHostPort(st.IPv4, "22"), 5*time.Second)
		if err == nil {
			conn.Close()
			// TCP is up; give sshd + cloud-init user creation a moment,
			// then confirm with a real SSH handshake + command.
			time.Sleep(3 * time.Second)
			if code, err := sshExec(st, "true"); err == nil && code == 0 {
				return nil
			} else {
				lastErr = err
			}
		} else {
			lastErr = err
		}
		time.Sleep(5 * time.Second)
	}
	return fmt.Errorf("timed out after %s, last error: %v", timeout, lastErr)
}

func sshExec(st *state, command string) (int, error) {
	session, client, err := sshSession(st)
	if err != nil {
		return -1, err
	}
	defer client.Close()
	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return runSession(session, command)
}

// sshExecScript pipes script to `bash -s` on the remote box, prefixed with
// `export` lines for env. Values are forwarded only over the SSH channel's
// stdin — never as a command-line argument on either end — so they don't
// show up in `ps` output on the box or in this process's argv.
func sshExecScript(st *state, script string, env map[string]string) (int, error) {
	session, client, err := sshSession(st)
	if err != nil {
		return -1, err
	}
	defer client.Close()
	defer session.Close()

	var buf strings.Builder
	for k, v := range env {
		fmt.Fprintf(&buf, "export %s=%s\n", k, shellQuote(v))
	}
	buf.WriteString(script)

	session.Stdin = strings.NewReader(buf.String())
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return runSession(session, "bash -s")
}

func shellQuote(v string) string {
	return "'" + strings.ReplaceAll(v, "'", `'\''`) + "'"
}

func sshSession(st *state) (*ssh.Session, *ssh.Client, error) {
	key, err := os.ReadFile(st.PrivKeyPath)
	if err != nil {
		return nil, nil, err
	}
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, nil, err
	}
	sshCfg := &ssh.ClientConfig{
		User:            "root",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(signer)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // ephemeral box, TOFU risk accepted for build workloads
		Timeout:         10 * time.Second,
	}
	client, err := ssh.Dial("tcp", net.JoinHostPort(st.IPv4, "22"), sshCfg)
	if err != nil {
		return nil, nil, err
	}
	session, err := client.NewSession()
	if err != nil {
		client.Close()
		return nil, nil, err
	}

	// A PTY makes the remote side (ninja, gclient, python) think it's
	// talking to a terminal, so it line-buffers instead of block-buffering
	// for a non-tty pipe -- without this, output for a multi-hour build can
	// arrive in laggy multi-KB chunks instead of streaming as it happens.
	// Combines stderr into stdout, which is fine for a CI log (one
	// interleaved stream is what you'd see watching a real terminal
	// anyway).
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}
	if err := session.RequestPty("xterm", 40, 200, modes); err != nil {
		session.Close()
		client.Close()
		return nil, nil, fmt.Errorf("requesting pty: %w", err)
	}

	return session, client, nil
}

func runSession(session *ssh.Session, command string) (int, error) {
	if err := session.Run(command); err != nil {
		if exitErr, ok := err.(*ssh.ExitError); ok {
			return exitErr.ExitStatus(), nil
		}
		return -1, err
	}
	return 0, nil
}

func ensureEphemeralKey() (pubAuthorizedKeys string, privPath string, err error) {
	dir, err := os.MkdirTemp("", "hetzner-runner-key-*")
	if err != nil {
		return "", "", err
	}
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return "", "", err
	}
	sshPub, err := ssh.NewPublicKey(pub)
	if err != nil {
		return "", "", err
	}
	pkcs8, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", "", err
	}
	privPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: pkcs8})

	privPath = filepath.Join(dir, "id_ed25519")
	if err := os.WriteFile(privPath, privPEM, 0600); err != nil {
		return "", "", err
	}
	return string(ssh.MarshalAuthorizedKey(sshPub)), privPath, nil
}

func saveState(st *state) {
	f, err := os.Create(stateFileName)
	if err != nil {
		fatal("writing state file: %v", err)
	}
	defer f.Close()
	json.NewEncoder(f).Encode(st)
}

func loadState() (*state, error) {
	f, err := os.Open(stateFileName)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var st state
	if err := json.NewDecoder(f).Decode(&st); err != nil {
		return nil, err
	}
	return &st, nil
}

func fatal(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "error: "+format+"\n", args...)
	os.Exit(1)
}