go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/gcsproxy/main.go

// Copyright 2019 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.

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"golang.org/x/time/rate"
)

const (
	// OAuth2 scope for reading GCS.
	// See https://cloud.google.com/storage/docs/authentication.
	gcsReadOnlyScope = "https://www.googleapis.com/auth/devstorage.read_only"

	// We want to allow at most `maxNumRequests` per IP to be serviced every
	// `refreshWindowMs` milliseconds.
	// In terms of the token bucket underlying rate.Limiter, this translates to a
	// new token refreshed every `tokenRefreshRate` with a burst size of
	// `burstSize`: this allows for a token pool in which we can check Allow()
	// and Wait() at the desired rates without consuming any tokens reserved for
	// servicing of other requests.
	// See https://godoc.org/golang.org/x/time/rate#Limiter for more details.
	maxNumRequests        = 20
	refreshWindowMs       = 200
	tokenBudgetPerRequest = 2
	tokenRefreshRate      = (refreshWindowMs / (tokenBudgetPerRequest * maxNumRequests)) * time.Millisecond
	tokenBurstSize        = 2 * maxNumRequests

	// Constants for retrying communication with GCS.
	retryAttempts = 10

	// finishTimeout is the timeout we allow the subprocess to complete in
	// after we send a SIGTERM.
	finishTimeout = 10 * time.Second
)

var (
	credentialsFile  string
	port             string
	allowedAddrsFile string
	// Only a var for testability.
	// See https://cloud.google.com/storage/docs/request-endpoints.
	gcsHost = "storage.googleapis.com"
	// Only a var for testability.
	retryBackoff = 100 * time.Millisecond
)

func usage() {
	fmt.Printf(`gcsproxy [flags] [subcommand]

Starts a proxy server that forwards requests to GCS with authentication.
If positional arguments are provided, they will be run as subprocess and
the lifetime of the server will be scoped to the lifetime of that process.
`)
}

func init() {
	flag.Usage = usage
	flag.StringVar(&credentialsFile, "credentials", "", "path to a credentials file in the Google Credentials File format; if none provided, default application credentials will be used.")
	flag.StringVar(&port, "port", "", "port at which the server should listen.")
	flag.StringVar(&allowedAddrsFile, "allowed", "", "a flat JSON list of remote addresses allowed to make requests of the proxy server; if not provided, all addresses will be allowed")
}

func main() {
	flag.Parse()

	log.SetPrefix("gcsproxy: ")
	log.SetFlags(log.Lmsgprefix | log.Ltime | log.Lmicroseconds | log.Lshortfile)

	// For a graceful teardown in the event of a canceling signal.
	ctx, cancel := context.WithCancel(context.Background())
	signals := make(chan os.Signal)
	defer close(signals)
	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)

	go func() {
		select {
		case <-signals:
			cancel()
		}
	}()

	if err := execute(ctx, credentialsFile, port, allowedAddrsFile, flag.Args()); err != nil {
		log.Fatal(err)
	}
}

func runSubcommand(ctx context.Context, subCmd []string) error {
	cmd := exec.Command(subCmd[0], subCmd[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// Set a process group ID so we can kill the entire group,
	// meaning the process and any of its children.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	// Spin off handler to exit subprocesses cleanly via SIGTERM.
	processDone := make(chan struct{})
	var processMu sync.Mutex
	go func() {
		select {
		case <-processDone:
		case <-ctx.Done():
			// We need to check if the process is nil because it won't exist if
			// it has been SIGKILL'd already.
			processMu.Lock()
			defer processMu.Unlock()
			if cmd.Process != nil {
				if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
					log.Printf("exited cmd with error %s", err)
				}
				// If the subprocess doesn't complete within the finishTimeout,
				// send a SIGKILL to force it to exit.
				go func() {
					select {
					case <-processDone:
					case <-time.After(finishTimeout):
						log.Printf("killing process %d", cmd.Process.Pid)
						// Negating the process ID means interpret it as a process group ID, so
						// we kill the subprocess and all of its children.
						if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
							log.Printf("killed cmd with error %s", err)
						}
					}
				}()
			}
		}
	}()

	// Ensure that the context still exists before running the subprocess.
	if ctx.Err() != nil {
		log.Print("context exited before starting subprocess")
		return ctx.Err()
	}

	// We need to make this a critical section because running Start changes
	// cmd.Process, which we attempt to access in the goroutine above. Not locking
	// causes a data race.
	processMu.Lock()
	log.Printf("starting: %s", subCmd)
	err := cmd.Start()
	processMu.Unlock()
	if err != nil {
		close(processDone)
		return err
	}
	// Since we wait for the command to complete even if we send a SIGTERM when the
	// context is canceled, it is up to the underlying command to exit with the
	// proper exit code after handling a SIGTERM.
	err = cmd.Wait()
	close(processDone)
	if err != nil && ctx.Err() != nil {
		// If the subprocess was terminated early and completed without terminating
		// its child processes, kill any remaining processes in the group.
		log.Printf("killing process %d", cmd.Process.Pid)
		syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		return ctx.Err()
	}
	return err
}

// Execute starts the proxy server.
func execute(ctx context.Context, credFile, port, addrsFile string, subCmd []string) error {
	if port == "" {
		return fmt.Errorf("-port is required")
	}

	client, err := httpClient(ctx, credFile)
	if err != nil {
		return err
	}

	var allowedAddrs []string
	if allowedAddrsFile != "" {
		b, err := os.ReadFile(allowedAddrsFile)
		if err != nil {
			return err
		}
		if err := json.Unmarshal(b, &allowedAddrs); err != nil {
			return err
		}
	}

	limiters := new(sync.Map)
	for _, addr := range allowedAddrs {
		limiters.Store(addr, newLimiter())
	}

	outdir := os.Getenv("FUCHSIA_TEST_OUTDIR")
	if outdir == "" {
		log.Printf("GCS Proxy could not find FUCHSIA_TEST_OUTDIR, using %s instead", os.TempDir())
		outdir = os.TempDir()
	}
	metricsFile := path.Join(outdir, "gcsproxy_metrics")
	log.Printf("GCS Proxy will log metrics to %s", metricsFile)

	redirect := &redirectHandler{
		client:        client,
		restrictAddrs: allowedAddrs != nil,
		limiters:      limiters,
		metricsFile:   metricsFile,
	}

	mux := http.NewServeMux()
	mux.Handle("/", redirect)
	ln, err := net.Listen("tcp", fmt.Sprintf(":%s", port))
	if err != nil {
		return err
	}
	s := http.Server{Handler: mux}

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Printf("starting a GCS proxy server at %s", ln.Addr().String())
		if err := s.Serve(ln); err != http.ErrServerClosed {
			log.Fatal(err)
		}
	}()
	defer wg.Wait()

	defer func() {
		log.Printf("Shutting down GCS proxy server at %s", ln.Addr().String())
		if err := s.Shutdown(ctx); err != nil {
			log.Fatal(err)
		}
	}()

	if len(subCmd) == 0 {
		log.Printf("Press Ctrl-C to abort")
		<-ctx.Done()
	}
	return runSubcommand(ctx, subCmd)
}

func newLimiter() *rate.Limiter {
	limit := rate.Every(tokenRefreshRate)
	return rate.NewLimiter(limit, tokenBurstSize)
}

// sleep sleeps until the context is done or d time elapsed.
//
// Returns true if the context was done.
func sleep(ctx context.Context, d time.Duration) bool {
	timer := time.NewTimer(d)
	defer timer.Stop()

	select {
	case <-ctx.Done():
		return true
	case <-timer.C:
		return false
	}
}

// RedirectHandler is a simple handler that redirects requests to GCS.
type redirectHandler struct {
	client *http.Client
	// Whether to only serve to addresses present in the limiter map.
	restrictAddrs bool
	// limiters is a map of string: *rate.Limiter.
	limiters *sync.Map
	// metricsFile is a path to a file to write metrics to.
	metricsFile string
}

func shouldRetry(statusCode int) bool {
	// We suspect we're seeing flaky NotFound errors, so retry in that case. See fxbug.dev/65035.
	return (statusCode == http.StatusNotFound ||
		// Suggested by https://cloud.google.com/storage/docs/request-rate#ramp-up
		statusCode == http.StatusTooManyRequests ||
		statusCode >= http.StatusInternalServerError)
}

func (h *redirectHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if req == nil {
		http.Error(w, "cannot handle nil request", http.StatusInternalServerError) // = 501
		return
	}
	ctx := context.Background()
	reqReceived := time.Now()

	limiter, ok := h.getLimiter(req.RemoteAddr)
	if !ok {
		http.Error(
			w,
			fmt.Sprintf("address %q is not authorized to make requests", req.RemoteAddr),
			http.StatusForbidden, // = 403
		)
		return
	}

	startedLimiterWait := time.Now()
	if !limiter.Allow() {
		if err := limiter.Wait(ctx); err != nil {
			http.Error(w, fmt.Sprintf("rate-limiting error: %s", err), http.StatusInternalServerError)
		}
	}
	limiterLatency := time.Since(startedLimiterWait).Milliseconds()

	req.Host = gcsHost
	req.URL.Host = gcsHost
	req.URL.Scheme = "https"
	// It is an error to set this field in an HTTP client request
	// See https://golang.org/pkg/net/http/#Request.
	req.RequestURI = ""

	var resp *http.Response
	var err error
	loggedReq := false
	var gcsLatency int64
	var attempt int
	var backoffLatency int64
	for attempt = 0; attempt < retryAttempts; attempt++ {
		startedGCSReq := time.Now()
		resp, err = h.client.Do(req)
		if err == nil {
			gcsLatency += time.Since(startedGCSReq).Milliseconds()
		}
		if err == nil && !shouldRetry(resp.StatusCode) || ctx.Err() != nil || attempt+1 == retryAttempts {
			break
		}
		if !loggedReq {
			var msg string
			if err != nil {
				msg = fmt.Sprintf("unexpected error: %s", err)
			} else if resp != nil {
				msg = fmt.Sprintf("unexpected response: %s", resp.Status)
			} else {
				msg = "both error and response are nil"
			}
			log.Printf("%s. Retrying request. URL: %s Header: %s", msg, req.URL, req.Header)
			loggedReq = true
		}
		backoff := retryBackoff * time.Duration(math.Pow(1.5, float64(attempt)))
		backoffLatency += backoff.Milliseconds()
		if sleep(ctx, backoff) {
			break
		}
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	} else if resp == nil {
		http.Error(w, "received a nil response", http.StatusInternalServerError)
	}

	startResponseTime := time.Now()
	for k, v := range resp.Header {
		for _, s := range v {
			w.Header().Add(k, s)
		}
	}
	w.WriteHeader(resp.StatusCode)

	blobSize, err := io.Copy(w, resp.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	metrics := fmt.Sprintf(
		"%s, %s, %d, %d, %d, %d, %d, %d, %d\n",
		reqReceived.String(),
		req.URL.String(),
		blobSize,
		time.Since(reqReceived).Milliseconds(),
		limiterLatency,
		gcsLatency,
		backoffLatency,
		time.Since(startResponseTime).Milliseconds(),
		attempt+1,
	)
	f, err := os.OpenFile(h.metricsFile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0o755)
	if err != nil {
		log.Printf("failed to open metrics file: %s", err)
		return
	}
	defer f.Close()
	f.WriteString(metrics)
}

// getLimiter returns the limiter associated with a given address and whether
// the address is allowed to make requests.
func (h *redirectHandler) getLimiter(addr string) (*rate.Limiter, bool) {
	var limiter any
	var ok bool
	if limiter, ok = h.limiters.Load(addr); !ok {
		// While the full address may have been in the `-allowed` list, its
		// hostname may be; in that case, dynamically add a limiter for that address.
		hostname := strings.Split(addr, ":")[0]
		_, hostOK := h.limiters.Load(hostname)

		if !h.restrictAddrs || hostOK {
			limiter = newLimiter()
			h.limiters.Store(addr, limiter)
			ok = true
		}
	}
	if ok {
		return limiter.(*rate.Limiter), ok
	}
	return nil, false
}

// Returns an HTTP client with the credentials to read from GCS. If no
// credential file is supplied, then the default application credentials
// will be used.
func httpClient(ctx context.Context, credFile string) (*http.Client, error) {
	var creds *google.Credentials
	var err error
	if credFile == "" {
		creds, err = google.FindDefaultCredentials(ctx, gcsReadOnlyScope)
		if err != nil {
			return nil, fmt.Errorf("failed to find default credentials: %w", err)
		}
	} else {
		contents, err := os.ReadFile(credFile)
		if err != nil {
			return nil, err
		}
		creds, err = google.CredentialsFromJSON(ctx, contents, gcsReadOnlyScope)
		if err != nil {
			return nil, fmt.Errorf("failed to derive the derive credentials from supplied file: %w", err)
		}
	}
	return oauth2.NewClient(ctx, creds.TokenSource), nil
}