go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/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"
	"io/fs"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"

	"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

	// Signal that relay process is ready to accept connections.
	ready chan struct{}
}

const (
	socketPollInterval   = 1 * time.Second
	socketCreateTimeout  = 5 * time.Second
	socketConnectTimeout = 5 * time.Second
)

// initSocketFile establishes a full path to a socket file.
// The directory 'tempdir' should already exist.
func (r *socketRelay) initSocketFile(tempDir string) (string, error) {
	if _, err := os.Stat(tempDir); errors.Is(err, fs.ErrNotExist) {
		return "", err
	}
	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)
	r.ready = make(chan struct{})
	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).
// `waitAfterLaunch` is a function to wait for a certain condition before
// resuming, for example, waiting for a socket to bind.  If no wait is needed,
// pass a function that just returns nil right away.
func runInterruptOnCancel(ctx context.Context, logPrefix string, waitAfterLaunch func() error, command ...string) error {
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	glog.V(0).Infof("%sSpawning process... %v", logPrefix, command)

	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
	}

	// Run process in the background.
	if err := cmd.Start(); err != nil {
		return err
	}
	glog.V(0).Infof("%sProcess has started.", logPrefix)

	// Wait for a condition before signaling back to the caller
	// that the relay process is ready to accept connections.
	if waitErr := waitAfterLaunch(); waitErr != nil {
		return waitErr
	}

	err := cmd.Wait()
	glog.V(0).Infof("%sProcess finished.", logPrefix)

	// Cancellation is expected.
	if interrupted {
		glog.V(1).Infof("%sInterrupted after work was completed.", logPrefix)
		return nil
	}
	if err == nil {
		// Program exited before interrupt.
		glog.V(1).Infof("%sProgram exited itself.", logPrefix)
		return nil
	}
	if errors.Is(err, context.Canceled) {
		glog.V(1).Infof("%sProgram's context was canceled.", logPrefix)
		return nil
	}
	if exitErr, ok := err.(*exec.ExitError); ok {
		if !exitErr.Exited() { // process was terminated
			glog.V(1).Infof("%sProgram was terminated.", logPrefix)
			return nil
		}
	}
	return fmt.Errorf("%sprocess error: %v", logPrefix, err)
}

// waitForSocketBind() waits for a socket to appear and be connect-able.
// `name` is just an identifier for diagnostics.
// `socketPath` is the path of the socket to wait for.
// An alternative would be to use something like inotify.
func waitForSocketBind(name string, socketPath string) error {
	// Wait for socket to be connect-able before returning.
	glog.V(1).Infof("[%s] waiting for socket to bind: %s", name, socketPath)
	c := make(chan struct{})
	go func() {
		for {
			if _, err := os.Stat(socketPath); err == nil {
				glog.V(1).Infof("[%s] socket now exists", name)
				break
			} else {
				glog.V(2).Infof("[%s] waiting for socket to bind...", name)
				time.Sleep(socketPollInterval)
			}
		}
		c <- struct{}{}
	}()

	// Wait for socket with timeout.
	select {
	case <-c:
		break
	case <-time.After(socketCreateTimeout):
		return fmt.Errorf("Socket file %s not found after %v", socketPath, socketCreateTimeout)
	}

	// Ensure that connecting to the socket works.
	glog.V(1).Infof("[%s] Dialing socket: %s", name, socketPath)
	conn, err := net.DialTimeout("unix", socketPath, socketConnectTimeout)
	if err != nil {
		return err
	}
	glog.V(1).Infof("[%s] Closing connection on socket: %s", name, socketPath)
	return conn.Close()
}

// 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),
	}
	glog.V(1).Infof("[%s] launching: %v", r.Name, command)
	return runInterruptOnCancel(ctx, fmt.Sprintf("[%s](relay) ", r.Name),
		func() error {
			err := waitForSocketBind(r.Name, r.socketFileFullPath)
			if err == nil {
				glog.V(1).Infof("[%s] After launching, relay is now ready to accept connections", r.Name)
				r.ready <- struct{}{}
				glog.V(1).Infof("[%s] Sent relay ready signal", r.Name)
			}
			return err
		},
		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
		}
	}

	// Setup wait for all relays to be ready to accept connections.
	var readyWg sync.WaitGroup
	for _, r := range relays {
		readyWg.Add(1)
		go func() {
			defer readyWg.Done()
			glog.V(2).Infof("[%s] Waiting for relay to be ready...", r.Name)
			<-r.ready
			glog.V(2).Infof("[%s] Received relay ready signal.", r.Name)
		}()
	}

	// Launch relays, running them in the background.
	// Stop relays before returning.
	glog.V(1).Infof("Starting relays in the background.")
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(ctx)
	for _, r := range relays {
		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).
	glog.V(1).Infof("Preparing environment for wrapped command")
	var env []string
	for _, r := range relays {
		v, err := r.env()
		if err != nil {
			return err
		}
		env = append(env, v)
	}
	glog.V(1).Infof("environment: %v ", env)

	// Wait for all relays to be ready to accept connections through sockets.
	glog.V(1).Infof("Waiting for all relays to be ready.")
	readyWg.Wait()
	glog.V(1).Infof("All relays are ready.")

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