cmd/go: restore go.mod files during toolchain selection

They have to be renamed to _go.mod to make a valid module.
Copy them back to go.mod so that 'go test cmd' has a better
chance of working.

For #57001.

Change-Id: Ied6f0dd77928996ab322a55c5606d7f75431e362
Reviewed-on: https://go-review.googlesource.com/c/go/+/504118
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Bryan Mills <bcmills@google.com>
Auto-Submit: Russ Cox <rsc@golang.org>
Run-TryBot: Russ Cox <rsc@golang.org>

3 files changed, 134 insertions, 0 deletions

diff --git a/src/cmd/go/internal/toolchain/select.go b/src/cmd/go/internal/toolchain/select.go
index 8eac03b339..8b1a0b94be 100644
--- a/src/cmd/go/internal/toolchain/select.go
+++ b/src/cmd/go/internal/toolchain/select.go
@@ -245,6 +245,8 @@ var TestVersionSwitch string
 func Exec(gotoolchain string) {
 	log.SetPrefix("go: ")
 
+	writeBits = sysWriteBits()
+
 	count, _ := strconv.Atoi(os.Getenv(countEnv))
 	if count >= maxSwitch-10 {
 		fmt.Fprintf(os.Stderr, "go: switching from go%v to %v [depth %d]\n", gover.Local(), gotoolchain, count)
@@ -357,10 +359,101 @@ func Exec(gotoolchain string) {
 		}
 	}
 
+	srcUGoMod := filepath.Join(dir, "src/_go.mod")
+	srcGoMod := filepath.Join(dir, "src/go.mod")
+	if size(srcGoMod) != size(srcUGoMod) {
+		err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+			if path == srcUGoMod {
+				// Leave for last, in case we are racing with another go command.
+				return nil
+			}
+			if pdir, name := filepath.Split(path); name == "_go.mod" {
+				if err := raceSafeCopy(path, pdir+"go.mod"); err != nil {
+					return err
+				}
+			}
+			return nil
+		})
+		// Handle src/go.mod; this is the signal to other racing go commands
+		// that everything is okay and they can skip this step.
+		if err == nil {
+			err = raceSafeCopy(srcUGoMod, srcGoMod)
+		}
+		if err != nil {
+			base.Fatalf("download %s: %v", gotoolchain, err)
+		}
+	}
+
 	// Reinvoke the go command.
 	execGoToolchain(gotoolchain, dir, filepath.Join(dir, "bin/go"))
 }
 
+func size(path string) int64 {
+	info, err := os.Stat(path)
+	if err != nil {
+		return -1
+	}
+	return info.Size()
+}
+
+var writeBits fs.FileMode
+
+// raceSafeCopy copies the file old to the file new, being careful to ensure
+// that if multiple go commands call raceSafeCopy(old, new) at the same time,
+// they don't interfere with each other: both will succeed and return and
+// later observe the correct content in new. Like in the build cache, we arrange
+// this by opening new without truncation and then writing the content.
+// Both go commands can do this simultaneously and will write the same thing
+// (old never changes content).
+func raceSafeCopy(old, new string) error {
+	oldInfo, err := os.Stat(old)
+	if err != nil {
+		return err
+	}
+	newInfo, err := os.Stat(new)
+	if err == nil && newInfo.Size() == oldInfo.Size() {
+		return nil
+	}
+	data, err := os.ReadFile(old)
+	if err != nil {
+		return err
+	}
+	// The module cache has unwritable directories by default.
+	// Restore the user write bit in the directory so we can create
+	// the new go.mod file. We clear it again at the end on a
+	// best-effort basis (ignoring failures).
+	dir := filepath.Dir(old)
+	info, err := os.Stat(dir)
+	if err != nil {
+		return err
+	}
+	if err := os.Chmod(dir, info.Mode()|writeBits); err != nil {
+		return err
+	}
+	defer os.Chmod(dir, info.Mode())
+	// Note: create the file writable, so that a racing go command
+	// doesn't get an error before we store the actual data.
+	f, err := os.OpenFile(new, os.O_CREATE|os.O_WRONLY, writeBits&^0o111)
+	if err != nil {
+		// If OpenFile failed because a racing go command completed our work
+		// (and then OpenFile failed because the directory or file is now read-only),
+		// count that as a success.
+		if size(old) == size(new) {
+			return nil
+		}
+		return err
+	}
+	defer os.Chmod(new, oldInfo.Mode())
+	if _, err := f.Write(data); err != nil {
+		f.Close()
+		return err
+	}
+	return f.Close()
+}
+
 // modGoToolchain finds the enclosing go.work or go.mod file
 // and returns the go version and toolchain lines from the file.
 // The toolchain line overrides the version line
diff --git a/src/cmd/go/internal/toolchain/umask_none.go b/src/cmd/go/internal/toolchain/umask_none.go
new file mode 100644
index 0000000000..b092fe8b7d
--- /dev/null
+++ b/src/cmd/go/internal/toolchain/umask_none.go
@@ -0,0 +1,13 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !(darwin || freebsd || linux || netbsd || openbsd)
+
+package toolchain
+
+import "io/fs"
+
+func sysWriteBits() fs.FileMode {
+	return 0700
+}
diff --git a/src/cmd/go/internal/toolchain/umask_unix.go b/src/cmd/go/internal/toolchain/umask_unix.go
new file mode 100644
index 0000000000..cbe4307311
--- /dev/null
+++ b/src/cmd/go/internal/toolchain/umask_unix.go
@@ -0,0 +1,28 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build darwin || freebsd || linux || netbsd || openbsd
+
+package toolchain
+
+import (
+	"io/fs"
+	"syscall"
+)
+
+// sysWriteBits determines which bits to OR into the mode to make a directory writable.
+// It must be called when there are no other file system operations happening.
+func sysWriteBits() fs.FileMode {
+	// Read current umask. There's no way to read it without also setting it,
+	// so set it conservatively and then restore the original one.
+	m := syscall.Umask(0o777)
+	syscall.Umask(m)    // restore bits
+	if m&0o22 == 0o22 { // group and world are unwritable by default
+		return 0o700
+	}
+	if m&0o2 == 0o2 { // group is writable by default, but not world
+		return 0o770
+	}
+	return 0o777 // everything is writable by default
+}