cmd/buildproxywrap/main.go

// Copyright 2023 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can
// found in the LICENSE file.

/**
buildproxywrap is a command wrapper that starts/stops build service
relays around a command that typically involves bazel.

Usage: buildproxywrap --cfg FILE -- command...

An example relay configuration file looks like:

[
  {
    "name": "sponge",
    "socket_file_name": "sponge.sock",
    "socket_path_env_var": "BAZEL_sponge_socket_path",
    "server_address": "buildeventservice-pa.googleapis.com:443"
  },
  {
    "name": "resultstore",
    "socket_file_name": "resultstore.sock",
    "socket_path_env_var": "BAZEL_resultstore_socket_path",
    "server_address": "buildeventservice.googleapis.com:443"
  },
  {
    "name": "RBE",
    "socket_file_name": "rbe.sock",
    "socket_path_env_var": "BAZEL_rbe_socket_path",
    "server_address": "remotebuildexecution.googleapis.com:443"
  }
]

The above configuration will setup:
	Service      -> Socket file environment
	----------------------------------------------------------------------
	sponge       -> BAZEL_sponge_socket_path       (for bazel --bes_proxy)
	resultstore  -> BAZEL_resultstore_socket_path  (for bazel --bes_proxy)
	RBE          -> BAZEL_rbe_socket_path          (for bazel --remote_proxy)

Example: run bazel using remote services through a proxy

	(assuming that fx bazel responds to the above environment variables)

	buildproxywrap ... -- fx bazel build --config=remote ...
**/
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"

	"github.com/golang/glog"
)

// wrapCommand runs `command` with environment `env` in a subprocess,
// and returns its exit code.
// The `env` environment contains paths to various socket files used
// by the relays.
// Forward all standard pipes.
func wrapCommand(ctx context.Context, command []string, env []string) error {
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	cmd.Env = append(os.Environ(), env...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	return cmd.Run()
}

// errorToExitCode converts an error to a program exit code.
func errorToExitCode(err error) int {
	if err == nil {
		return 0
	}
	if exiterr, ok := err.(*exec.ExitError); ok {
		exitCode := exiterr.ExitCode()
		glog.V(0).Infof("Command exited %d", exitCode)
		return exitCode
	}
	glog.Errorf("Error: %v", err)
	return 2 // Some other error.
}

// innerMain is the main routine, that returns an error from the subprocess.
// Separating this function from main() ensures that all defer calls are
// executed before os.Exit().
func innerMain(ctx context.Context) error {
	var socatPath string
	var socketDir string
	var configFile string
	flag.StringVar(&socatPath, "socat", "socat", "Path to the 'socat' tool.")
	flag.StringVar(&socketDir, "socket_dir", "", "Temporary directory for sockets.  If empty, this directory will be automatically chosen.  In all cases, this directory will be cleaned up on exit.")
	flag.StringVar(&configFile, "cfg", "", "Relay configuration file (required).")
	flag.Parse()
	command := flag.Args()

	// Load relay configuration.
	if configFile == "" {
		return fmt.Errorf("missing required --cfg flag")
	}
	cfgData, err := os.ReadFile(configFile)
	if err != nil {
		return err
	}

	var relays []*socketRelay
	if err := json.Unmarshal(cfgData, &relays); err != nil {
		return fmt.Errorf("failed to parse JSON file %s: %v", configFile, err)
	}

	// Setup a temporary directory for sockets.
	if socketDir == "" {
		var err error
		socketDir, err = os.MkdirTemp("", "bazel_service_proxy.*")
		if err != nil {
			return err
		}
	}
	defer os.RemoveAll(socketDir)

	finished := make(chan int, 0) // sent after wrapped command completes
	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
	defer stop()
	cmdErr := multiRelayWrap(ctx, relays, socketDir, socatPath, func(env []string) error {
		defer func() {
			close(finished)
		}()
		return wrapCommand(ctx, command, env)
	})
	// Wait for either completion or interrupt
	select {
	case <-finished:
	case <-ctx.Done():
	}
	return cmdErr
}

func main() {
	os.Exit(errorToExitCode(innerMain(context.Background())))
}