fix(devcontainer): parse user in multi-stage builds

Fixes #231

devcontainer/devcontainer.go

@@ -15,6 +15,8 @@ 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/instructions"
+	"github.com/moby/buildkit/frontend/dockerfile/parser"
 	"github.com/moby/buildkit/frontend/dockerfile/shell"
 	"github.com/tailscale/hujson"
 )
@@ -202,16 +204,9 @@ 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 = UserFromDockerfile(params.DockerfileContent)
-	}
-	if params.User == "" {
-		imageRef, err := ImageFromDockerfile(params.DockerfileContent)
+		params.User, err = UserFromDockerfile(params.DockerfileContent)
 		if err != nil {
-			return nil, fmt.Errorf("parse image from dockerfile: %w", err)
-		}
-		params.User, err = UserFromImage(imageRef)
-		if err != nil {
-			return nil, fmt.Errorf("get user from image: %w", err)
+			return nil, fmt.Errorf("user from dockerfile: %w", err)
 		}
 	}
 	remoteUser := s.RemoteUser
@@ -313,17 +308,82 @@ func (s *Spec) compileFeatures(fs billy.Filesystem, devcontainerDir, scratchDir
 
 // UserFromDockerfile inspects the contents of a provided Dockerfile
 // and returns the user that will be used to run the container.
-func UserFromDockerfile(dockerfileContent string) 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, "USER ") {
+func UserFromDockerfile(dockerfileContent string) (user string, err error) {
+	res, err := parser.Parse(strings.NewReader(dockerfileContent))
+	if err != nil {
+		return "", fmt.Errorf("parse dockerfile: %w", err)
+	}
+
+	// Parse stages and user commands to determine the relevant user
+	// from the final stage.
+	var (
+		stages       []*instructions.Stage
+		stageNames   = make(map[string]*instructions.Stage)
+		stageUser    = make(map[*instructions.Stage]*instructions.UserCommand)
+		currentStage *instructions.Stage
+	)
+	for _, child := range res.AST.Children {
+		inst, err := instructions.ParseInstruction(child)
+		if err != nil {
+			return "", fmt.Errorf("parse instruction: %w", err)
+		}
+
+		switch i := inst.(type) {
+		case *instructions.Stage:
+			stages = append(stages, i)
+			if i.Name != "" {
+				stageNames[i.Name] = i
+			}
+			currentStage = i
+		case *instructions.UserCommand:
+			if currentStage == nil {
+				continue
+			}
+			stageUser[currentStage] = i
+		}
+	}
+
+	// Iterate over stages in bottom-up order to find the user,
+	// skipping any stages not referenced by the final stage.
+	lookupStage := stages[len(stages)-1]
+	for i := len(stages) - 1; i >= 0; i-- {
+		stage := stages[i]
+		if stage != lookupStage {
 			continue
 		}
-		return strings.TrimSpace(strings.TrimPrefix(line, "USER "))
+
+		if user, ok := stageUser[stage]; ok {
+			return user.User, nil
+		}
+
+		// If we reach the scratch stage, we can't determine the user.
+		if stage.BaseName == "scratch" {
+			return "", nil
+		}
+
+		// Check if this FROM references another stage.
+		if stage.BaseName != "" {
+			var ok bool
+			lookupStage, ok = stageNames[stage.BaseName]
+			if ok {
+				continue
+			}
+		}
+
+		// If we can't find a user command, try to find the user from
+		// the image.
+		ref, err := name.ParseReference(strings.TrimSpace(stage.BaseName))
+		if err != nil {
+			return "", fmt.Errorf("parse image ref %q: %w", stage.BaseName, err)
+		}
+		user, err := UserFromImage(ref)
+		if err != nil {
+			return "", fmt.Errorf("user from image %s: %w", ref.Name(), err)
+		}
+		return user, nil
 	}
-	return ""
+
+	return "", nil
 }
 
 // ImageFromDockerfile inspects the contents of a provided Dockerfile

devcontainer/devcontainer_test.go

@@ -190,12 +190,6 @@ func TestCompileDevContainer(t *testing.T) {
 	})
 }
 
-func TestUserFromDockerfile(t *testing.T) {
-	t.Parallel()
-	user := devcontainer.UserFromDockerfile("FROM ubuntu\nUSER kyle")
-	require.Equal(t, "kyle", user)
-}
-
 func TestImageFromDockerfile(t *testing.T) {
 	t.Parallel()
 	for _, tc := range []struct {
@@ -224,27 +218,156 @@ func TestImageFromDockerfile(t *testing.T) {
 	}
 }
 
-func TestUserFromImage(t *testing.T) {
+func TestUserFrom(t *testing.T) {
 	t.Parallel()
-	registry := registrytest.New(t)
-	image, err := partial.UncompressedToImage(emptyImage{configFile: &v1.ConfigFile{
-		Config: v1.Config{
-			User: "example",
-		},
-	}})
-	require.NoError(t, err)
 
-	parsed, err := url.Parse(registry)
-	require.NoError(t, err)
-	parsed.Path = "coder/test:latest"
-	ref, err := name.ParseReference(strings.TrimPrefix(parsed.String(), "http://"))
-	require.NoError(t, err)
-	err = remote.Write(ref, image)
-	require.NoError(t, err)
+	t.Run("Image", func(t *testing.T) {
+		t.Parallel()
+		registry := registrytest.New(t)
+		image, err := partial.UncompressedToImage(emptyImage{configFile: &v1.ConfigFile{
+			Config: v1.Config{
+				User: "example",
+			},
+		}})
+		require.NoError(t, err)
 
-	user, err := devcontainer.UserFromImage(ref)
-	require.NoError(t, err)
-	require.Equal(t, "example", user)
+		parsed, err := url.Parse(registry)
+		require.NoError(t, err)
+		parsed.Path = "coder/test:latest"
+		ref, err := name.ParseReference(strings.TrimPrefix(parsed.String(), "http://"))
+		require.NoError(t, err)
+		err = remote.Write(ref, image)
+		require.NoError(t, err)
+
+		user, err := devcontainer.UserFromImage(ref)
+		require.NoError(t, err)
+		require.Equal(t, "example", user)
+	})
+
+	t.Run("Dockerfile table based test", func(t *testing.T) {
+		t.Parallel()
+		tests := []struct {
+			name    string
+			content string
+			user    string
+		}{
+			{
+				name:    "Empty",
+				content: "FROM scratch",
+				user:    "",
+			},
+			{
+				name:    "User",
+				content: "FROM scrach\nUSER kyle",
+				user:    "kyle",
+			},
+			{
+				name:    "Env with default",
+				content: "FROM scratch\nENV MYUSER=maf\nUSER ${MYUSER}",
+				user:    "${MYUSER}", // This should be "maf" but the current implementation doesn't support this.
+			},
+			{
+				name:    "Env var with default",
+				content: "FROM scratch\nUSER ${MYUSER:-maf}",
+				user:    "${MYUSER:-maf}", // This should be "maf" but the current implementation doesn't support this.
+			},
+			{
+				name:    "Arg",
+				content: "FROM scratch\nARG MYUSER\nUSER ${MYUSER}",
+				user:    "${MYUSER}", // This should be "" or populated but the current implementation doesn't support this.
+			},
+			{
+				name:    "Arg with default",
+				content: "FROM scratch\nARG MYUSER=maf\nUSER ${MYUSER}",
+				user:    "${MYUSER}", // This should be "maf" but the current implementation doesn't support this.
+			},
+		}
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				t.Parallel()
+				user, err := devcontainer.UserFromDockerfile(tt.content)
+				require.NoError(t, err)
+				require.Equal(t, tt.user, user)
+			})
+		}
+	})
+
+	t.Run("Multi-stage", func(t *testing.T) {
+		t.Parallel()
+
+		registry := registrytest.New(t)
+		for tag, user := range map[string]string{
+			"one": "maf",
+			"two": "fam",
+		} {
+			image, err := partial.UncompressedToImage(emptyImage{configFile: &v1.ConfigFile{
+				Config: v1.Config{
+					User: user,
+				},
+			}})
+			require.NoError(t, err)
+			parsed, err := url.Parse(registry)
+			require.NoError(t, err)
+			parsed.Path = "coder/test:" + tag
+			ref, err := name.ParseReference(strings.TrimPrefix(parsed.String(), "http://"))
+			fmt.Println(ref)
+			require.NoError(t, err)
+			err = remote.Write(ref, image)
+			require.NoError(t, err)
+		}
+
+		tests := []struct {
+			name    string
+			images  map[string]string
+			content string
+			user    string
+		}{
+			{
+				name:    "Single",
+				content: "FROM coder/test:one",
+				user:    "maf",
+			},
+			{
+				name:    "Multi",
+				content: "FROM ubuntu AS u\nFROM coder/test:two",
+				user:    "fam",
+			},
+			{
+				name:    "Multi-2",
+				content: "FROM coder/test:two AS two\nUSER maffam\nFROM coder/test:one AS one",
+				user:    "maf",
+			},
+			{
+				name:    "Multi-3",
+				content: "FROM coder/test:two AS two\nFROM coder/test:one AS one\nUSER fammaf",
+				user:    "fammaf",
+			},
+			{
+				name: "Multi-4",
+				content: `FROM ubuntu AS a
+USER root
+RUN useradd --create-home pickme
+USER pickme
+FROM a AS other
+USER root
+RUN useradd --create-home notme
+USER notme
+FROM a`,
+				user: "pickme",
+			},
+		}
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				t.Parallel()
+
+				content := strings.ReplaceAll(tt.content, "coder/test", strings.TrimPrefix(registry, "http://")+"/coder/test")
+
+				user, err := devcontainer.UserFromDockerfile(content)
+				require.NoError(t, err)
+				require.Equal(t, tt.user, user)
+			})
+		}
+	})
 }
 
 type emptyImage struct {

integration/integration_test.go

@@ -102,6 +102,63 @@ func TestInitScriptInitCommand(t *testing.T) {
 	require.NoError(t, ctx.Err(), "init script did not execute for legacy env vars")
 }
 
+func TestDanglingBuildStage(t *testing.T) {
+	t.Parallel()
+
+	// Ensures that a Git repository with a devcontainer.json is cloned and built.
+	srv := gittest.CreateGitServer(t, gittest.Options{
+		Files: map[string]string{
+			"devcontainer.json": `{
+				"name": "Test",
+				"build": {
+					"dockerfile": "Dockerfile"
+				},
+			}`,
+			"Dockerfile": fmt.Sprintf(`FROM %s as a
+RUN date > /root/date.txt`, testImageUbuntu),
+		},
+	})
+	ctr, err := runEnvbuilder(t, runOpts{env: []string{
+		envbuilderEnv("GIT_URL", srv.URL),
+	}})
+	require.NoError(t, err)
+
+	output := execContainer(t, ctr, "cat /date.txt")
+	require.NotEmpty(t, strings.TrimSpace(output))
+}
+
+func TestUserFromMultistage(t *testing.T) {
+	t.Parallel()
+
+	// Ensures that a Git repository with a devcontainer.json is cloned and built.
+	srv := gittest.CreateGitServer(t, gittest.Options{
+		Files: map[string]string{
+			"devcontainer.json": `{
+				"name": "Test",
+				"build": {
+					"dockerfile": "Dockerfile"
+				},
+			}`,
+			"Dockerfile": fmt.Sprintf(`FROM %s AS a
+USER root
+RUN useradd --create-home pickme
+USER pickme
+FROM a AS other
+USER root
+RUN useradd --create-home notme
+USER notme
+FROM a AS b`, testImageUbuntu),
+		},
+	})
+	ctr, err := runEnvbuilder(t, runOpts{env: []string{
+		envbuilderEnv("GIT_URL", srv.URL),
+	}})
+	require.NoError(t, err)
+
+	output := execContainer(t, ctr, "ps aux | awk '/^pickme / {print $1}' | sort -u")
+	require.Equal(t, "pickme", strings.TrimSpace(output))
+}
+
 func TestUidGid(t *testing.T) {
 	t.Parallel()
 	t.Run("MultiStage", func(t *testing.T) {