Merge pull request #2480 from dgageot/board/code-review-find-bugs-and-security-issue-d9ce6d2b

fix: harden artifact store, skills loader, hooks, shell and agent warnings

Author: dgageot

pkg/agent/agent.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log/slog"
 	"math/rand"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -39,8 +40,13 @@ type Agent struct {
 	addPromptFiles          []string
 	tools                   []tools.Tool
 	commands                types.Commands
-	pendingWarnings         []string
 	hooks                   *latest.HooksConfig
+
+	// warningsMu guards pendingWarnings. addToolWarning and DrainWarnings
+	// may be called concurrently from the runtime loop, the MCP server,
+	// the TUI and session manager.
+	warningsMu      sync.Mutex
+	pendingWarnings []string
 }
 
 // New creates a new agent
@@ -286,14 +292,15 @@ func (a *Agent) addToolWarning(msg string) {
 	if msg == "" {
 		return
 	}
+	a.warningsMu.Lock()
 	a.pendingWarnings = append(a.pendingWarnings, msg)
+	a.warningsMu.Unlock()
 }
 
 // DrainWarnings returns pending warnings and clears them.
 func (a *Agent) DrainWarnings() []string {
-	if len(a.pendingWarnings) == 0 {
-		return nil
-	}
+	a.warningsMu.Lock()
+	defer a.warningsMu.Unlock()
 	warnings := a.pendingWarnings
 	a.pendingWarnings = nil
 	return warnings

pkg/agent/agent_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"errors"
 	"log/slog"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -296,3 +298,54 @@ func TestAgentNoDuplicateStartWarnings(t *testing.T) {
 	require.NoError(t, err)
 	assert.Empty(t, a.DrainWarnings(), "turn 3: no duplicate warning on repeated failure")
 }
+
+// TestAgentWarningsConcurrentAccess exercises the warnings queue from
+// multiple goroutines to catch regressions in locking. Run with -race to
+// actually detect a regression.
+func TestAgentWarningsConcurrentAccess(t *testing.T) {
+	t.Parallel()
+
+	a := New("root", "test")
+
+	const writers = 8
+	const drainers = 4
+	const perWriter = 200
+
+	var wg sync.WaitGroup
+	wg.Add(writers + drainers)
+
+	for range writers {
+		go func() {
+			defer wg.Done()
+			for range perWriter {
+				a.addToolWarning("boom")
+			}
+		}()
+	}
+
+	stop := make(chan struct{})
+	for range drainers {
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-stop:
+					// One final drain so we can assert a total count.
+					_ = a.DrainWarnings()
+					return
+				default:
+					_ = a.DrainWarnings()
+				}
+			}
+		}()
+	}
+
+	// Give writers a little time to finish, then signal drainers to stop.
+	time.Sleep(20 * time.Millisecond)
+	close(stop)
+	wg.Wait()
+
+	// A successful run means no data race and no panic; we don't assert a
+	// specific number of warnings drained because drainers run concurrently
+	// with writers.
+}

pkg/content/store.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"time"
 
@@ -18,6 +19,23 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
 )
 
+// ErrInvalidDigest indicates that an identifier shaped like a digest
+// (e.g. "sha256:...") does not match the expected format. Rejecting these
+// prevents path-traversal when the digest is used as a filename component.
+var ErrInvalidDigest = errors.New("invalid artifact digest")
+
+// sha256DigestRe matches a well-formed sha256 digest: 64 lowercase hex chars.
+var sha256DigestRe = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
+
+// validateDigest returns nil if digest is a well-formed sha256 digest,
+// or a wrapped ErrInvalidDigest otherwise.
+func validateDigest(digest string) error {
+	if !sha256DigestRe.MatchString(digest) {
+		return fmt.Errorf("%w: %q", ErrInvalidDigest, digest)
+	}
+	return nil
+}
+
 // ErrStoreCorrupted indicates that the local artifact store is in an
 // inconsistent or partially missing state (e.g. missing tar, refs or metadata).
 // Callers may safely attempt to re-fetch the artifact from the remote source.
@@ -79,6 +97,11 @@ func (s *Store) StoreArtifact(img v1.Image, reference string) (string, error) {
 
 	digestStr := digest.String()
 
+	// Validate the digest before using it in filesystem paths (defense-in-depth).
+	if err := validateDigest(digestStr); err != nil {
+		return "", err
+	}
+
 	tarPath := filepath.Join(s.baseDir, digestStr+".tar")
 
 	if err := crane.Save(img, reference, tarPath); err != nil {
@@ -266,27 +289,32 @@ func (s *Store) DeleteArtifact(identifier string) error {
 
 // resolveIdentifier resolves a user-provided identifier (digest or reference)
 // into a concrete content digest stored in the local artifact store.
+//
+// The returned digest is always strictly validated ("sha256:" + 64 hex chars)
+// so it can be safely used as a filename component without enabling path
+// traversal.
 func (s *Store) resolveIdentifier(identifier string) (string, error) {
-	// If the identifier is already a bare digest, return it directly.
+	// Bare digest, e.g. "sha256:abc123...".
 	if strings.HasPrefix(identifier, "sha256:") {
+		if err := validateDigest(identifier); err != nil {
+			return "", err
+		}
 		return identifier, nil
 	}
 
-	// If the identifier is a digest reference (e.g. "repo@sha256:abc..."),
-	// extract and return the digest portion directly. Digest references
-	// are content-addressable, so the digest alone identifies the artifact.
+	// Digest reference, e.g. "repo@sha256:abc123...".
 	if i := strings.LastIndex(identifier, "@sha256:"); i >= 0 {
-		return identifier[i+1:], nil
+		digest := identifier[i+1:]
+		if err := validateDigest(digest); err != nil {
+			return "", err
+		}
+		return digest, nil
 	}
 
-	// If no tag is provided, default to ":latest".
-	// This mirrors standard OCI reference semantics.
+	// Tagged or tag-less reference. Default to ":latest" per OCI semantics.
 	if !strings.Contains(identifier, ":") {
 		identifier += ":latest"
 	}
-
-	// Resolve the reference to a digest via the refs store.
-	// Any failure here indicates the local store is missing or inconsistent.
 	return s.resolveReference(identifier)
 }
 
@@ -312,8 +340,14 @@ func (s *Store) resolveReference(reference string) (string, error) {
 		return "", fmt.Errorf("reading reference file: %w", err)
 	}
 
-	// The file content is expected to be the digest string.
-	return strings.TrimSpace(string(data)), nil
+	// The file content is expected to be the digest string. Refs files are
+	// generated by us, but we validate defense-in-depth in case the refs
+	// directory is ever tampered with.
+	digest := strings.TrimSpace(string(data))
+	if err := validateDigest(digest); err != nil {
+		return "", fmt.Errorf("ref %q: %w", reference, err)
+	}
+	return digest, nil
 }
 
 // createReferenceLink creates a link from reference to digest

pkg/content/store_test.go

@@ -2,6 +2,8 @@ package content
 
 import (
 	"fmt"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/google/go-containerregistry/pkg/v1/empty"
@@ -140,3 +142,42 @@ func TestStoreResolution_DigestReference(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, digest, meta.Digest)
 }
+
+// TestStoreResolution_RejectsMalformedDigest ensures that identifiers
+// shaped like a digest but carrying path-traversal sequences are rejected
+// before they can be joined into a filesystem path.
+func TestStoreResolution_RejectsMalformedDigest(t *testing.T) {
+	baseDir := t.TempDir()
+	store, err := NewStore(WithBaseDir(baseDir))
+	require.NoError(t, err)
+
+	// Create a sentinel file outside baseDir to confirm we don't touch it.
+	sentinelDir := t.TempDir()
+	sentinel := filepath.Join(sentinelDir, "sentinel.tar")
+	require.NoError(t, os.WriteFile(sentinel, []byte("keep me"), 0o600))
+
+	malformed := []string{
+		"sha256:../../etc/passwd",
+		"sha256:../" + filepath.Base(sentinelDir) + "/sentinel",
+		"sha256:",
+		"sha256:deadbeef", // too short
+		// non-hex char in an otherwise 64-char body
+		"sha256:z0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde",
+		"repo@sha256:../../oops",
+	}
+
+	for _, id := range malformed {
+		_, err := store.GetArtifactImage(id)
+		require.ErrorIsf(t, err, ErrInvalidDigest, "id=%q", id)
+
+		_, err = store.GetArtifactPath(id)
+		require.ErrorIsf(t, err, ErrInvalidDigest, "id=%q", id)
+
+		err = store.DeleteArtifact(id)
+		require.ErrorIsf(t, err, ErrInvalidDigest, "id=%q", id)
+	}
+
+	// Sentinel must be untouched.
+	_, err = os.Stat(sentinel)
+	require.NoError(t, err, "sentinel file should not have been affected")
+}

pkg/hooks/executor.go

@@ -281,6 +281,19 @@ func (e *Executor) executeHook(ctx context.Context, hook Hook, inputJSON []byte)
 	// Run command
 	err := cmd.Run()
 
+	// A fired timeout or parent-context cancellation surfaces as a non-nil
+	// error whose Go type varies (often *exec.ExitError with ExitCode()==-1).
+	// Normalize to a plain execution error so PreToolUse gates can fail
+	// closed rather than look at a meaningless exit code.
+	if ctxErr := timeoutCtx.Err(); ctxErr != nil {
+		reason := "cancelled"
+		if errors.Is(ctxErr, context.DeadlineExceeded) {
+			reason = fmt.Sprintf("timed out after %s", hook.GetTimeout())
+		}
+		return nil, stdout.String(), stderr.String(), -1,
+			fmt.Errorf("hook %q %s: %w", hook.Command, reason, ctxErr)
+	}
+
 	exitCode := 0
 	if err != nil {
 		if exitErr, ok := errors.AsType[*exec.ExitError](err); ok {
@@ -318,7 +331,18 @@ func (e *Executor) aggregateResults(results []hookResult, eventType EventType) (
 
 	for _, r := range results {
 		if r.err != nil {
-			slog.Warn("Hook execution error", "error", r.err)
+			// PreToolUse is a security boundary: if a hook fails to
+			// produce a verdict (timeout, spawn failure, missing binary),
+			// deny the tool call rather than silently letting it through.
+			if eventType == EventPreToolUse {
+				slog.Warn("PreToolUse hook failed to execute; denying tool call", "error", r.err)
+				finalResult.Allowed = false
+				finalResult.ExitCode = -1
+				finalResult.Stderr = r.stderr
+				messages = append(messages, fmt.Sprintf("PreToolUse hook failed to execute: %v", r.err))
+			} else {
+				slog.Warn("Hook execution error", "error", r.err)
+			}
 			continue
 		}
 

pkg/hooks/hooks_test.go

@@ -632,6 +632,72 @@ func TestExecuteHooksWithContextCancellation(t *testing.T) {
 
 	result, err := exec.ExecutePreToolUse(ctx, input)
 	require.NoError(t, err)
-	// Should be allowed because the hook timed out (non-blocking error)
+	// PreToolUse is a security boundary: when the hook fails to run to
+	// completion (here, the parent context was cancelled before the hook
+	// could report a verdict), the tool call must be denied rather than
+	// silently allowed.
+	assert.False(t, result.Allowed)
+	assert.Equal(t, -1, result.ExitCode)
+	assert.Contains(t, result.Message, "PreToolUse hook failed to execute")
+}
+
+// A hook that exits with a non-zero, non-2 code is a non-blocking error:
+// it is reported as such in the result but does not deny the tool call.
+// Pair this with TestExecuteHooksWithContextCancellation, which asserts the
+// opposite for execution failures (timeout, spawn error).
+func TestExecutePreToolUseAllowsNonBlockingExitCode(t *testing.T) {
+	t.Parallel()
+
+	config := &Config{
+		PreToolUse: []MatcherConfig{
+			{
+				Matcher: "*",
+				Hooks: []Hook{
+					{Type: HookTypeCommand, Command: "exit 1", Timeout: 5},
+				},
+			},
+		},
+	}
+
+	exec := NewExecutor(config, t.TempDir(), nil)
+	input := &Input{
+		SessionID: "test-session",
+		ToolName:  "shell",
+		ToolUseID: "test-id",
+	}
+
+	result, err := exec.ExecutePreToolUse(t.Context(), input)
+	require.NoError(t, err)
+	assert.True(t, result.Allowed)
+}
+
+func TestExecutePostToolUseDoesNotFailClosedOnError(t *testing.T) {
+	t.Parallel()
+
+	config := &Config{
+		PostToolUse: []MatcherConfig{
+			{
+				Matcher: "*",
+				Hooks: []Hook{
+					{Type: HookTypeCommand, Command: "sleep 10", Timeout: 30},
+				},
+			},
+		},
+	}
+
+	exec := NewExecutor(config, t.TempDir(), nil)
+	input := &Input{
+		SessionID: "test-session",
+		ToolName:  "shell",
+		ToolUseID: "test-id",
+	}
+
+	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
+	defer cancel()
+
+	result, err := exec.ExecutePostToolUse(ctx, input)
+	require.NoError(t, err)
+	// Post-tool-use is observational only: a failed hook must not block
+	// the already-completed tool call.
 	assert.True(t, result.Allowed)
 }