Skip to content

Commit e134f01

Browse files
authored
L8: MCP client robustness (timeouts, names, terminal escapes) (#80)
- internal/mcpclient/client.go: default timeout for caller-deadline-free calls; validate server names; reject tool names containing '__'; DefaultTimeout var for test overrides. - internal/mcpclient/client_test.go: tests for invalid server names, double underscore tool names, default timeout. - internal/mcpclient/testdata/main.go: FAKE_DELAY env for timeout tests. - cmd/odek/mcp_approval.go: sanitize tool descriptions in approval prompt. - cmd/odek/mcp_approval_test.go: sanitize tests. - docs/SECURITY.md + AGENTS.md: document L8 mitigations.
1 parent 73b0a29 commit e134f01

7 files changed

Lines changed: 175 additions & 3 deletions

File tree

AGENTS.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,7 @@ Layered prompt-injection / approval-fatigue defenses. Full reference: [docs/SECU
182182
- **odek self-invocation gate** (`internal/danger/classifier.go`) — any shell stage whose program basename is `odek` is classified as `system_write`. This prevents a prompt-injected agent from reaching the human-gated trust mutations (`odek memory promote`, `odek memory extended confirm`, `odek skill promote --force`) through the shell tool and flipping its own taint gates.
183183
- **MCP inputSchema hardening** (`cmd/odek/mcp_approval.go`) — every string in an MCP tool's `inputSchema` is recursively guard-scanned for injection patterns; schemas larger than 256 KiB are rejected; the interactive approval prompt shows a SHA-256 hash and byte size of the schema so operators can detect changes.
184184
- **MCP tool batch classification** (`internal/loop/loop.go`) — MCP tools (`<server>__<tool>`) are classified as `unknown` by `classifyToolCall`, so the batch approval gate shows them and untrusted sub-agents force them to `deny`.
185+
- **MCP client robustness** (`internal/mcpclient/client.go`, `cmd/odek/mcp_approval.go`) — MCP calls use a default timeout when the caller supplies no deadline; server names and tool names are validated (tool names are rejected if they contain `__` to prevent collisions with odek's `<server>__<tool>` namespace); the interactive approval prompt sanitises tool descriptions to strip ANSI escape sequences and terminal control characters.
185186
- **Sub-agent trust defaults + delegate_tasks gate** (`cmd/odek/subagent.go`, `internal/loop/loop.go`) — a missing `trust_level` in `delegate_tasks` defaults to `untrusted`; `delegate_tasks` itself is classified as `system_write` so it requires explicit approval before spawning child processes.
186187
- **Memory add/replace pipe-to-shell filter** (`internal/memory/memory.go`) — `AddFact` and `ReplaceFact` now run `FactLooksUnsafe` in addition to the general guard scan, blocking agent-driven planting of download-and-execute facts.
187188
- **Skill learn-loop provenance propagation** (`internal/skills/learnloop.go`) — conversation-extracted suggestions and LLM-enhanced suggestions both retain the session's `SkillProvenance`, so tainted sessions cannot produce clean-looking auto-saved skills.

cmd/odek/mcp_approval.go

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"io"
1212
"os"
1313
"path/filepath"
14+
"regexp"
1415
"sort"
1516
"strings"
1617

@@ -243,7 +244,7 @@ func approveMCPToolsWithTTY(projectDir, serverName string, cfg mcpclient.ServerC
243244

244245
fmt.Fprintf(stdout, "\nMCP server %q wants to register tool %q\n", serverName, def.Name)
245246
if def.Description != "" {
246-
fmt.Fprintf(stdout, " description: %s\n", truncateDescription(def.Description, 200))
247+
fmt.Fprintf(stdout, " description: %s\n", sanitizeTerminal(truncateDescription(def.Description, 200)))
247248
}
248249
fmt.Fprintf(stdout, " schema: sha256:%s (%d bytes)\n", schemaHash, schemaSize)
249250
fmt.Fprintf(stdout, "Approve? [y/N] ")
@@ -344,6 +345,28 @@ func truncateDescription(desc string, max int) string {
344345
return desc[:max-3] + "..."
345346
}
346347

348+
// sanitizeTerminal removes ANSI escape sequences and replaces other terminal
349+
// control characters with a replacement character so a malicious MCP server
350+
// cannot disguise an approval prompt with cursor movement or colour codes.
351+
func sanitizeTerminal(s string) string {
352+
// Strip ANSI escape sequences: ESC [ ... m (and similar).
353+
ansi := regexp.MustCompile("\x1b\\[[0-9;]*[a-zA-Z]")
354+
s = ansi.ReplaceAllString(s, "")
355+
// Replace remaining control characters (except tab/newline) with �.
356+
var b strings.Builder
357+
for _, r := range s {
358+
switch {
359+
case r == '\t' || r == '\n' || r == '\r':
360+
b.WriteRune(r)
361+
case r < 0x20 || r == 0x7f:
362+
b.WriteRune('�')
363+
default:
364+
b.WriteRune(r)
365+
}
366+
}
367+
return b.String()
368+
}
369+
347370
// maxMCPSchemaBytes caps the serialized JSON schema size for a single MCP tool.
348371
// Schemas are part of the model's tool catalogue, so an oversized schema can be
349372
// used for prompt stuffing. Real-world MCP schemas are typically small; 256 KiB

cmd/odek/mcp_approval_test.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,3 +301,18 @@ func TestApproveMCPTools_OversizeSchemaSkipped(t *testing.T) {
301301
t.Errorf("approved %d tools, want 0 (oversized schema should be skipped)", len(got))
302302
}
303303
}
304+
305+
func TestSanitizeTerminal_StripsANSIAndControlChars(t *testing.T) {
306+
// ANSI red colour + cursor-up sequence, then a bell and normal text.
307+
input := "\x1b[31m\x1b[2A\x07normal"
308+
got := sanitizeTerminal(input)
309+
if strings.Contains(got, "\x1b") {
310+
t.Errorf("ANSI escapes should be stripped, got: %q", got)
311+
}
312+
if strings.Contains(got, "\x07") {
313+
t.Errorf("control characters should be removed, got: %q", got)
314+
}
315+
if !strings.Contains(got, "normal") {
316+
t.Errorf("normal text should be preserved, got: %q", got)
317+
}
318+
}

docs/SECURITY.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -634,6 +634,16 @@ Now:
634634
The per-turn audit log and divergence heuristic only inspected `tool` messages for the nonce'd `<untrusted_content_*>` wrapper. @-references, `--ctx` files, and Web-UI attachments were also wrapped before entering the user message, but:
635635

636636
- `cmd/odek/refs.go` called `wrapUntrusted` with `context.Background()`, so `recordIngest` found no active recorder.
637+
638+
### 39k. MCP client robustness (timeouts, name validation, terminal output)
639+
640+
Three MCP client hardening fixes close availability and spoofing issues:
641+
642+
1. **Request timeout**`internal/mcpclient/client.go` declared `DefaultTimeout` but never applied it. A hung MCP server would block `Discover` or `CallTool` indefinitely. The timeout is now applied automatically when the caller does not supply a context deadline.
643+
644+
2. **Server-name validation** — MCP server names are used as the prefix in registered tool names (`<server>__<tool>`). Names are now validated to be non-empty, ≤ 64 characters, ASCII letters/digits/underscore/hyphen only, and must not contain `__`. This prevents invalid API identifiers from killing every LLM turn and blocks the collision where server `a` + tool `b__c` produces the same effective name as server `a__b` + tool `c`. Tool names are also rejected if they contain `__`.
645+
646+
3. **Terminal-sanitized approval prompt** — the interactive MCP tool approval prompt printed the server-supplied description verbatim. A malicious server could hide cursor movement or colour codes in the description to disguise what was being approved. Descriptions are now passed through `sanitizeTerminal`, which strips ANSI escape sequences and replaces other control characters before printing.
637647
- `cmd/odek/serve.go` resolved @-refs and attachments before attaching the per-session ingest recorder.
638648
- `recordTurnAudit` only scanned `tool` messages, so `ingested_untrusted` stayed false for these vectors.
639649
- The enriched prompt (including injected resource literals) was passed as the "user message" to the divergence check, making attacker resources count as user-mentioned and disabling the heuristic.

internal/mcpclient/client.go

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,7 @@ import (
4444
// ── Protocol Constants ──────────────────────────────────────────────────
4545

4646
const (
47-
ProtocolVersion = "2025-03-26"
48-
DefaultTimeout = 30 * time.Second
47+
ProtocolVersion = "2025-03-26"
4948
// maxMCPResponseLine caps the size of a single JSON-RPC response line
5049
// from an MCP server. A malicious or broken server that emits a huge
5150
// line without a newline would otherwise be buffered entirely in memory
@@ -54,6 +53,10 @@ const (
5453
maxMCPResponseLine = 10 << 20 // 10 MiB
5554
)
5655

56+
// DefaultTimeout bounds each MCP request when the caller does not supply a
57+
// context deadline. It is a var so tests can temporarily lower it.
58+
var DefaultTimeout = 30 * time.Second
59+
5760
// ── JSON-RPC Types ─────────────────────────────────────────────────────
5861

5962
// request is a JSON-RPC 2.0 request sent to the MCP server.
@@ -119,6 +122,9 @@ func validateToolName(name string) error {
119122
if len(name) > 64 {
120123
return fmt.Errorf("tool name %q exceeds 64 characters", name)
121124
}
125+
if strings.Contains(name, "__") {
126+
return fmt.Errorf("tool name %q cannot contain %q", name, "__")
127+
}
122128
for _, r := range name {
123129
switch {
124130
case r >= 'a' && r <= 'z':
@@ -192,9 +198,38 @@ type Client struct {
192198
pending map[int]chan callResponse // routes responses to waiting callers
193199
}
194200

201+
// validateName checks that an MCP server or tool name is safe to use as part
202+
// of an odek tool identifier. Names must be non-empty, ≤64 chars, contain only
203+
// ASCII letters/digits/underscore/hyphen, and must not contain the double-
204+
// underscore separator used to qualify tool names.
205+
func validateName(kind, name string) error {
206+
if name == "" {
207+
return fmt.Errorf("%s name cannot be empty", kind)
208+
}
209+
if len(name) > 64 {
210+
return fmt.Errorf("%s name %q too long (%d > 64)", kind, name, len(name))
211+
}
212+
if strings.Contains(name, "__") {
213+
return fmt.Errorf("%s name %q cannot contain %q", kind, name, "__")
214+
}
215+
for i, r := range name {
216+
isLetter := (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z')
217+
isDigit := r >= '0' && r <= '9'
218+
if isLetter || isDigit || r == '_' || r == '-' {
219+
continue
220+
}
221+
return fmt.Errorf("%s name %q contains invalid character %q at position %d", kind, name, r, i)
222+
}
223+
return nil
224+
}
225+
195226
// New spawns an MCP server process and returns a client connected to it.
196227
// The server process is started immediately and cleaned up on Close().
197228
func New(name string, cfg ServerConfig) (*Client, error) {
229+
if err := validateName("server", name); err != nil {
230+
return nil, err
231+
}
232+
198233
cmd := exec.Command(cfg.Command, cfg.Args...)
199234

200235
// Apply env overrides. Always build a sanitized environment so MCP children
@@ -430,6 +465,14 @@ func (c *Client) call(ctx context.Context, method string, params json.RawMessage
430465
if ctx == nil {
431466
ctx = context.Background()
432467
}
468+
// Bound the request with the package default timeout unless the caller
469+
// already supplied a deadline. A hung MCP server must not deadlock the
470+
// agent loop or startup discovery.
471+
if _, ok := ctx.Deadline(); !ok {
472+
var cancel context.CancelFunc
473+
ctx, cancel = context.WithTimeout(ctx, DefaultTimeout)
474+
defer cancel()
475+
}
433476

434477
// Assign unique ID and register a response channel.
435478
respCh := make(chan callResponse, 1)

internal/mcpclient/client_test.go

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -504,6 +504,78 @@ func TestClient_CloseWithoutDiscover(t *testing.T) {
504504
}
505505
}
506506

507+
// ── L-8 regression tests ────────────────────────────────────────────────
508+
509+
func TestNew_RejectsInvalidServerName(t *testing.T) {
510+
cases := []struct {
511+
name string
512+
}{
513+
{""},
514+
{"has space"},
515+
{"has__double"},
516+
{"tool?name"},
517+
{"a" + strings.Repeat("x", 64)},
518+
}
519+
for _, tc := range cases {
520+
t.Run(tc.name, func(t *testing.T) {
521+
_, err := New(tc.name, ServerConfig{Command: "echo"})
522+
if err == nil {
523+
t.Fatal("expected error for invalid server name")
524+
}
525+
})
526+
}
527+
}
528+
529+
func TestDiscover_RejectsToolWithDoubleUnderscore(t *testing.T) {
530+
client, err := New("valid", ServerConfig{
531+
Command: fakeServerPath(t),
532+
Env: map[string]string{
533+
"FAKE_TOOLS": `[{"name":"a__b","description":"bad"}]`,
534+
},
535+
})
536+
if err != nil {
537+
t.Fatalf("New: %v", err)
538+
}
539+
defer client.Close()
540+
541+
_, err = client.Discover(context.Background())
542+
if err == nil {
543+
t.Fatal("expected Discover to reject tool name containing __")
544+
}
545+
}
546+
547+
func TestCall_DefaultTimeout(t *testing.T) {
548+
// A client whose read loop never answers must time out via DefaultTimeout.
549+
client, err := New("timeout", ServerConfig{
550+
Command: fakeServerPath(t),
551+
Env: map[string]string{
552+
"FAKE_DELAY": "10s",
553+
"FAKE_TOOLS": `[{"name":"x","description":"y"}]`,
554+
},
555+
})
556+
if err != nil {
557+
t.Fatalf("New: %v", err)
558+
}
559+
defer client.Close()
560+
561+
// Restore DefaultTimeout after the test.
562+
orig := DefaultTimeout
563+
DefaultTimeout = 100 * time.Millisecond
564+
defer func() { DefaultTimeout = orig }()
565+
566+
start := time.Now()
567+
_, err = client.Discover(context.Background())
568+
if err == nil {
569+
t.Fatal("expected timeout error")
570+
}
571+
if !contains(err.Error(), "context deadline exceeded") {
572+
t.Errorf("expected context deadline exceeded, got: %v", err)
573+
}
574+
if time.Since(start) > 2*time.Second {
575+
t.Errorf("timeout took too long: %v", time.Since(start))
576+
}
577+
}
578+
507579
// ── Helpers ────────────────────────────────────────────────────────────
508580

509581
func contains(s, substr string) bool {

internal/mcpclient/testdata/main.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77
"os"
88
"strings"
9+
"time"
910
)
1011

1112
type request struct {
@@ -31,6 +32,7 @@ type rpcError struct {
3132
var (
3233
toolsJSON = os.Getenv("FAKE_TOOLS")
3334
errorOnCall = os.Getenv("FAKE_ERROR_ON_CALL")
35+
delayStr = os.Getenv("FAKE_DELAY")
3436
)
3537

3638
func main() {
@@ -46,6 +48,12 @@ func main() {
4648
continue
4749
}
4850

51+
if delayStr != "" {
52+
if d, err := time.ParseDuration(delayStr); err == nil {
53+
time.Sleep(d)
54+
}
55+
}
56+
4957
switch req.Method {
5058
case "initialize":
5159
initResult, _ := json.Marshal(map[string]any{

0 commit comments

Comments
 (0)