go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/buildproxywrap/main.go (about)

     1  // Copyright 2023 The Fuchsia Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style license that can
     3  // found in the LICENSE file.
     4  
     5  /*
     6  *
     7  buildproxywrap is a command wrapper that starts/stops build service
     8  relays around a command that typically involves bazel.
     9  
    10  Usage: buildproxywrap --cfg FILE -- command...
    11  
    12  An example relay configuration file looks like:
    13  
    14  [
    15  
    16  	{
    17  	  "name": "sponge",
    18  	  "socket_file_name": "sponge.sock",
    19  	  "socket_path_env_var": "BAZEL_sponge_socket_path",
    20  	  "server_address": "buildeventservice-pa.googleapis.com:443"
    21  	},
    22  	{
    23  	  "name": "resultstore",
    24  	  "socket_file_name": "resultstore.sock",
    25  	  "socket_path_env_var": "BAZEL_resultstore_socket_path",
    26  	  "server_address": "buildeventservice.googleapis.com:443"
    27  	},
    28  	{
    29  	  "name": "RBE",
    30  	  "socket_file_name": "rbe.sock",
    31  	  "socket_path_env_var": "BAZEL_rbe_socket_path",
    32  	  "server_address": "remotebuildexecution.googleapis.com:443"
    33  	}
    34  
    35  ]
    36  
    37  The above configuration will setup:
    38  
    39  	Service      -> Socket file environment
    40  	----------------------------------------------------------------------
    41  	sponge       -> BAZEL_sponge_socket_path       (for bazel --bes_proxy)
    42  	resultstore  -> BAZEL_resultstore_socket_path  (for bazel --bes_proxy)
    43  	RBE          -> BAZEL_rbe_socket_path          (for bazel --remote_proxy)
    44  
    45  Example: run bazel using remote services through a proxy
    46  
    47  	(assuming that fx bazel responds to the above environment variables)
    48  
    49  	buildproxywrap ... -- fx bazel build --config=remote ...
    50  
    51  *
    52  */
    53  package main
    54  
    55  import (
    56  	"context"
    57  	"encoding/json"
    58  	"flag"
    59  	"fmt"
    60  	"os"
    61  	"os/exec"
    62  	"os/signal"
    63  
    64  	"github.com/golang/glog"
    65  )
    66  
    67  // wrapCommand runs `command` with environment `env` in a subprocess,
    68  // and returns its exit code.
    69  // The `env` environment contains paths to various socket files used
    70  // by the relays.
    71  // Forward all standard pipes.
    72  func wrapCommand(ctx context.Context, command []string, env []string) error {
    73  	glog.V(1).Infof("wrapCommand: %v", command)
    74  	// If there is no wrapped command to run, just return right away.
    75  	if len(command) == 0 {
    76  		return nil
    77  	}
    78  	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
    79  	cmd.Env = append(os.Environ(), env...)
    80  	cmd.Stdout = os.Stdout
    81  	cmd.Stderr = os.Stderr
    82  	cmd.Stdin = os.Stdin
    83  	return cmd.Run()
    84  }
    85  
    86  // errorToExitCode converts an error to a program exit code.
    87  func errorToExitCode(err error) int {
    88  	if err == nil {
    89  		return 0
    90  	}
    91  	if exiterr, ok := err.(*exec.ExitError); ok {
    92  		exitCode := exiterr.ExitCode()
    93  		glog.V(0).Infof("Command exited %d", exitCode)
    94  		return exitCode
    95  	}
    96  	glog.Errorf("Error: %v", err)
    97  	return 2 // Some other error.
    98  }
    99  
   100  // innerMain is the main routine, that returns an error from the subprocess.
   101  // Separating this function from main() ensures that all defer calls are
   102  // executed before os.Exit().
   103  func innerMain(ctx context.Context) error {
   104  	var socatPath string
   105  	var socketDir string
   106  	var configFile string
   107  	flag.StringVar(&socatPath, "socat", "socat", "Path to the 'socat' tool.  If not provided, PATH will be searched for one.")
   108  	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.")
   109  	flag.StringVar(&configFile, "cfg", "", "Relay configuration file (required).")
   110  	flag.Parse()
   111  	command := flag.Args()
   112  
   113  	// Load relay configuration.
   114  	if configFile == "" {
   115  		return fmt.Errorf("missing required --cfg flag")
   116  	}
   117  	cfgData, err := os.ReadFile(configFile)
   118  	if err != nil {
   119  		return err
   120  	}
   121  
   122  	var relays []*socketRelay
   123  	if err := json.Unmarshal(cfgData, &relays); err != nil {
   124  		return fmt.Errorf("failed to parse JSON file %s: %v", configFile, err)
   125  	}
   126  
   127  	// Setup a temporary directory for sockets.
   128  	if socketDir == "" {
   129  		var err error
   130  		socketDir, err = os.MkdirTemp("", "bazel_service_proxy.*")
   131  		if err != nil {
   132  			return err
   133  		}
   134  	}
   135  	defer os.RemoveAll(socketDir)
   136  
   137  	finished := make(chan int, 0) // sent after wrapped command completes
   138  	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
   139  	defer stop()
   140  	cmdErr := multiRelayWrap(ctx, relays, socketDir, socatPath, func(env []string) error {
   141  		defer func() {
   142  			close(finished)
   143  		}()
   144  		return wrapCommand(ctx, command, env)
   145  	})
   146  	// Wait for either completion or interrupt
   147  	select {
   148  	case <-finished:
   149  	case <-ctx.Done():
   150  	}
   151  	return cmdErr
   152  }
   153  
   154  func main() {
   155  	os.Exit(errorToExitCode(innerMain(context.Background())))
   156  }