Harden exec.Command inputs for scanner Docker invocations and upgrade re-exec

pkg/cli/docker_args_validation.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"regexp"
 	"strings"
 	"unicode"
 
@@ -14,6 +15,16 @@ import (
 
 var dockerArgsValidationLog = logger.New("cli:docker_args_validation")
 
+var (
+	dockerImageNamePattern = regexp.MustCompile(`^(?:[a-zA-Z0-9.-]+(?::[0-9]+)?/)?[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*(?:/[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*)*$`)
+	dockerImageTagPattern  = regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}$`)
+)
+
+var dockerImageDigestAlgorithms = map[string]int{
+	"sha256": 64,
+	"sha512": 128,
+}
+
 func containsControlCharacters(value string) bool {
 	return strings.IndexFunc(value, func(r rune) bool {
 		return unicode.IsControl(r) || unicode.In(r, unicode.Cf) || r == '\u2028' || r == '\u2029'
@@ -63,9 +74,60 @@ func validateDockerImageRef(imageRef string) (string, error) {
 	if strings.HasPrefix(imageRef, "-") {
 		return "", fmt.Errorf("grant image reference cannot start with '-'. Example: ghcr.io/example/image:tag. Got: %q", imageRef)
 	}
+
+	imageRefWithoutDigest := imageRef
+	if strings.Count(imageRef, "@") > 1 {
+		return "", fmt.Errorf("grant image reference has multiple digest separators. Example: ghcr.io/example/image@sha256:<digest>. Got: %q", imageRef)
+	}
+	nameWithOptionalTag, digest, hasDigest := strings.Cut(imageRef, "@")
+	if hasDigest {
+		if digest == "" || !isAllowedDockerImageDigest(digest) {
+			return "", fmt.Errorf("grant image reference has an invalid digest format. Example: ghcr.io/example/image@sha256:<digest>. Got: %q", imageRef)
+		}
+		imageRefWithoutDigest = nameWithOptionalTag
+	}
+	if imageRefWithoutDigest == "" {
+		return "", fmt.Errorf("grant image reference is missing an image name. Example: ghcr.io/example/image:tag. Got: %q", imageRef)
+	}
+
+	imageName := imageRefWithoutDigest
+	if colon := strings.LastIndex(imageRefWithoutDigest, ":"); colon > strings.LastIndex(imageRefWithoutDigest, "/") {
+		tag := imageRefWithoutDigest[colon+1:]
+		if !dockerImageTagPattern.MatchString(tag) {
+			return "", fmt.Errorf("grant image reference has an invalid tag format. Example: ghcr.io/example/image:tag. Got: %q", imageRef)
+		}
+		imageName = imageRefWithoutDigest[:colon]
+	}
+
+	if imageName == "" || strings.HasSuffix(imageName, "/") || !dockerImageNamePattern.MatchString(imageName) {
+		return "", fmt.Errorf("grant image reference must match an allow-listed image pattern. Example: ghcr.io/example/image:tag. Got: %q", imageRef)
+	}
 	return imageRef, nil
 }
 
+func isAllowedDockerImageDigest(digest string) bool {
+	algorithm, hexDigest, ok := strings.Cut(digest, ":")
+	if !ok {
+		return false
+	}
+
+	expectedLength, ok := dockerImageDigestAlgorithms[algorithm]
+	if !ok || len(hexDigest) != expectedLength {
+		return false
+	}
+
+	for _, r := range hexDigest {
+		switch {
+		case r >= '0' && r <= '9':
+		case r >= 'a' && r <= 'f':
+		case r >= 'A' && r <= 'F':
+		default:
+			return false
+		}
+	}
+	return true
+}
+
 func isWindowsDrivePath(hostPath string) bool {
 	if len(hostPath) < 3 {
 		return false

pkg/cli/docker_args_validation_test.go

@@ -88,6 +88,12 @@ func TestValidateDockerImageRefRejectsUnsafeCharacters(t *testing.T) {
 		{name: "trailing newline", imageRef: "alpine:latest\n", wantErr: "invalid whitespace/control characters"},
 		{name: "unicode line separator", imageRef: "alpine\u2028latest", wantErr: "invalid whitespace/control characters"},
 		{name: "unicode bidi override", imageRef: "alpine\u202elatest", wantErr: "invalid whitespace/control characters"},
+		{name: "multiple digests", imageRef: "ghcr.io/org/image@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa@sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", wantErr: "multiple digest separators"},
+		{name: "invalid digest", imageRef: "ghcr.io/org/image@sha256:nothex", wantErr: "invalid digest format"},
+		{name: "invalid digest algorithm", imageRef: "ghcr.io/org/image@sha1:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", wantErr: "invalid digest format"},
+		{name: "invalid sha256 digest length", imageRef: "ghcr.io/org/image@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", wantErr: "invalid digest format"},
+		{name: "invalid tag", imageRef: "ghcr.io/org/image:-tag", wantErr: "invalid tag format"},
+		{name: "invalid image name characters", imageRef: "ghcr.io/org/im;age:latest", wantErr: "allow-listed image pattern"},
 	}
 
 	for _, tt := range testCases {
@@ -98,3 +104,23 @@ func TestValidateDockerImageRefRejectsUnsafeCharacters(t *testing.T) {
 		})
 	}
 }
+
+func TestValidateDockerImageRefAcceptsCommonReferences(t *testing.T) {
+	testCases := []string{
+		"alpine:latest",
+		"ghcr.io/github/gh-aw:1.2.3",
+		"localhost:5000/org/image_name:tag-1",
+		"registry.example.com/team/my__image:latest",
+		"team/my--image:latest",
+		"ghcr.io/org/image@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+		"ghcr.io/org/image@sha512:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+	}
+
+	for _, imageRef := range testCases {
+		t.Run(imageRef, func(t *testing.T) {
+			validated, err := validateDockerImageRef(imageRef)
+			require.NoError(t, err)
+			require.Equal(t, imageRef, validated)
+		})
+	}
+}

pkg/cli/upgrade_command.go

@@ -7,7 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"strings"
+	"slices"
 
 	"github.com/github/gh-aw/pkg/console"
 	"github.com/github/gh-aw/pkg/constants"
@@ -458,6 +458,13 @@ func updateCopilotArtifacts(ctx context.Context, verbose bool) error {
 // path because os.Executable() returns a "(deleted)"-suffixed path after the binary
 // has been renamed out of the way during the upgrade.
 func relaunchWithSameArgs(extraFlag string, exeOverride string) error {
+	allowedExtraFlags := map[string]struct{}{
+		"--skip-extension-upgrade": {},
+	}
+	if _, ok := allowedExtraFlags[extraFlag]; !ok {
+		return fmt.Errorf("invalid relaunch flag %q: expected --skip-extension-upgrade", extraFlag)
+	}
+
 	var exe string
 	if exeOverride != "" {
 		exe = exeOverride
@@ -479,10 +486,8 @@ func relaunchWithSameArgs(extraFlag string, exeOverride string) error {
 	// Explicitly copy os.Args[1:] so appending the extra flag does not modify
 	// the original slice backing array.
 	newArgs := append(append([]string(nil), os.Args[1:]...), extraFlag)
-	for _, arg := range newArgs {
-		if strings.ContainsRune(arg, '\x00') {
-			return errors.New("invalid relaunch arguments: argument contains NUL byte. Example: compile .github/workflows/example.md")
-		}
+	if slices.ContainsFunc(newArgs, containsControlCharacters) {
+		return errors.New("invalid relaunch arguments: argument contains invalid control characters. Example: compile .github/workflows/example.md")
 	}
 	upgradeLog.Printf("Re-launching with new binary: %s %v", exe, newArgs)
 

pkg/cli/upgrade_command_test.go

@@ -141,7 +141,7 @@ func TestUpgradeCommandRepoDispatchWithPR(t *testing.T) {
 }
 
 func TestRelaunchWithSameArgsRejectsRelativeExecutableOverride(t *testing.T) {
-	err := relaunchWithSameArgs("--post-upgrade", "relative/gh-aw")
+	err := relaunchWithSameArgs("--skip-extension-upgrade", "relative/gh-aw")
 	require.Error(t, err)
 	require.ErrorContains(t, err, "invalid executable path")
 }
@@ -151,7 +151,22 @@ func TestRelaunchWithSameArgsRejectsNullByteArgument(t *testing.T) {
 	t.Cleanup(func() { os.Args = origArgs })
 	os.Args = []string{"gh-aw", "compile", "bad\x00arg"}
 
-	err := relaunchWithSameArgs("--post-upgrade", "/bin/echo")
+	err := relaunchWithSameArgs("--skip-extension-upgrade", "/bin/echo")
 	require.Error(t, err)
-	require.ErrorContains(t, err, "argument contains NUL byte")
+	require.ErrorContains(t, err, "argument contains invalid control characters")
+}
+
+func TestRelaunchWithSameArgsAllowsEmptyForwardedArgument(t *testing.T) {
+	origArgs := os.Args
+	t.Cleanup(func() { os.Args = origArgs })
+	os.Args = []string{"gh-aw", "compile", ""}
+
+	err := relaunchWithSameArgs("--skip-extension-upgrade", "/bin/echo")
+	require.NoError(t, err)
+}
+
+func TestRelaunchWithSameArgsRejectsUnknownExtraFlag(t *testing.T) {
+	err := relaunchWithSameArgs("--unknown-flag", "/bin/echo")
+	require.Error(t, err)
+	require.ErrorContains(t, err, "invalid relaunch flag")
 }