fix: substitute devcontainer.json build args into Dockerfile ARG variables

devcontainer/devcontainer.go

@@ -204,7 +204,7 @@ func (s *Spec) Compile(fs billy.Filesystem, devcontainerDir, scratchDir string,
 		// We should make a best-effort attempt to find the user.
 		// Features must be executed as root, so we need to swap back
 		// to the running user afterwards.
-		params.User, err = UserFromDockerfile(params.DockerfileContent)
+		params.User, err = UserFromDockerfile(params.DockerfileContent, BuildArgsMap(params.BuildArgs))
 		if err != nil {
 			return nil, fmt.Errorf("user from dockerfile: %w", err)
 		}
@@ -306,14 +306,46 @@ func (s *Spec) compileFeatures(fs billy.Filesystem, devcontainerDir, scratchDir
 	return strings.Join(lines, "\n"), featureContexts, err
 }
 
+// BuildArgsMap converts a slice of "KEY=VALUE" strings to a map.
+func BuildArgsMap(buildArgs []string) map[string]string {
+	m := make(map[string]string, len(buildArgs))
+	for _, arg := range buildArgs {
+		if key, val, ok := strings.Cut(arg, "="); ok {
+			m[key] = val
+		}
+	}
+	return m
+}
+
 // UserFromDockerfile inspects the contents of a provided Dockerfile
 // and returns the user that will be used to run the container.
-func UserFromDockerfile(dockerfileContent string) (user string, err error) {
+func UserFromDockerfile(dockerfileContent string, buildArgs map[string]string) (user string, err error) {
 	res, err := parser.Parse(strings.NewReader(dockerfileContent))
 	if err != nil {
 		return "", fmt.Errorf("parse dockerfile: %w", err)
 	}
 
+	// Collect ARG values (defaults + overrides from buildArgs) for
+	// substitution into FROM image refs.
+	lexer := shell.NewLex('\\')
+	var argEnvs []string
+	for _, child := range res.AST.Children {
+		if !strings.EqualFold(child.Value, "arg") || child.Next == nil {
+			continue
+		}
+		if key, val, ok := strings.Cut(child.Next.Value, "="); ok {
+			if override, has := buildArgs[key]; has {
+				val = override
+			}
+			argEnvs = append(argEnvs, key+"="+val)
+		} else {
+			arg := child.Next.Value
+			if val, has := buildArgs[arg]; has {
+				argEnvs = append(argEnvs, arg+"="+val)
+			}
+		}
+	}
+
 	// Parse stages and user commands to determine the relevant user
 	// from the final stage.
 	var (
@@ -330,6 +362,12 @@ func UserFromDockerfile(dockerfileContent string) (user string, err error) {
 
 		switch i := inst.(type) {
 		case *instructions.Stage:
+			// Substitute ARG values in the base image name.
+			baseName, _, err := lexer.ProcessWord(i.BaseName, shell.EnvsFromSlice(argEnvs))
+			if err != nil {
+				return "", fmt.Errorf("processing ARG substitution in FROM %q: %w", i.BaseName, err)
+			}
+			i.BaseName = baseName
 			stages = append(stages, i)
 			if i.Name != "" {
 				stageNames[i.Name] = i
@@ -388,7 +426,7 @@ func UserFromDockerfile(dockerfileContent string) (user string, err error) {
 
 // ImageFromDockerfile inspects the contents of a provided Dockerfile
 // and returns the image that will be used to run the container.
-func ImageFromDockerfile(dockerfileContent string) (name.Reference, error) {
+func ImageFromDockerfile(dockerfileContent string, buildArgs map[string]string) (name.Reference, error) {
 	lexer := shell.NewLex('\\')
 	var args []string
 	var imageRef string
@@ -398,17 +436,25 @@ func ImageFromDockerfile(dockerfileContent string) (name.Reference, error) {
 		line := lines[i]
 		if arg, ok := strings.CutPrefix(line, "ARG "); ok {
 			arg = strings.TrimSpace(arg)
-			if strings.Contains(arg, "=") {
-				parts := strings.SplitN(arg, "=", 2)
-				key, _, err := lexer.ProcessWord(parts[0], shell.EnvsFromSlice(args))
+			if key, val, ok := strings.Cut(arg, "="); ok {
+				key, _, err := lexer.ProcessWord(key, shell.EnvsFromSlice(args))
 				if err != nil {
 					return nil, fmt.Errorf("processing %q: %w", line, err)
 				}
-				val, _, err := lexer.ProcessWord(parts[1], shell.EnvsFromSlice(args))
+				val, _, err := lexer.ProcessWord(val, shell.EnvsFromSlice(args))
 				if err != nil {
 					return nil, fmt.Errorf("processing %q: %w", line, err)
 				}
+				// Allow buildArgs to override Dockerfile ARG defaults.
+				if override, has := buildArgs[key]; has {
+					val = override
+				}
 				args = append(args, key+"="+val)
+			} else {
+				// ARG without a default — look up in buildArgs.
+				if val, has := buildArgs[arg]; has {
+					args = append(args, arg+"="+val)
+				}
 			}
 			continue
 		}

devcontainer/devcontainer_test.go

@@ -213,13 +213,71 @@ func TestImageFromDockerfile(t *testing.T) {
 		tc := tc
 		t.Run(tc.image, func(t *testing.T) {
 			t.Parallel()
-			ref, err := devcontainer.ImageFromDockerfile(tc.content)
+			ref, err := devcontainer.ImageFromDockerfile(tc.content, nil)
 			require.NoError(t, err)
 			require.Equal(t, tc.image, ref.Name())
 		})
 	}
 }
 
+func TestImageFromDockerfile_BuildArgs(t *testing.T) {
+	t.Parallel()
+
+	// Test that build args override ARG defaults.
+	t.Run("OverridesDefault", func(t *testing.T) {
+		t.Parallel()
+		content := "ARG VARIANT=3.10\nFROM mcr.microsoft.com/devcontainers/python:0-${VARIANT}"
+		ref, err := devcontainer.ImageFromDockerfile(content, map[string]string{"VARIANT": "3.11-bookworm"})
+		require.NoError(t, err)
+		require.Equal(t, "mcr.microsoft.com/devcontainers/python:0-3.11-bookworm", ref.Name())
+	})
+
+	// Test that build args supply values for ARGs without defaults.
+	t.Run("SuppliesArgWithoutDefault", func(t *testing.T) {
+		t.Parallel()
+		content := "ARG VARIANT\nFROM mcr.microsoft.com/devcontainers/python:1-${VARIANT}"
+		ref, err := devcontainer.ImageFromDockerfile(content, map[string]string{"VARIANT": "3.11-bookworm"})
+		require.NoError(t, err)
+		require.Equal(t, "mcr.microsoft.com/devcontainers/python:1-3.11-bookworm", ref.Name())
+	})
+
+	// Test that a missing build arg for an ARG without default results
+	// in the variable being substituted as empty string.
+	t.Run("MissingBuildArgUsesEmpty", func(t *testing.T) {
+		t.Parallel()
+		content := "ARG VARIANT\nFROM mcr.microsoft.com/devcontainers/python:1-${VARIANT}"
+		ref, err := devcontainer.ImageFromDockerfile(content, nil)
+		require.NoError(t, err)
+		require.Equal(t, "mcr.microsoft.com/devcontainers/python:1-", ref.Name())
+	})
+}
+
+func TestUserFromDockerfile_BuildArgs(t *testing.T) {
+	t.Parallel()
+
+	t.Run("SubstitutesARGInFROM", func(t *testing.T) {
+		t.Parallel()
+		registry := registrytest.New(t)
+		image, err := partial.UncompressedToImage(emptyImage{configFile: &v1.ConfigFile{
+			Config: v1.Config{
+				User: "testuser",
+			},
+		}})
+		require.NoError(t, err)
+		ref := strings.TrimPrefix(registry, "http://") + "/coder/test:latest"
+		parsed, err := name.ParseReference(ref)
+		require.NoError(t, err)
+		err = remote.Write(parsed, image)
+		require.NoError(t, err)
+
+		// Dockerfile uses ARG without default for the image ref.
+		content := fmt.Sprintf("ARG TAG\nFROM %s/coder/test:${TAG}", strings.TrimPrefix(registry, "http://"))
+		user, err := devcontainer.UserFromDockerfile(content, map[string]string{"TAG": "latest"})
+		require.NoError(t, err)
+		require.Equal(t, "testuser", user)
+	})
+}
+
 func TestUserFrom(t *testing.T) {
 	t.Parallel()
 
@@ -287,7 +345,7 @@ func TestUserFrom(t *testing.T) {
 		for _, tt := range tests {
 			t.Run(tt.name, func(t *testing.T) {
 				t.Parallel()
-				user, err := devcontainer.UserFromDockerfile(tt.content)
+				user, err := devcontainer.UserFromDockerfile(tt.content, nil)
 				require.NoError(t, err)
 				require.Equal(t, tt.user, user)
 			})
@@ -364,7 +422,7 @@ FROM a`,
 
 				content := strings.ReplaceAll(tt.content, "coder/test", strings.TrimPrefix(registry, "http://")+"/coder/test")
 
-				user, err := devcontainer.UserFromDockerfile(content)
+				user, err := devcontainer.UserFromDockerfile(content, nil)
 				require.NoError(t, err)
 				require.Equal(t, tt.user, user)
 			})

envbuilder.go

@@ -493,7 +493,7 @@ func run(ctx context.Context, opts options.Options, execArgs *execArgsInfo) erro
 			defer cleanupBuildContext()
 			if runtimeData.Built && opts.SkipRebuild {
 				endStage := startStage("🏗️ Skipping build because of cache...")
-				imageRef, err := devcontainer.ImageFromDockerfile(buildParams.DockerfileContent)
+				imageRef, err := devcontainer.ImageFromDockerfile(buildParams.DockerfileContent, devcontainer.BuildArgsMap(buildParams.BuildArgs))
 				if err != nil {
 					return nil, fmt.Errorf("image from dockerfile: %w", err)
 				}