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,12 @@ 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(gomodGeneratorOpts{
+			sourceName:  "src",
+			gen:         gen,
+			sourceState: base,
+			worker:      worker,
+		})
 		assert.NilError(t, err)
 		assert.Assert(t, st != nil)
 		return *st

frontend/gateway.go

@@ -134,6 +134,13 @@ func SourceOptFromUIClient(ctx context.Context, c gwclient.Client, dc *dockerui.
 		return loadSourceFilterConfig(ctx, c, sOpt.GetContext)
 	})
 
+	if gomodProxy, ok := GetBuildArg(c, dalec.BuildArgDalecGomodProxy); ok {
+		if sOpt.ExtraEnvs == nil {
+			sOpt.ExtraEnvs = map[string]string{}
+		}
+		sOpt.ExtraEnvs["GOPROXY"] = gomodProxy
+	}
+
 	return sOpt
 }
 

frontend/gateway_test.go

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

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,16 +66,35 @@ 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 {
+type gomodGeneratorOpts struct {
+	sourceName  string
+	gen         *SourceGenerator
+	sourceState llb.State
+	worker      llb.State
+	credHelper  llb.RunOption
+	extraEnvs   map[string]string
+	constraints []llb.ConstraintsOpt
+}
+
+func (opts gomodGeneratorOpts) gomodProxy() string {
+	if opts.extraEnvs == nil {
+		return ""
+	}
+	return opts.extraEnvs["GOPROXY"]
+}
+
+func withGomod(gomodOpts gomodGeneratorOpts) func(llb.State) llb.State {
 	return func(in llb.State) llb.State {
 		const (
 			workDir                      = "/work/src"
 			scriptMountpoint             = "/tmp/dalec/internal/gomod"
 			gomodDownloadWrapperBasename = "go_mod_download.sh"
 		)
 
-		joinedWorkDir := filepath.Join(workDir, subPath, g.Subpath)
-		srcMount := llb.AddMount(workDir, srcSt)
+		g := gomodOpts.gen
+		opts := gomodOpts.constraints
+		joinedWorkDir := filepath.Join(workDir, gomodOpts.sourceName, g.Subpath)
+		srcMount := llb.AddMount(workDir, gomodOpts.sourceState)
 
 		paths := g.Gomod.Paths
 		if g.Gomod.Paths == nil {
@@ -86,23 +108,30 @@ 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"),
-				credHelper,
+				gomodOpts.credHelper,
 				llb.AddEnv("TMP_GOMODCACHE", proxyPath),
 				llb.AddEnv("GIT_SSH_COMMAND", "ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"),
 				llb.Dir(filepath.Join(joinedWorkDir, path)),
 				srcMount,
 				llb.AddMount(proxyPath, llb.Scratch(), llb.AsPersistentCacheDir(GomodCacheKey, llb.CacheMountShared)),
 				WithConstraints(opts...),
 				g.Gomod._sourceMap.GetLocation(in),
+			}
+			if gomodProxy := gomodOpts.gomodProxy(); gomodProxy != "" {
+				runOpts = append(runOpts, llb.AddEnv("GOPROXY", gomodProxy))
+			}
+
+			in = gomodOpts.worker.Run(
+				runOpts...,
 			).AddMount(gomodCacheDir, in)
 		}
 
@@ -238,7 +267,15 @@ 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(gomodGeneratorOpts{
+						sourceName:  key,
+						gen:         gen,
+						sourceState: patched[key],
+						worker:      worker,
+						credHelper:  credHelperRunOpt,
+						extraEnvs:   sOpt.ExtraEnvs,
+						constraints: 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) {
@@ -296,3 +300,158 @@ func TestGomodEditArgs_DropRequire(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(gomodGeneratorOpts{
+		sourceName:  "src",
+		gen:         gen,
+		sourceState: llb.Scratch(),
+		worker:      llb.Scratch(),
+		extraEnvs: map[string]string{
+			"GOPROXY": 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 {
+	sOpt := SourceOpts{
+		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
+		},
+	}
+	if proxy != "" {
+		sOpt.ExtraEnvs = map[string]string{
+			"GOPROXY": proxy,
+		}
+	}
+	return sOpt
+}
+
+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

@@ -53,6 +53,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
 	}