generator/gomod: support GOPROXY build arg

Add DALEC_GOMOD_PROXY as a known frontend build arg and carry it through
SourceOpts into gomod dependency generation and gomod patch preprocessing.

Signed-off-by: Kartik Joshi <kartikjoshi@microsoft.com>

Author: kartikjoshi21

determinism_test.go

@@ -222,7 +222,7 @@ func TestGomodPatchEnvProducesDeterministicLLB(t *testing.T) {
 	worker := llb.Scratch()
 
 	build := func() llb.State {
-		st, err := spec.generateGomodPatchStateForSource("src", gen, base, worker, nil)
+		st, err := spec.generateGomodPatchStateForSource("src", gen, base, worker, nil, "")
 		assert.NilError(t, err)
 		assert.Assert(t, st != nil)
 		return *st

frontend/gateway.go

@@ -134,6 +134,10 @@ func SourceOptFromUIClient(ctx context.Context, c gwclient.Client, dc *dockerui.
 		return loadSourceFilterConfig(ctx, c, sOpt.GetContext)
 	})
 
+	if gomodProxy, ok := GetBuildArg(c, dalec.BuildArgDalecGomodProxy); ok {
+		sOpt.GomodProxy = gomodProxy
+	}
+
 	return sOpt
 }
 

frontend/gateway_test.go

@@ -0,0 +1,21 @@
+package frontend
+
+import (
+	"context"
+	"testing"
+
+	"github.com/project-dalec/dalec"
+)
+
+func TestSourceOptFromUIClientReadsGomodProxyBuildArg(t *testing.T) {
+	t.Parallel()
+
+	const proxy = "http://proxy.example:5000,direct"
+	client := newStubClient()
+	client.opts["build-arg:"+dalec.BuildArgDalecGomodProxy] = proxy
+
+	sOpt := SourceOptFromUIClient(context.Background(), client, nil, nil)
+	if sOpt.GomodProxy != proxy {
+		t.Fatalf("expected GomodProxy %q, got %q", proxy, sOpt.GomodProxy)
+	}
+}

generator_gomod.go

@@ -19,6 +19,9 @@ const (
 	// GoModCacheKey is the key used to identify the go module cache in the buildkit cache.
 	// It is exported only for testing purposes.
 	GomodCacheKey = "dalec-gomod-proxy-cache"
+	// BuildArgDalecGomodProxy is the frontend build arg used to override GOPROXY
+	// for gomod dependency generation.
+	BuildArgDalecGomodProxy = "DALEC_GOMOD_PROXY"
 )
 
 func (g *GeneratorGomod) processBuildArgs(args map[string]string, allowArg func(key string) bool) error {
@@ -63,7 +66,7 @@ func (s *Spec) HasGomods() bool {
 	return false
 }
 
-func withGomod(g *SourceGenerator, srcSt, worker llb.State, subPath string, credHelper llb.RunOption, opts ...llb.ConstraintsOpt) func(llb.State) llb.State {
+func withGomod(g *SourceGenerator, srcSt, worker llb.State, subPath string, credHelper llb.RunOption, gomodProxy string, opts ...llb.ConstraintsOpt) func(llb.State) llb.State {
 	return func(in llb.State) llb.State {
 		const (
 			workDir                      = "/work/src"
@@ -86,12 +89,12 @@ func withGomod(g *SourceGenerator, srcSt, worker llb.State, subPath string, cred
 		scriptPath := filepath.Join(scriptMountpoint, gomodDownloadWrapperBasename)
 
 		for _, path := range paths {
-			in = worker.Run(
+			runOpts := []llb.RunOption{
 				// First download the go module deps to our persistent cache
 				// Then set the GOPROXY to the cache dir so that we can extract just the deps we need
 				// This allows us to persist the module cache across builds and avoid downloading
 				// the same deps over and over again.
-				ShArgs(`set -e; GOMODCACHE="${TMP_GOMODCACHE}" `+scriptPath+`; GOPROXY="file://${TMP_GOMODCACHE}/cache/download" `+scriptPath),
+				ShArgs(`set -e; GOMODCACHE="${TMP_GOMODCACHE}" ` + scriptPath + `; GOPROXY="file://${TMP_GOMODCACHE}/cache/download" ` + scriptPath),
 				g.withGomodSecretsAndSockets(),
 				llb.AddMount(scriptMountpoint, script),
 				llb.AddEnv("GOPATH", "/go"),
@@ -103,6 +106,13 @@ func withGomod(g *SourceGenerator, srcSt, worker llb.State, subPath string, cred
 				llb.AddMount(proxyPath, llb.Scratch(), llb.AsPersistentCacheDir(GomodCacheKey, llb.CacheMountShared)),
 				WithConstraints(opts...),
 				g.Gomod._sourceMap.GetLocation(in),
+			}
+			if gomodProxy != "" {
+				runOpts = append(runOpts, llb.AddEnv("GOPROXY", gomodProxy))
+			}
+
+			in = worker.Run(
+				runOpts...,
 			).AddMount(gomodCacheDir, in)
 		}
 
@@ -238,7 +248,7 @@ func (s *Spec) GomodDeps(sOpt SourceOpts, worker llb.State, opts ...llb.Constrai
 		deps = deps.With(func(in llb.State) llb.State {
 			for _, gen := range src.Generate {
 				if gen.Gomod != nil {
-					in = in.With(withGomod(gen, patched[key], worker, key, credHelperRunOpt, opts...))
+					in = in.With(withGomod(gen, patched[key], worker, key, credHelperRunOpt, sOpt.GomodProxy, opts...))
 				}
 			}
 			return in

generator_gomod_test.go

@@ -1,10 +1,14 @@
 package dalec
 
 import (
+	"context"
 	"encoding/json"
+	"slices"
+	"strings"
 	"testing"
 
 	"github.com/goccy/go-yaml"
+	"github.com/moby/buildkit/client/llb"
 )
 
 func TestGomodReplaceUnmarshal(t *testing.T) {
@@ -210,3 +214,145 @@ func TestGomodReplaceGoModEditArg(t *testing.T) {
 		})
 	}
 }
+
+func TestGomodDepsUsesGomodProxy(t *testing.T) {
+	t.Parallel()
+
+	const proxy = "http://proxy.example:5000,direct"
+	spec := testGomodProxySpec()
+	st := spec.GomodDeps(testGomodProxySourceOpts(proxy), llb.Scratch())
+	if st == nil {
+		t.Fatal("gomod generator succeeded but returned nil state")
+	}
+
+	env := gomodDownloadExecEnv(context.Background(), t, *st)
+	if !slices.Contains(env, "GOPROXY="+proxy) {
+		t.Fatalf("expected gomod exec env to include GOPROXY=%q, got %v", proxy, env)
+	}
+}
+
+func TestGomodDepsSkipsEmptyGomodProxy(t *testing.T) {
+	t.Parallel()
+
+	spec := testGomodProxySpec()
+	st := spec.GomodDeps(testGomodProxySourceOpts(""), llb.Scratch())
+	if st == nil {
+		t.Fatal("gomod generator succeeded but returned nil state")
+	}
+
+	env := gomodDownloadExecEnv(context.Background(), t, *st)
+	for _, item := range env {
+		if strings.HasPrefix(item, "GOPROXY=") {
+			t.Fatalf("expected empty gomod proxy to omit GOPROXY, got %v", env)
+		}
+	}
+}
+
+func TestGomodProxyBuildArgIsKnown(t *testing.T) {
+	t.Parallel()
+
+	spec := &Spec{}
+	err := spec.SubstituteArgs(map[string]string{
+		BuildArgDalecGomodProxy: "http://proxy.example:5000",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestGomodPatchUsesGomodProxy(t *testing.T) {
+	t.Parallel()
+
+	const proxy = "http://proxy.example:5000,direct"
+	gen := &SourceGenerator{
+		Gomod: &GeneratorGomod{
+			Edits: &GomodEdits{
+				Replace: []GomodReplace{
+					{Original: "example.com/old", Update: "example.com/new v1.2.3"},
+				},
+			},
+		},
+	}
+
+	st, err := (&Spec{}).generateGomodPatchStateForSource("src", gen, llb.Scratch(), llb.Scratch(), nil, proxy)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if st == nil {
+		t.Fatal("gomod patch generation succeeded but returned nil state")
+	}
+
+	env := gomodPatchExecEnv(context.Background(), t, *st)
+	if !slices.Contains(env, "GOPROXY="+proxy) {
+		t.Fatalf("expected gomod patch exec env to include GOPROXY=%q, got %v", proxy, env)
+	}
+}
+
+func testGomodProxySpec() *Spec {
+	return &Spec{
+		Sources: map[string]Source{
+			"src": {
+				Git: &SourceGit{
+					URL:    "https://example.com/repo.git",
+					Commit: "0123456789abcdef",
+				},
+				Generate: []*SourceGenerator{
+					{Gomod: &GeneratorGomod{}},
+				},
+			},
+		},
+	}
+}
+
+func testGomodProxySourceOpts(proxy string) SourceOpts {
+	return SourceOpts{
+		GomodProxy: proxy,
+		GetContext: func(name string, opts ...llb.LocalOption) (*llb.State, error) {
+			st := llb.Local(name, opts...)
+			return &st, nil
+		},
+		GitCredHelperOpt: func() (llb.RunOption, error) {
+			st := llb.Scratch().File(llb.Mkfile("/frontend", 0o755, []byte("#!/usr/bin/env bash\nexit 0\n")))
+			return RunOptFunc(func(ei *llb.ExecInfo) {
+				llb.AddMount("/usr/local/bin/frontend", st, llb.SourcePath("/frontend")).SetRunOption(ei)
+			}), nil
+		},
+	}
+}
+
+func gomodDownloadExecEnv(ctx context.Context, t *testing.T, st llb.State) []string {
+	t.Helper()
+
+	for _, op := range sourceOpsFromState(ctx, t, st) {
+		exec := op.GetExec()
+		if exec == nil {
+			continue
+		}
+
+		env := exec.Meta.Env
+		if slices.Contains(env, "GOPATH=/go") && slices.Contains(env, "TMP_GOMODCACHE=/tmp/dalec/gomod-proxy-cache") {
+			return env
+		}
+	}
+
+	t.Fatal("expected gomod download exec")
+	return nil
+}
+
+func gomodPatchExecEnv(ctx context.Context, t *testing.T, st llb.State) []string {
+	t.Helper()
+
+	for _, op := range sourceOpsFromState(ctx, t, st) {
+		exec := op.GetExec()
+		if exec == nil {
+			continue
+		}
+
+		if slices.Contains(exec.Meta.Args, "/gomod-patch.sh") {
+			return exec.Meta.Env
+		}
+	}
+
+	t.Fatal("expected gomod patch exec")
+	return nil
+}

load.go

@@ -51,6 +51,8 @@ func knownArg(key string) bool {
 		return true
 	case "DALEC_SOURCE_FILTER_CONFIG_CONTEXT_NAME":
 		return true
+	case BuildArgDalecGomodProxy:
+		return true
 	case KeyDalecTarget:
 		return true
 	}

preprocess.go

@@ -68,7 +68,7 @@ func (s *Spec) preprocessGomodEdits(sOpt SourceOpts, worker llb.State, opts ...l
 			}
 
 			// Generate patch state (LLB state, not solved bytes)
-			patchSt, err := s.generateGomodPatchStateForSource(sourceName, gen, baseState, worker, credHelper, opts...)
+			patchSt, err := s.generateGomodPatchStateForSource(sourceName, gen, baseState, worker, credHelper, sOpt.GomodProxy, opts...)
 			if err != nil {
 				return errors.Wrapf(err, "failed to generate gomod patch state for source %s", sourceName)
 			}
@@ -220,7 +220,7 @@ func buildGomodPatchEnv(editArgs string, paths []string, gen *SourceGenerator, s
 // generateGomodPatchStateForSource generates a single merged patch LLB state for all paths
 // in a gomod generator by running go mod edit + tidy and capturing the diff.
 // Returns the LLB state containing the patch file, or nil if no changes are needed.
-func (s *Spec) generateGomodPatchStateForSource(sourceName string, gen *SourceGenerator, baseState llb.State, worker llb.State, credHelper llb.RunOption, opts ...llb.ConstraintsOpt) (*llb.State, error) {
+func (s *Spec) generateGomodPatchStateForSource(sourceName string, gen *SourceGenerator, baseState llb.State, worker llb.State, credHelper llb.RunOption, gomodProxy string, opts ...llb.ConstraintsOpt) (*llb.State, error) {
 	editArgs, err := gomodEditArgs(gen.Gomod)
 	if err != nil {
 		return nil, err
@@ -274,6 +274,10 @@ func (s *Spec) generateGomodPatchStateForSource(sourceName string, gen *SourceGe
 		WithConstraints(opts...),
 	}
 
+	if gomodProxy != "" {
+		runOpts = append(runOpts, llb.AddEnv("GOPROXY", gomodProxy))
+	}
+
 	// Add environment variables from the script
 	for key, value := range SortedMapIter(envVars) {
 		runOpts = append(runOpts, llb.AddEnv(key, value))

source.go

@@ -155,6 +155,10 @@ type SourceOpts struct {
 	TargetPlatform   *ocispecs.Platform
 	GitCredHelperOpt func() (llb.RunOption, error)
 	SourceFilter     func() (SourceFilterConfig, error)
+	// GomodProxy, when set, overrides GOPROXY while gomod generators download
+	// modules from the network. The generator still uses its local file proxy
+	// for the cache extraction pass.
+	GomodProxy string
 }
 
 var errInvalidMountConfig = errors.New("invalid mount config")