cmd/buildproxywrap/relay.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.

// 'relay.go' provides a simple interface for running a socket relay
// to a remote server, using the 'socat' tool.
package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sync"

	"github.com/golang/glog"
)

// socketRelay contains the parameters for operating a single socket relay.
type socketRelay struct {
	// Identifier for this relay, used for logging.
	Name string `json:"name"`
	// Name of socket file (to be created inside a temporary directory).
	SocketFileName string `json:"socket_file_name"`
	// Name of environment variable to export, pointing to the socket file.
	SocketPathEnvVar string `json:"socket_path_env_var"`
	// Remote address of server to connect.
	ServerAddr string `json:"server_address"`

	// Internal fields follow:

	// Full path to socket file.
	socketFileFullPath string
}

// initSocketFile establishes a full path to a socket file.
// The directory 'tempdir' should already exist.
func (r *socketRelay) initSocketFile(tempDir string) (string, error) {
	r.socketFileFullPath = filepath.Join(tempDir, r.SocketFileName)
	glog.V(1).Infof("[%s] using socket file %s", r.Name, r.socketFileFullPath)
	// remove existing socket first
	err := os.RemoveAll(r.socketFileFullPath)
	return r.socketFileFullPath, err
}

// env returns the environment variable string X=Y for the socket file.
func (r *socketRelay) env() (string, error) {
	if r.socketFileFullPath == "" {
		return "", fmt.Errorf("must call initSocketFile() before env().")
	}
	env := fmt.Sprintf("%s=%s", r.SocketPathEnvVar, r.socketFileFullPath)
	glog.V(1).Infof("[%s] env %s", r.Name, env)
	return env, nil
}

// runInterruptOnCancel runs a command that is expected to be terminated
// only by an interrupt signal (Ctrl-C).
func runInterruptOnCancel(ctx context.Context, logPrefix string, command ...string) error {
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	glog.V(0).Infof("%sSpawning process", logPrefix)

	interrupted := false
	cmd.Cancel = func() error { // .Cancel requires go-1.20
		// Send SIGINT (instead of the default SIGKILL) to allow socat to
		// clean-up before gracefully exiting.
		interrupted = true
		if cmd.Process != nil {
			glog.V(0).Infof("%sShutting down with SIGINT", logPrefix)
			return cmd.Process.Signal(os.Interrupt)
		}
		return nil
	}

	err := cmd.Run()

	// Cancellation is expected.
	if interrupted {
		return nil
	}
	if err == nil {
		// Program exited before interrupt.
		return nil
	}
	if errors.Is(err, context.Canceled) {
		return nil
	}
	if exitErr, ok := err.(*exec.ExitError); ok {
		if !exitErr.Exited() { // process was terminated
			return nil
		}
	}
	return fmt.Errorf("%sprocess error: %v", logPrefix, err)
}

// Run launches a relay.  Caller should invoke this in a go-routine to run it in the background.
// 'socatPath' is the path to the 'socat' tool for the relay subprocess.
func (r *socketRelay) run(ctx context.Context, socatPath string) error {
	if r.socketFileFullPath == "" {
		return fmt.Errorf("Must call initSocketFile() before run().")
	}
	command := []string{
		socatPath,
		// For verbose debugging, pass repeated -d options.
		fmt.Sprintf("UNIX-LISTEN:%s,unlink-early,fork", r.socketFileFullPath),
		fmt.Sprintf("TCP:%s", r.ServerAddr),
	}
	return runInterruptOnCancel(ctx, fmt.Sprintf("[%s](relay) ", r.Name), command...)
}

// multiRelayWrap sets up multiple socket relay processes, sets up an
// environment, and invokes a function 'f', which usually invokes
// a wrapped subprocess using the new environment, and then shuts down
// the relay processes.
// 'relays' are the set of connections to relay through sockets.
// 'tempDir' is a writeable directory where socket files will be created.
// 'socatPath' is the location of the `socat` tool.
// 'f' is any function or subprocesses that operate while the relays are run.
// Returns the exit code of the wrapped function.
func multiRelayWrap(ctx context.Context, relays []*socketRelay, tempDir string, socatPath string, f func(env []string) error) error {
	// Initialize socket files.
	for _, r := range relays {
		if _, err := r.initSocketFile(tempDir); err != nil {
			return err
		}
	}

	// Launch relays, running them in the background.
	// Stop relays before returning.
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(ctx)
	for _, r := range relays {
		r := r
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := r.run(ctx, socatPath); err != nil {
				glog.Errorf("[%s] error running relay: %v", r.Name, err)
			}
		}()
	}

	defer func() {
		cancel()
		wg.Wait() // Wait for all relay go-routines to finish exiting.
	}()

	// Setup a modified environment for the function (usually subprocess).
	var env []string
	for _, r := range relays {
		v, err := r.env()
		if err != nil {
			return err
		}
		env = append(env, v)
	}

	// Call the wrapped function/subprocesses.
	return f(env)
}