fix/batches: validate shell probe paths

Amp-Thread-ID: https://ampcode.com/threads/T-01a05d3c-cfaf-7540-9b4d-43c5e402c15b
Co-authored-by: Amp <amp@ampcode.com>

Author: cbrnrd

internal/batches/executor/run_steps.go

@@ -9,7 +9,9 @@ import (
 	"maps"
 	"os"
 	"os/exec"
+	"path"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"time"
 
@@ -348,21 +350,33 @@ func executeSingleStep(
 		scriptWorkDir = workDir + "/" + opts.Task.Path
 	}
 
+	if err := validateContainerTempPath(containerTemp); err != nil {
+		return bytes.Buffer{}, bytes.Buffer{}, errors.Wrap(err, "validating run script target")
+	}
+	runScriptMount, err := dockerBindMount(runScriptFile, containerTemp)
+	if err != nil {
+		return bytes.Buffer{}, bytes.Buffer{}, errors.Wrap(err, "creating run script mount")
+	}
+
 	args := append([]string{
 		"run",
 		"--rm",
 		"--init",
 		"--cidfile", cidFile,
 		"--workdir", scriptWorkDir,
-		"--mount", fmt.Sprintf("type=bind,source=%s,target=%s,ro", runScriptFile, containerTemp),
+		"--mount", runScriptMount,
 	}, workspaceOpts...)
 
 	if opts.ForceRoot {
 		args = append(args, "--user", "0:0")
 	}
 
 	for target, source := range filesToMount {
-		args = append(args, "--mount", fmt.Sprintf("type=bind,source=%s,target=%s,ro", source.Name(), target))
+		mountArg, err := dockerBindMount(source.Name(), target)
+		if err != nil {
+			return bytes.Buffer{}, bytes.Buffer{}, errors.Wrap(err, "creating files mount")
+		}
+		args = append(args, "--mount", mountArg)
 	}
 
 	// Mount any paths on the local system to the docker container. The paths have already been validated during parsing.
@@ -371,7 +385,11 @@ func executeSingleStep(
 		if err != nil {
 			return bytes.Buffer{}, bytes.Buffer{}, err
 		}
-		args = append(args, "--mount", fmt.Sprintf("type=bind,source=%s,target=%s,ro", workspaceFilePath, mount.Mountpoint))
+		mountArg, err := dockerBindMount(workspaceFilePath, mount.Mountpoint)
+		if err != nil {
+			return bytes.Buffer{}, bytes.Buffer{}, errors.Wrap(err, "creating host mount")
+		}
+		args = append(args, "--mount", mountArg)
 	}
 
 	for k, v := range env {
@@ -479,6 +497,8 @@ func setOutputs(stepOutputs batcheslib.Outputs, global map[string]any, stepCtx *
 	return nil
 }
 
+var containerTempPathPattern = regexp.MustCompile(`^/[A-Za-z0-9_./-]+$`)
+
 func probeImageForShell(ctx context.Context, image string) (shell, tempfile string, err error) {
 	// We need to know two things to be able to run a shell script:
 	//
@@ -508,9 +528,14 @@ func probeImageForShell(ctx context.Context, image string) (shell, tempfile stri
 		if runErr := cmd.Run(); runErr != nil {
 			err = errors.Append(err, errors.Wrapf(runErr, "probing shell %q:\n%s", shell, stderr.String()))
 		} else {
+			tempfile, runErr = parseContainerTempPath(stdout.String())
+			if runErr != nil {
+				err = errors.Append(err, errors.Wrapf(runErr, "probing shell %q", shell))
+				continue
+			}
+
 			// Even if there were previous errors, we can now ignore them.
 			err = nil
-			tempfile = strings.TrimSpace(stdout.String())
 			return
 		}
 	}
@@ -520,6 +545,39 @@ func probeImageForShell(ctx context.Context, image string) (shell, tempfile stri
 	return
 }
 
+func parseContainerTempPath(output string) (string, error) {
+	// mktemp writes one path followed by a newline. Remove that one expected
+	// delimiter, then reject any additional output rather than allowing the
+	// image to inject Docker's comma-delimited mount grammar.
+	tempfile := strings.TrimSuffix(output, "\n")
+	if tempfile == "" {
+		return "", errors.New("mktemp returned an empty path")
+	}
+	if strings.ContainsAny(tempfile, "\r\n") {
+		return "", errors.New("mktemp returned more than one line")
+	}
+	if err := validateContainerTempPath(tempfile); err != nil {
+		return "", err
+	}
+	return tempfile, nil
+}
+
+func validateContainerTempPath(tempfile string) error {
+	if !path.IsAbs(tempfile) || !containerTempPathPattern.MatchString(tempfile) {
+		return errors.Newf("mktemp returned invalid path %q", tempfile)
+	}
+	return nil
+}
+
+func dockerBindMount(source, target string) (string, error) {
+	for name, value := range map[string]string{"source": source, "target": target} {
+		if value == "" || strings.ContainsAny(value, ",\r\n\x00") {
+			return "", errors.Newf("invalid Docker mount %s %q", name, value)
+		}
+	}
+	return fmt.Sprintf("type=bind,source=%s,target=%s,ro", source, target), nil
+}
+
 // createFilesToMount creates temporary files with the contents of Step.Files
 // that are to be mounted into the container that executes the step.
 func createFilesToMount(tempDir string, step batcheslib.Step, stepContext *template.StepContext) (map[string]*os.File, func(), error) {

internal/batches/executor/run_steps_test.go

@@ -1,6 +1,10 @@
 package executor
 
 import (
+	"context"
+	"os"
+	"path/filepath"
+	"runtime"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -9,6 +13,54 @@ import (
 	"github.com/sourcegraph/sourcegraph/lib/batches/template"
 )
 
+func TestParseContainerTempPath(t *testing.T) {
+	for _, valid := range []string{"/tmp/tmp.abc-123_456", "/tmp/tmp.abc-123_456\n"} {
+		t.Run("valid_"+valid, func(t *testing.T) {
+			got, err := parseContainerTempPath(valid)
+			require.NoError(t, err)
+			require.Equal(t, "/tmp/tmp.abc-123_456", got)
+		})
+	}
+
+	for _, invalid := range []string{
+		"",
+		"\n",
+		"relative/path\n",
+		"/tmp/first\n/tmp/second\n",
+		"/tmp/path\r\n",
+		"/tmp/x,source=/var/run/docker.sock,target=/var/run/docker.sock\n",
+		"/tmp/path with spaces\n",
+		`/tmp/path"quoted`,
+		`/tmp/path\backslash`,
+	} {
+		t.Run("invalid_"+invalid, func(t *testing.T) {
+			_, err := parseContainerTempPath(invalid)
+			require.Error(t, err)
+		})
+	}
+}
+
+func TestProbeImageForShellRejectsMountInjection(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("test uses a shell script as a fake docker executable")
+	}
+
+	dir := t.TempDir()
+	docker := filepath.Join(dir, "docker")
+	err := os.WriteFile(docker, []byte("#!/bin/sh\nprintf '%s\\n' '/tmp/x,source=/var/run/docker.sock,target=/var/run/docker.sock'\n"), 0o755)
+	require.NoError(t, err)
+	t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
+
+	_, _, err = probeImageForShell(context.Background(), "malicious-image")
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "mktemp returned invalid path")
+}
+
+func TestDockerBindMountRejectsMountGrammar(t *testing.T) {
+	_, err := dockerBindMount("/tmp/script", "/tmp/x,source=/var/run/docker.sock")
+	require.Error(t, err)
+}
+
 func TestCreateFilesToMount_RejectsCommaInTargetPath(t *testing.T) {
 	step := batcheslib.Step{
 		Files: map[string]string{

internal/batches/executor/testdata/dummydocker/docker

@@ -7,7 +7,7 @@
 # created a tempfile, or it executes the script supplied as the last arg to
 # "run".
 
-dummy_temp_file="DUMMYDOCKER-TEMP-FILE"
+dummy_temp_file="/tmp/DUMMYDOCKER-TEMP-FILE"
 
 workdir_prefix="/work/"