Merge pull request #13 from peg/staging

Security hardening: token exposure, resource limits, WebSocket deadlines

Author: peg

README.md

@@ -414,6 +414,8 @@ rules:
 
 The webhook receives the full tool call context and returns `{"decision": "allow"}` or `{"decision": "deny", "reason": "..."}`. Fail-open by default so a down webhook doesn't break your agent.
 
+**Reference implementation**: See [`rampart-verify`](https://github.com/peg/rampart-verify) — an optional sidecar that uses LLMs (gpt-4o-mini, Claude Haiku, or local Ollama) to classify ambiguous commands. Pattern matching handles 95% of decisions for free; the sidecar reviews the rest at ~$0.0001/call.
+
 ---
 
 ## Integration

cmd/rampart/cli/daemon.go

@@ -70,7 +70,7 @@ Example:
 			}
 
 			// Ensure audit directory exists.
-			if err := os.MkdirAll(auditDir, 0o755); err != nil {
+			if err := os.MkdirAll(auditDir, 0o700); err != nil {
 				return fmt.Errorf("daemon: create audit dir: %w", err)
 			}
 

cmd/rampart/cli/wrap.go

@@ -95,7 +95,7 @@ func newWrapCmd(opts *rootOptions, deps *wrapDeps) *cobra.Command {
 				}
 				auditDir = filepath.Join(home, ".rampart", "audit")
 			}
-			if err := os.MkdirAll(auditDir, 0o755); err != nil {
+			if err := os.MkdirAll(auditDir, 0o700); err != nil {
 				return fmt.Errorf("wrap: create audit dir %s: %w", auditDir, err)
 			}
 
@@ -167,6 +167,7 @@ func newWrapCmd(opts *rootOptions, deps *wrapDeps) *cobra.Command {
 			}
 			defer func() {
 				_ = os.Remove(shimPath)
+				_ = os.Remove(shimPath + ".tok")
 			}()
 
 			child := exec.Command(args[0], args[1:]...)
@@ -185,7 +186,7 @@ func newWrapCmd(opts *rootOptions, deps *wrapDeps) *cobra.Command {
 				childEnv = append(childEnv, e)
 			}
 			// Create PATH-based shell wrappers for agents that ignore $SHELL
-			shimDir, err := createShellWrappers(proxyURL, proxyServer.Token(), mode)
+			shimDir, err := createShellWrappers(proxyURL, shimPath+".tok", mode)
 			if err != nil {
 				return fmt.Errorf("wrap: create shell wrappers: %w", err)
 			}
@@ -355,11 +356,20 @@ func createShellShim(proxyURL, token, mode, realShell string) (string, error) {
 		return "", fmt.Errorf("wrap: create shell shim: %w", err)
 	}
 
+	// Write token to a separate file (0600) so it's not visible in the shim
+	// script or /proc/*/cmdline.
+	tokenFile := tmp.Name() + ".tok"
+	if err := os.WriteFile(tokenFile, []byte(token), 0o600); err != nil {
+		_ = tmp.Close()
+		_ = os.Remove(tmp.Name())
+		return "", fmt.Errorf("wrap: write token file: %w", err)
+	}
+
 	script := fmt.Sprintf(`#!/usr/bin/env bash
 # Rampart shell shim - auto-generated.
 REAL_SHELL=%q
 RAMPART_URL=%q
-RAMPART_TOKEN=%q
+RAMPART_TOKEN=$(cat %q 2>/dev/null)
 RAMPART_MODE=%q
 
 # Parse shell flags — collect flags before -c, extract the command after -c.
@@ -420,7 +430,7 @@ if [ "$FOUND_C" = "true" ]; then
 fi
 
 exec "$REAL_SHELL" $ORIG_ARGS
-`, realShell, proxyURL, token, mode)
+`, realShell, proxyURL, tokenFile, mode)
 
 	if _, err := io.WriteString(tmp, script); err != nil {
 		_ = tmp.Close()
@@ -446,11 +456,16 @@ exec "$REAL_SHELL" $ORIG_ARGS
 // through Rampart policy before executing. If not set, passes straight through
 // to the real shell. This catches agents that hardcode /bin/bash or /bin/zsh
 // instead of reading $SHELL.
-func createShellWrappers(proxyURL, token, mode string) (string, error) {
+func createShellWrappers(proxyURL, tokenFile, mode string) (string, error) {
 	dir, err := os.MkdirTemp("", "rampart-shells-*")
 	if err != nil {
 		return "", fmt.Errorf("create shell wrapper dir: %w", err)
 	}
+	// Restrict directory permissions — contains scripts with token file paths.
+	if err := os.Chmod(dir, 0o700); err != nil {
+		_ = os.RemoveAll(dir)
+		return "", fmt.Errorf("chmod shell wrapper dir: %w", err)
+	}
 
 	shells := []struct {
 		name     string
@@ -496,8 +511,9 @@ if [ "$FOUND_C" = "true" ]; then
 
     ENCODED=$(printf '%%s' "$CMD" | base64 | tr -d '\n\r')
     PAYLOAD=$(printf '{"agent":"wrapped","session":"wrap","params":{"command_b64":"%%s"}}' "$ENCODED")
+    RAMPART_TOKEN=$(cat %q 2>/dev/null)
     DECISION=$(curl -sfS -X POST "%s/v1/preflight/exec" \
-        -H "Authorization: Bearer %s" \
+        -H "Authorization: Bearer ${RAMPART_TOKEN}" \
         -H "Content-Type: application/json" \
         -d "$PAYLOAD" 2>/dev/null)
 
@@ -520,7 +536,7 @@ fi
 
 # Interactive or non -c usage — pass through directly
 exec "$REAL" $ORIG_ARGS
-`, s.name, s.realPath, proxyURL, token, mode)
+`, s.name, s.realPath, tokenFile, proxyURL, mode)
 
 		wrapperPath := filepath.Join(dir, s.name)
 		if err := os.WriteFile(wrapperPath, []byte(script), 0o755); err != nil {

internal/approval/store.go

@@ -253,6 +253,8 @@ func (s *Store) watchExpiry(req *Request) {
 	select {
 	case <-req.done:
 		return // Already resolved.
+	case <-s.stop:
+		return // Store is shutting down.
 	case <-timer.C:
 		s.mu.Lock()
 		if req.Status == StatusPending {

internal/daemon/daemon.go

@@ -142,6 +142,37 @@ func (d *Daemon) connectAndListen(ctx context.Context) error {
 	d.seq.Store(0)
 	d.mu.Unlock()
 
+	// Set up ping/pong to detect dead connections.
+	const pongWait = 90 * time.Second
+	const pingInterval = 30 * time.Second
+	conn.SetReadDeadline(time.Now().Add(pongWait))
+	conn.SetPongHandler(func(string) error {
+		conn.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+
+	// Start ping ticker in background.
+	pingDone := make(chan struct{})
+	go func() {
+		defer close(pingDone)
+		ticker := time.NewTicker(pingInterval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				d.mu.Lock()
+				err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second))
+				d.mu.Unlock()
+				if err != nil {
+					return
+				}
+			}
+		}
+	}()
+	defer func() { <-pingDone }()
+
 	// Perform handshake.
 	if err := d.handshake(ctx); err != nil {
 		return fmt.Errorf("daemon: handshake: %w", err)
@@ -162,6 +193,8 @@ func (d *Daemon) connectAndListen(ctx context.Context) error {
 			return fmt.Errorf("daemon: read: %w", err)
 		}
 
+		// Reset deadline on every successful read.
+		conn.SetReadDeadline(time.Now().Add(pongWait))
 		d.handleMessage(ctx, message)
 	}
 }

internal/engine/policy.go

@@ -135,8 +135,13 @@ type WebhookActionConfig struct {
 }
 
 // EffectiveTimeout returns the configured timeout or 5s default.
+// Capped at 30s to prevent resource exhaustion.
 func (c *WebhookActionConfig) EffectiveTimeout() time.Duration {
+	const maxTimeout = 30 * time.Second
 	if c.Timeout.Duration > 0 {
+		if c.Timeout.Duration > maxTimeout {
+			return maxTimeout
+		}
 		return c.Timeout.Duration
 	}
 	return 5 * time.Second

internal/mcp/proxy.go

@@ -100,8 +100,9 @@ type Proxy struct {
 }
 
 type pendingCall struct {
-	call    engine.ToolCall
-	request map[string]any
+	call      engine.ToolCall
+	request   map[string]any
+	createdAt time.Time
 }
 
 // NewProxy creates a new MCP stdio proxy.
@@ -287,13 +288,41 @@ func (p *Proxy) handleToolsCall(req Request, rawLine []byte) error {
 	if HasID(req.ID) {
 		id := NormalizedID(req.ID)
 		p.pendingMu.Lock()
-		p.pendingCalls[id] = pendingCall{call: call, request: requestData}
+		p.pendingCalls[id] = pendingCall{call: call, request: requestData, createdAt: time.Now()}
+		p.evictStalePendingCalls()
 		p.pendingMu.Unlock()
 	}
 
 	return p.writeToChild(rawLine)
 }
 
+// evictStalePendingCalls removes pending calls older than 5 minutes.
+// Must be called with pendingMu held.
+func (p *Proxy) evictStalePendingCalls() {
+	const maxAge = 5 * time.Minute
+	const maxPending = 1000
+	now := time.Now()
+	for id, pc := range p.pendingCalls {
+		if now.Sub(pc.createdAt) > maxAge {
+			delete(p.pendingCalls, id)
+		}
+	}
+	// Hard cap: if still too many, evict oldest
+	if len(p.pendingCalls) > maxPending {
+		var oldestID string
+		var oldestTime time.Time
+		for id, pc := range p.pendingCalls {
+			if oldestID == "" || pc.createdAt.Before(oldestTime) {
+				oldestID = id
+				oldestTime = pc.createdAt
+			}
+		}
+		if oldestID != "" {
+			delete(p.pendingCalls, oldestID)
+		}
+	}
+}
+
 func (p *Proxy) handleChildLine(line []byte, parentOut io.Writer) error {
 	trimmed := bytes.TrimSpace(line)
 

internal/proxy/server.go

@@ -110,8 +110,8 @@ func New(eng *engine.Engine, sink audit.AuditSink, opts ...Option) *Server {
 		s.token = generateToken(s.logger)
 	}
 
-	if len(s.token) > 8 {
-		s.logger.Info("proxy: auth token", "token", s.token[:8]+"...")
+	if len(s.token) > 4 {
+		s.logger.Info("proxy: auth token", "prefix", s.token[:4]+"…")
 	}
 	return s
 }
@@ -683,6 +683,11 @@ func enrichParams(toolName string, params map[string]any) {
 		if cmd, ok := decodeBase64Command(params); ok {
 			params["command"] = cmd
 		}
+		// Strip leading shell comment lines (e.g. "# description\nactual command")
+		// so that command_matches patterns work against the real command.
+		if cmd, ok := params["command"].(string); ok {
+			params["command"] = stripLeadingComments(cmd)
+		}
 	}
 
 	if toolName == "fetch" || toolName == "http" || toolName == "web_fetch" {
@@ -706,12 +711,42 @@ func enrichParams(toolName string, params map[string]any) {
 	}
 }
 
+// stripLeadingComments removes leading lines that start with # (shell comments)
+// from multi-line command strings. Agent frameworks often prepend descriptive
+// comments (e.g. "# Check disk space\ndf -h") which break command_matches
+// patterns that expect the actual command at the start of the string.
+func stripLeadingComments(cmd string) string {
+	lines := strings.Split(cmd, "\n")
+	start := 0
+	for start < len(lines) {
+		trimmed := strings.TrimSpace(lines[start])
+		if trimmed == "" || strings.HasPrefix(trimmed, "#") {
+			start++
+			continue
+		}
+		break
+	}
+	if start == 0 {
+		return cmd
+	}
+	if start >= len(lines) {
+		return cmd // all comments — return original
+	}
+	return strings.Join(lines[start:], "\n")
+}
+
 func decodeBase64Command(params map[string]any) (string, bool) {
 	encoded, _ := params["command_b64"].(string)
 	if strings.TrimSpace(encoded) == "" {
 		return "", false
 	}
 
+	// Cap encoded input at 1MB to prevent memory exhaustion.
+	const maxBase64Len = 1 << 20
+	if len(encoded) > maxBase64Len {
+		return "", false
+	}
+
 	decoded, err := base64.StdEncoding.DecodeString(encoded)
 	if err != nil {
 		return "", false

internal/proxy/server_test.go

@@ -371,3 +371,31 @@ func TestListenAndServeTimeouts(t *testing.T) {
 	assert.Equal(t, 30*time.Second, httpSrv.WriteTimeout)
 	assert.Equal(t, 120*time.Second, httpSrv.IdleTimeout)
 }
+
+func TestStripLeadingComments(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{"no comments", "ls -la", "ls -la"},
+		{"single comment", "# list files\nls -la", "ls -la"},
+		{"multiple comments", "# step 1\n# step 2\nls -la", "ls -la"},
+		{"comment with blank line", "# desc\n\nls -la", "ls -la"},
+		{"no stripping needed", "git push origin main", "git push origin main"},
+		{"all comments returns original", "# just a comment\n# another", "# just a comment\n# another"},
+		{"inline comment preserved", "ls -la # list files", "ls -la # list files"},
+		{"multiline command", "# build\ndocker build -t app .\ndocker push app", "docker build -t app .\ndocker push app"},
+		{"empty string", "", ""},
+		{"whitespace comment", "  # padded comment\necho hi", "echo hi"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := stripLeadingComments(tt.input)
+			if got != tt.want {
+				t.Errorf("stripLeadingComments(%q) = %q, want %q", tt.input, got, tt.want)
+			}
+		})
+	}
+}