dockerimg: restore go mod download functionality

Restore the go mod download functionality that was lost during the
transition to only copying git objects.

This pre-populates the Go module cache during image build time.

Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: s0941bfb3f9ba2251k
diff --git a/dockerimg/dockerimg.go b/dockerimg/dockerimg.go
index 98ce389..18284b1 100644
--- a/dockerimg/dockerimg.go
+++ b/dockerimg/dockerimg.go
@@ -806,7 +806,9 @@
 // The implementation here copies the git objects into the base image.
 // That enables fast clones into the container, because most of the git objects are already there.
 // It also avoids copying uncommitted changes, configs/hooks, etc.
-// TODO: We should also set up fake temporary Go module(s) so we can run "go mod download".
+// We also set up fake temporary Go module(s) so we can run "go mod download".
+// TODO: maybe 'go list ./...' and then do a build as well to populate the build cache.
+// TODO: 'npm install', etc? We have the rails for it.
 // This is an ok compromise, but a power user might want
 // less caching or more caching, depending on their use case. One approach we could take
 // is to punt entirely if /app/.git already exists. If the user has provided a -base-image with
@@ -824,14 +826,35 @@
 //
 // repoPath is the current working directory where sketch is being run from.
 func buildLayeredImage(ctx context.Context, imgName, baseImage, gitRoot string, verbose bool) error {
-	// Shove a bunch of git objects into the image for faster future cloning.
-	dockerfileContent := fmt.Sprintf(`FROM %s
-COPY . /git-ref
-WORKDIR /app
-# TODO: restore go.mod download
-# RUN if [ -f go.mod ]; then go mod download; fi
-CMD ["/bin/sketch"]
-`, baseImage)
+	goModules, err := collectGoModules(ctx, gitRoot)
+	if err != nil {
+		return fmt.Errorf("failed to collect go modules: %w", err)
+	}
+
+	buf := new(strings.Builder)
+	line := func(msg string, args ...any) {
+		fmt.Fprintf(buf, msg+"\n", args...)
+	}
+
+	line("FROM %s", baseImage)
+	line("COPY . /git-ref")
+
+	for _, module := range goModules {
+		line("RUN mkdir -p /go-module")
+		line("RUN git --git-dir=/git-ref --work-tree=/go-module cat-file blob %s > /go-module/go.mod", module.modSHA)
+		if module.sumSHA != "" {
+			line("RUN git --git-dir=/git-ref --work-tree=/go-module cat-file blob %s > /go-module/go.sum", module.sumSHA)
+		}
+		// drop any replaced modules
+		line("RUN cd /go-module && go mod edit -json | jq -r '.Replace? // [] | .[] | .Old.Path' | xargs -r -I{} go mod edit -dropreplace={} -droprequire={}")
+		// grab what’s left, best effort only to avoid breaking on (say) private modules
+		line("RUN cd /go-module && go mod download || true")
+		line("RUN rm -rf /go-module")
+	}
+
+	line("WORKDIR /app")
+	line(`CMD ["/bin/sketch"]`)
+	dockerfileContent := buf.String()
 
 	// Create a temporary directory for the Dockerfile
 	tmpDir, err := os.MkdirTemp("", "sketch-docker-*")
@@ -952,6 +975,62 @@
 	return gitCommonDir, nil
 }
 
+// goModuleInfo represents a Go module with its file paths and blob SHAs
+type goModuleInfo struct {
+	// modPath is the path to the go.mod file, for debugging
+	modPath string
+	// modSHA is the git blob SHA of the go.mod file
+	modSHA string
+	// sumSHA is the git blob SHA of the go.sum file, empty if no go.sum exists
+	sumSHA string
+}
+
+// collectGoModules returns all go.mod files in the git repository with their blob SHAs.
+func collectGoModules(ctx context.Context, gitRoot string) ([]goModuleInfo, error) {
+	cmd := exec.CommandContext(ctx, "git", "ls-files", "-z", "*.mod")
+	cmd.Dir = gitRoot
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		return nil, fmt.Errorf("git ls-files -z *.mod: %s: %w", out, err)
+	}
+
+	modFiles := strings.Split(string(out), "\x00")
+	var modules []goModuleInfo
+	for _, file := range modFiles {
+		if filepath.Base(file) != "go.mod" {
+			continue
+		}
+
+		modSHA, err := getGitBlobSHA(ctx, gitRoot, file)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get blob SHA for %s: %w", file, err)
+		}
+
+		// If corresponding go.sum exists, get its SHA
+		sumFile := filepath.Join(filepath.Dir(file), "go.sum")
+		sumSHA, _ := getGitBlobSHA(ctx, gitRoot, sumFile) // best effort
+
+		modules = append(modules, goModuleInfo{
+			modPath: file,
+			modSHA:  modSHA,
+			sumSHA:  sumSHA,
+		})
+	}
+
+	return modules, nil
+}
+
+// getGitBlobSHA returns the git blob SHA for a file at HEAD
+func getGitBlobSHA(ctx context.Context, gitRoot, filePath string) (string, error) {
+	cmd := exec.CommandContext(ctx, "git", "rev-parse", "HEAD:"+filePath)
+	cmd.Dir = gitRoot
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		return "", fmt.Errorf("git rev-parse HEAD:%s: %s: %w", filePath, out, err)
+	}
+	return strings.TrimSpace(string(out)), nil
+}
+
 // getEnvForwardingFromGitConfig retrieves environment variables to pass through to Docker
 // from git config using the sketch.envfwd multi-valued key.
 func getEnvForwardingFromGitConfig(ctx context.Context) []string {
diff --git a/dockerimg/dockerimg_test.go b/dockerimg/dockerimg_test.go
index 2b98c25..e8a5d5a 100644
--- a/dockerimg/dockerimg_test.go
+++ b/dockerimg/dockerimg_test.go
@@ -6,6 +6,7 @@
 	"crypto/sha256"
 	"encoding/hex"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"testing"
 )
@@ -148,3 +149,154 @@
 		t.Error("Different content should produce different hash")
 	}
 }
+
+func TestCollectGoModules(t *testing.T) {
+	// Create a temporary directory with test files
+	tempDir := t.TempDir()
+
+	// Initialize a git repository
+	cmd := exec.Command("git", "init", ".")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to init git repo: %v", err)
+	}
+
+	// Create test go.mod files
+	modContent := "module test\n\ngo 1.19\n"
+	sumContent := "example.com/test v1.0.0 h1:abc\n"
+
+	// Root go.mod
+	if err := os.WriteFile(filepath.Join(tempDir, "go.mod"), []byte(modContent), 0o644); err != nil {
+		t.Fatalf("Failed to create go.mod: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(tempDir, "go.sum"), []byte(sumContent), 0o644); err != nil {
+		t.Fatalf("Failed to create go.sum: %v", err)
+	}
+
+	// Subdirectory go.mod
+	subDir := filepath.Join(tempDir, "subdir")
+	if err := os.MkdirAll(subDir, 0o755); err != nil {
+		t.Fatalf("Failed to create subdir: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(subDir, "go.mod"), []byte(modContent), 0o644); err != nil {
+		t.Fatalf("Failed to create subdir/go.mod: %v", err)
+	}
+	// No go.sum for subdir to test the case where go.sum is missing
+
+	// Add files to git
+	cmd = exec.Command("git", "add", ".")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to add files to git: %v", err)
+	}
+
+	// Configure git user for the test repo
+	cmd = exec.Command("git", "config", "user.email", "test@example.com")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to set git user email: %v", err)
+	}
+	cmd = exec.Command("git", "config", "user.name", "Test User")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to set git user name: %v", err)
+	}
+
+	// Commit the files
+	cmd = exec.Command("git", "commit", "-m", "test commit")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to commit files: %v", err)
+	}
+
+	// Collect go modules
+	ctx := context.Background()
+	modules, err := collectGoModules(ctx, tempDir)
+	if err != nil {
+		t.Fatalf("collectGoModules failed: %v", err)
+	}
+
+	// Verify results
+	if len(modules) != 2 {
+		t.Fatalf("Expected 2 modules, got %d", len(modules))
+	}
+
+	// Check root module
+	root := modules[0]
+	if root.modPath != "go.mod" {
+		t.Errorf("Expected root modPath to be 'go.mod', got %s", root.modPath)
+	}
+	if root.modSHA == "" {
+		t.Errorf("Expected root modSHA to be non-empty")
+	}
+	if root.sumSHA == "" {
+		t.Errorf("Expected root sumSHA to be non-empty")
+	}
+
+	// Check subdir module
+	sub := modules[1]
+	if sub.modPath != "subdir/go.mod" {
+		t.Errorf("Expected subdir modPath to be 'subdir/go.mod', got %s", sub.modPath)
+	}
+	if sub.modSHA == "" {
+		t.Errorf("Expected subdir modSHA to be non-empty")
+	}
+	if sub.sumSHA != "" {
+		t.Errorf("Expected subdir sumSHA to be empty, got %s", sub.sumSHA)
+	}
+}
+
+func TestCollectGoModulesNoModFiles(t *testing.T) {
+	// Create a temporary directory with no go.mod files
+	tempDir := t.TempDir()
+
+	// Initialize a git repository
+	cmd := exec.Command("git", "init", ".")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to init git repo: %v", err)
+	}
+
+	// Create a non-go.mod file
+	if err := os.WriteFile(filepath.Join(tempDir, "README.md"), []byte("# Test"), 0o644); err != nil {
+		t.Fatalf("Failed to create README.md: %v", err)
+	}
+
+	// Add files to git
+	cmd = exec.Command("git", "add", ".")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to add files to git: %v", err)
+	}
+
+	// Configure git user for the test repo
+	cmd = exec.Command("git", "config", "user.email", "test@example.com")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to set git user email: %v", err)
+	}
+	cmd = exec.Command("git", "config", "user.name", "Test User")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to set git user name: %v", err)
+	}
+
+	// Commit the files
+	cmd = exec.Command("git", "commit", "-m", "test commit")
+	cmd.Dir = tempDir
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Failed to commit files: %v", err)
+	}
+
+	// Collect go modules
+	ctx := context.Background()
+	modules, err := collectGoModules(ctx, tempDir)
+	if err != nil {
+		t.Fatalf("collectGoModules failed: %v", err)
+	}
+
+	// Verify no modules found
+	if len(modules) != 0 {
+		t.Fatalf("Expected 0 modules, got %d", len(modules))
+	}
+}