cmd/go-cache-plugin: add optional module proxy support

The "serve" command now has flags to enable a --modproxy running as an HTTP
server alongside the build cache, which can optionally also proxy for (and
cache) Go sum database queries (with --sumdb).

cmd/go-cache-plugin/commands.go

@@ -3,12 +3,17 @@ package main
 import (
 	"context"
 	"errors"
+	"expvar"
 	"fmt"
 	"io"
 	"log"
 	"net"
+	"net/http"
 	"os"
 	"os/signal"
+	"path"
+	"path/filepath"
+	"strings"
 	"syscall"
 	"time"
 
@@ -18,7 +23,10 @@ import (
 	"github.com/creachadair/gocache"
 	"github.com/creachadair/gocache/cachedir"
 	"github.com/creachadair/taskgroup"
+	"github.com/goproxy/goproxy"
 	"github.com/tailscale/go-cache-plugin/s3cache"
+	"github.com/tailscale/go-cache-plugin/s3proxy"
+	"tailscale.com/tsweb"
 )
 
 var flags struct {
@@ -35,26 +43,26 @@ var flags struct {
 	DebugLog      bool          `flag:"debug,default=$GOCACHE_DEBUG,Enable detailed per-request debug logging (noisy)"`
 }
 
-func initCacheServer(env *command.Env) (*gocache.Server, error) {
+func initCacheServer(env *command.Env) (*gocache.Server, *s3.Client, error) {
 	switch {
 	case flags.CacheDir == "":
-		return nil, env.Usagef("you must provide a --cache-dir")
+		return nil, nil, env.Usagef("you must provide a --cache-dir")
 	case flags.S3Bucket == "":
-		return nil, env.Usagef("you must provide an S3 --bucket name")
+		return nil, nil, env.Usagef("you must provide an S3 --bucket name")
 	}
 	region, err := getBucketRegion(env.Context(), flags.S3Bucket)
 	if err != nil {
-		return nil, env.Usagef("you must provide an S3 --region name")
+		return nil, nil, env.Usagef("you must provide an S3 --region name")
 	}
 
 	dir, err := cachedir.New(flags.CacheDir)
 	if err != nil {
-		return nil, fmt.Errorf("create local cache: %w", err)
+		return nil, nil, fmt.Errorf("create local cache: %w", err)
 	}
 
 	cfg, err := config.LoadDefaultConfig(env.Context(), config.WithRegion(region))
 	if err != nil {
-		return nil, fmt.Errorf("laod AWS config: %w", err)
+		return nil, nil, fmt.Errorf("laod AWS config: %w", err)
 	}
 
 	vprintf("local cache directory: %s", flags.CacheDir)
@@ -67,6 +75,8 @@ func initCacheServer(env *command.Env) (*gocache.Server, error) {
 		MinUploadSize:     flags.MinUploadSize,
 		UploadConcurrency: flags.S3Concurrency,
 	}
+	cache.SetMetrics(env.Context(), expvar.NewMap("gocache_host"))
+
 	close := cache.Close
 	if flags.Expiration > 0 {
 		dirClose := dir.Cleanup(flags.Expiration)
@@ -83,13 +93,14 @@ func initCacheServer(env *command.Env) (*gocache.Server, error) {
 		Logf:        vprintf,
 		LogRequests: flags.DebugLog,
 	}
-	return s, nil
+	expvar.Publish("gocache_server", s.Metrics().Get("server"))
+	return s, cache.S3Client, nil
 }
 
 // runDirect runs a cache communicating on stdin/stdout, for use as a direct
 // GOCACHEPROG plugin.
 func runDirect(env *command.Env) error {
-	s, err := initCacheServer(env)
+	s, _, err := initCacheServer(env)
 	if err != nil {
 		return err
 	}
@@ -102,33 +113,35 @@ func runDirect(env *command.Env) error {
 	return nil
 }
 
-var remoteFlags struct {
-	Socket string `flag:"socket,default=$GOCACHE_SOCKET,Socket path (required)"`
+var serveFlags struct {
+	Socket   string `flag:"socket,default=$GOCACHE_SOCKET,Socket path (required)"`
+	ModProxy string `flag:"modproxy,default=$GOCACHE_MODPROXY,Module proxy service address ([host]:port)"`
+	SumDB    string `flag:"sumdb,default=$GOCACHE_SUMDB,SumDB servers to proxy for (comma-separated)"`
 }
 
 func noopClose(context.Context) error { return nil }
 
-// runRemote runs a cache communicating over a Unix-domain socket.
-func runRemote(env *command.Env) error {
-	if remoteFlags.Socket == "" {
+// runServe runs a cache communicating over a Unix-domain socket.
+func runServe(env *command.Env) error {
+	if serveFlags.Socket == "" {
 		return env.Usagef("you must provide a --socket path")
 	}
 
 	// Initialize the cache server. Unlike a direct server, only close down and
 	// wait for cache cleanup when the whole process exits.
-	s, err := initCacheServer(env)
+	s, s3c, err := initCacheServer(env)
 	if err != nil {
 		return err
 	}
 	closeHook := s.Close
 	s.Close = noopClose
 
 	// Listen for connections from the Go toolchain on the specified socket.
-	lst, err := net.Listen("unix", remoteFlags.Socket)
+	lst, err := net.Listen("unix", serveFlags.Socket)
 	if err != nil {
 		return fmt.Errorf("listen: %w", err)
 	}
-	defer os.Remove(remoteFlags.Socket) // best-effort
+	defer os.Remove(serveFlags.Socket) // best-effort
 
 	ctx, cancel := signal.NotifyContext(env.Context(), syscall.SIGINT, syscall.SIGTERM)
 	defer cancel()
@@ -138,7 +151,60 @@ func runRemote(env *command.Env) error {
 		lst.Close()
 	}()
 
+	// If a module proxy is enabled, start it.
 	var g taskgroup.Group
+	if serveFlags.ModProxy != "" {
+		modCachePath := filepath.Join(flags.CacheDir, "module")
+		if err := os.MkdirAll(modCachePath, 0700); err != nil {
+			lst.Close()
+			return fmt.Errorf("create module cache: %w", err)
+		}
+		cacher := &s3proxy.Cacher{
+			Local:       modCachePath,
+			S3Client:    s3c,
+			S3Bucket:    flags.S3Bucket,
+			KeyPrefix:   path.Join(flags.KeyPrefix, "module"),
+			MaxTasks:    flags.S3Concurrency,
+			LogRequests: flags.DebugLog,
+			Logf:        vprintf,
+		}
+		defer func() {
+			vprintf("close cacher (err=%v)", cacher.Close())
+		}()
+		proxy := &goproxy.Goproxy{
+			Fetcher: &goproxy.GoFetcher{
+				// As configured, the fetcher should never shell out to the go
+				// tool. Specifically, because we set GOPROXY and do not set any
+				// bypass via GONOPROXY, GOPRIVATE, etc., we will only attempt to
+				// proxy for the specific server(s) listed in Env.
+				GoBin: "/bin/false",
+				Env:   []string{"GOPROXY=https://proxy.golang.org"},
+			},
+			Cacher: cacher,
+		}
+		if serveFlags.SumDB != "" {
+			proxy.ProxiedSumDBs = strings.Split(serveFlags.SumDB, ",")
+			vprintf("enabling sum DB proxy for %s", strings.Join(proxy.ProxiedSumDBs, ", "))
+		}
+		expvar.Publish("modcache", cacher.Metrics())
+
+		// Run an HTTP server exporting the proxy and debug metrics.
+		mux := http.NewServeMux()
+		mux.Handle("/", proxy)
+		tsweb.Debugger(mux)
+		srv := &http.Server{
+			Addr:    serveFlags.ModProxy,
+			Handler: mux,
+		}
+		g.Go(srv.ListenAndServe)
+		vprintf("started module proxy at %q", serveFlags.ModProxy)
+		go func() {
+			<-ctx.Done()
+			vprintf("signal received, stopping module proxy")
+			srv.Shutdown(context.Background())
+		}()
+	}
+
 	for {
 		conn, err := lst.Accept()
 		if err != nil {

cmd/go-cache-plugin/go-cache-plugin.go

@@ -45,10 +45,14 @@ the --cache-dir flag or GOCACHE_DIR environment.`,
 
 In this mode, the cache server listens for connections on a socket instead of
 serving directly over stdin/stdout. The "connect" command adapts the direct
-interface to this one.`,
+interface to this one.
 
-				SetFlags: command.Flags(flax.MustBind, &remoteFlags),
-				Run:      command.Adapt(runRemote),
+By default, only the build cache is exported via the --socket path.
+If --modcache is set, the server also exports a caching module proxy at the
+specified address.`,
+
+				SetFlags: command.Flags(flax.MustBind, &serveFlags),
+				Run:      command.Adapt(runServe),
 			},
 			{
 				Name:  "connect",