refactor: address PR review comments

- Change buildArgs parameter from []string to map[string]string
- Add exported BuildArgsMap helper for call-site conversion
- Use child.Next.Value instead of fragile child.Original slicing
- Use strings.Cut instead of strings.Contains + SplitN
- Return ProcessWord errors instead of silently swallowing them
- Add UserFromDockerfile_BuildArgs test with registry
- Add MissingBuildArgUsesEmpty test case

Author: kacpersaw

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.BuildArgs)
+		params.User, err = UserFromDockerfile(params.DockerfileContent, BuildArgsMap(params.BuildArgs))
 		if err != nil {
 			return nil, fmt.Errorf("user from dockerfile: %w", err)
 		}
@@ -306,9 +306,20 @@ 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, buildArgs []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)
@@ -319,27 +330,18 @@ func UserFromDockerfile(dockerfileContent string, buildArgs []string) (user stri
 	lexer := shell.NewLex('\\')
 	var argEnvs []string
 	for _, child := range res.AST.Children {
-		if strings.EqualFold(child.Value, "arg") {
-			arg := strings.TrimSpace(child.Original[len("ARG "):])
-			if strings.Contains(arg, "=") {
-				parts := strings.SplitN(arg, "=", 2)
-				key := parts[0]
-				val := parts[1]
-				// Allow buildArgs to override.
-				for _, ba := range buildArgs {
-					if k, v, ok := strings.Cut(ba, "="); ok && k == key {
-						val = v
-						break
-					}
-				}
-				argEnvs = append(argEnvs, key+"="+val)
-			} else {
-				for _, ba := range buildArgs {
-					if k, v, ok := strings.Cut(ba, "="); ok && k == arg {
-						argEnvs = append(argEnvs, k+"="+v)
-						break
-					}
-				}
+		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)
 			}
 		}
 	}
@@ -362,9 +364,10 @@ func UserFromDockerfile(dockerfileContent string, buildArgs []string) (user stri
 		case *instructions.Stage:
 			// Substitute ARG values in the base image name.
 			baseName, _, err := lexer.ProcessWord(i.BaseName, shell.EnvsFromSlice(argEnvs))
-			if err == nil {
-				i.BaseName = baseName
+			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
@@ -423,7 +426,7 @@ func UserFromDockerfile(dockerfileContent string, buildArgs []string) (user stri
 
 // ImageFromDockerfile inspects the contents of a provided Dockerfile
 // and returns the image that will be used to run the container.
-func ImageFromDockerfile(dockerfileContent string, buildArgs []string) (name.Reference, error) {
+func ImageFromDockerfile(dockerfileContent string, buildArgs map[string]string) (name.Reference, error) {
 	lexer := shell.NewLex('\\')
 	var args []string
 	var imageRef string
@@ -433,31 +436,24 @@ func ImageFromDockerfile(dockerfileContent string, buildArgs []string) (name.Ref
 		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.
-				for _, ba := range buildArgs {
-					if k, v, ok := strings.Cut(ba, "="); ok && k == key {
-						val = v
-						break
-					}
+				if override, has := buildArgs[key]; has {
+					val = override
 				}
 				args = append(args, key+"="+val)
 			} else {
 				// ARG without a default — look up in buildArgs.
-				for _, ba := range buildArgs {
-					if k, v, ok := strings.Cut(ba, "="); ok && k == arg {
-						args = append(args, k+"="+v)
-						break
-					}
+				if val, has := buildArgs[arg]; has {
+					args = append(args, arg+"="+val)
 				}
 			}
 			continue

devcontainer/devcontainer_test.go

@@ -227,7 +227,7 @@ func TestImageFromDockerfile_BuildArgs(t *testing.T) {
 	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, []string{"VARIANT=3.11-bookworm"})
+		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())
 	})
@@ -236,10 +236,46 @@ func TestImageFromDockerfile_BuildArgs(t *testing.T) {
 	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, []string{"VARIANT=3.11-bookworm"})
+		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) {

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, buildParams.BuildArgs)
+				imageRef, err := devcontainer.ImageFromDockerfile(buildParams.DockerfileContent, devcontainer.BuildArgsMap(buildParams.BuildArgs))
 				if err != nil {
 					return nil, fmt.Errorf("image from dockerfile: %w", err)
 				}