gomodfs: add go.mod copies to golang.org/toolchain modules

gomodfs.go

@@ -38,6 +38,7 @@ import (
 	"golang.org/x/mod/sumdb/dirhash"
 	"golang.org/x/sync/singleflight"
 	"tailscale.com/types/result"
+	"tailscale.com/util/set"
 	"tailscale.com/util/testenv"
 )
 
@@ -335,10 +336,52 @@ func (fs *FS) downloadZip(ctx context.Context, mv store.ModuleVersion) (store.Mo
 			zf:   f,
 		})
 	}
+	if mv.Module == toolchainModule {
+		put.Files = addToolchainGoMods(put.Files)
+	}
 
 	return fs.Store.PutModule(ctx, mv, put)
 }
 
+// toolchainModule is the module path of the Go toolchain zips that cmd/go
+// downloads for GOTOOLCHAIN switching. See https://go.dev/doc/toolchain.
+const toolchainModule = "golang.org/toolchain"
+
+// addToolchainGoMods returns files with a "go.mod" copy added beside
+// each "_go.mod" file that lacks one.
+//
+// Toolchain zips ship their go.mod files renamed to _go.mod (so the module
+// zip doesn't contain nested modules), and after downloading a toolchain,
+// cmd/go copies each _go.mod back to go.mod in the module cache. That
+// requires making module cache directories writable, which fails on our
+// read-only filesystem. But cmd/go skips that step if src/go.mod already
+// exists with the same size as src/_go.mod, so we provide the copies
+// ourselves.
+func addToolchainGoMods(files []store.PutFile) []store.PutFile {
+	have := set.Set[string]{}
+	for _, f := range files {
+		have.Add(f.Path())
+	}
+	for _, f := range files {
+		dir, ok := strings.CutSuffix(f.Path(), "_go.mod")
+		if !ok || (dir != "" && !strings.HasSuffix(dir, "/")) {
+			continue
+		}
+		if !have.Contains(dir + "go.mod") {
+			files = append(files, renamedPutFile{f, dir + "go.mod"})
+		}
+	}
+	return files
+}
+
+// renamedPutFile is a PutFile with a different path.
+type renamedPutFile struct {
+	store.PutFile
+	path string
+}
+
+func (f renamedPutFile) Path() string { return f.path }
+
 // downloadZipSet downloads a module version's zip, info, and mod
 // files from the module proxy URL prefix baseURL.
 func (fs *FS) downloadZipSet(ctx context.Context, baseURL string) (map[string][]byte, error) {
@@ -463,6 +506,13 @@ func (fs *FS) getZipRoot(ctx context.Context, mv store.ModuleVersion) (mh store.
 		}
 
 		root, err := fs.Store.GetZipRoot(ctx, mv)
+		if err == nil && mv.Module == toolchainModule && !fs.ReadOnly {
+			// Toolchains stored before addToolchainGoMods existed lack
+			// the go.mod copies that cmd/go needs. Re-download them.
+			if _, serr := fs.Store.Stat(ctx, root, "src/go.mod"); errors.Is(serr, os.ErrNotExist) {
+				err = store.ErrCacheMiss
+			}
+		}
 		if err == nil {
 			fs.setZipRootCache(mv, root)
 			return root, nil

gomodfs_test.go

@@ -5,8 +5,10 @@ package gomodfs
 
 import (
 	"fmt"
+	"io"
 	"net/http"
 	"os"
+	"slices"
 	"strings"
 	"testing"
 	"time"
@@ -182,3 +184,32 @@ func TestModuleProxyURLs(t *testing.T) {
 		}
 	}
 }
+
+type fakePutFile string
+
+func (f fakePutFile) Path() string                 { return string(f) }
+func (f fakePutFile) Mode() os.FileMode            { return 0644 }
+func (f fakePutFile) Size() int64                  { return 0 }
+func (f fakePutFile) Open() (io.ReadCloser, error) { return io.NopCloser(strings.NewReader("")), nil }
+
+func TestAddToolchainGoMods(t *testing.T) {
+	var in []store.PutFile
+	for _, p := range []string{
+		"_go.mod",
+		"src/_go.mod",
+		"src/cmd/_go.mod",
+		"src/cmd/go.mod", // already present; no copy added
+		"src/foo/x_go.mod",
+		"src/main.go",
+	} {
+		in = append(in, fakePutFile(p))
+	}
+	var got []string
+	for _, f := range addToolchainGoMods(in)[len(in):] {
+		got = append(got, f.Path())
+	}
+	want := []string{"go.mod", "src/go.mod"}
+	if !slices.Equal(got, want) {
+		t.Errorf("added %q; want %q", got, want)
+	}
+}