Fix Dockerfile ARG parsing

devcontainer/devcontainer.go

@@ -15,6 +15,7 @@ import (
 	"github.com/go-git/go-billy/v5"
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/moby/buildkit/frontend/dockerfile/shell"
 	"github.com/tailscale/hujson"
 )
 
@@ -317,30 +318,43 @@ func UserFromDockerfile(dockerfileContent string) string {
 // 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) {
-	args := map[string]string{}
+	lexer := shell.NewLex('\\')
+	var args []string
 	var imageRef string
 	lines := strings.Split(dockerfileContent, "\n")
 	// Iterate over lines in reverse
 	for i := len(lines) - 1; i >= 0; i-- {
 		line := lines[i]
-		if strings.HasPrefix(line, "ARG ") {
-			arg := strings.TrimSpace(strings.TrimPrefix(line, "ARG "))
+		if arg, ok := strings.CutPrefix(line, "ARG "); ok {
+			arg = strings.TrimSpace(arg)
 			if strings.Contains(arg, "=") {
 				parts := strings.SplitN(arg, "=", 2)
-				args[parts[0]] = parts[1]
+				key, err := lexer.ProcessWord(parts[0], args)
+				if err != nil {
+					return nil, fmt.Errorf("processing %q: %w", line, err)
+				}
+				val, err := lexer.ProcessWord(parts[1], args)
+				if err != nil {
+					return nil, fmt.Errorf("processing %q: %w", line, err)
+				}
+				args = append(args, key+"="+val)
 			}
 			continue
 		}
-		if imageRef == "" && strings.HasPrefix(line, "FROM ") {
-			imageRef = strings.TrimPrefix(line, "FROM ")
+		if imageRef == "" {
+			if fromArgs, ok := strings.CutPrefix(line, "FROM "); ok {
+				imageRef = fromArgs
+			}
 		}
 	}
 	if imageRef == "" {
 		return nil, fmt.Errorf("no FROM directive found")
 	}
-	image, err := name.ParseReference(os.Expand(imageRef, func(s string) string {
-		return args[s]
-	}))
+	imageRef, err := lexer.ProcessWord(imageRef, args)
+	if err != nil {
+		return nil, fmt.Errorf("processing %q: %w", imageRef, err)
+	}
+	image, err := name.ParseReference(strings.TrimSpace(imageRef))
 	if err != nil {
 		return nil, fmt.Errorf("parse image ref %q: %w", imageRef, err)
 	}

devcontainer/devcontainer_test.go

@@ -167,8 +167,14 @@ func TestImageFromDockerfile(t *testing.T) {
 		content: "FROM ubuntu",
 		image:   "index.docker.io/library/ubuntu:latest",
 	}, {
-		content: "ARG VARIANT=ionic\nFROM ubuntu:$VARIANT",
-		image:   "index.docker.io/library/ubuntu:ionic",
+		content: "ARG VARIANT=bionic\nFROM ubuntu:$VARIANT",
+		image:   "index.docker.io/library/ubuntu:bionic",
+	}, {
+		content: "ARG VARIANT=\"3.10\"\nFROM mcr.microsoft.com/devcontainers/python:0-${VARIANT}",
+		image:   "mcr.microsoft.com/devcontainers/python:0-3.10",
+	}, {
+		content: "ARG VARIANT=\"3.10\"\nFROM mcr.microsoft.com/devcontainers/python:0-$VARIANT ",
+		image:   "mcr.microsoft.com/devcontainers/python:0-3.10",
 	}} {
 		tc := tc
 		t.Run(tc.image, func(t *testing.T) {

go.mod

@@ -30,6 +30,7 @@ require (
 	github.com/google/go-containerregistry v0.15.2
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
 	github.com/mattn/go-isatty v0.0.20
+	github.com/moby/buildkit v0.11.6
 	github.com/otiai10/copy v1.14.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.7.0
@@ -184,7 +185,6 @@ require (
 	github.com/mitchellh/go-wordwrap v1.0.1 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.2 // indirect
-	github.com/moby/buildkit v0.11.6 // indirect
 	github.com/moby/locker v1.0.1 // indirect
 	github.com/moby/patternmatcher v0.5.0 // indirect
 	github.com/moby/swarmkit/v2 v2.0.0-20230315203717-e28e8ba9bc83 // indirect