diff --git a/README.md b/README.md index d0b5f5c..5d30eaf 100644 --- a/README.md +++ b/README.md @@ -124,8 +124,9 @@ Using these rules, everyone can build their own copy of Bantam from scratch in l 2. Append the assistant's response message object to `messages`. If non-streaming and response has reasoning tokens (`reasoning_content` or `reasoning`), output them wrapped in `--- reasoning start ---` / `--- reasoning end ---` markers. 3. If there are pending `tool_calls` in the assistant response: - For each tool call, output a trace log (`[tool call: name(args)]`). - - Validate JSON arguments. If invalid, format a tool-error response so the LLM can self-correct. + - Sanitize tool arguments to filter out non-printable and space-like Unicode characters (protecting against indirect prompt injection), and validate JSON arguments. If invalid, format a tool-error response so the LLM can self-correct. - Execute tool action (`shell_exec` or `run_subagent`). + - Sanitize the tool result output to strip any non-printable and space-like Unicode characters (leaving only ASCII space, tab, newline, and printable Unicode characters). - Output a trace log of the result (`[tool result: name]`). - Append tool result message (`role: "tool"`, `tool_call_id`, `content`: result string) to `messages`. - Loop back to step 1. @@ -173,6 +174,7 @@ MicroBantam (`mb`) is a compressed Perl 5 reference implementation of the same a ### Features - Full agentic loop: LLM calls, `shell_exec` / `run_subagent` tool execution with JSON argument validation (invalid args are fed back so the model can self-correct), and the 5-level subagent recursion depth limit +- Indirect prompt injection defense: sanitizes tool parameters and tool outputs by filtering non-printable and space-like Unicode characters, preserving standard space, tab, newline, and printable Unicode characters - A `...requesting...` in-flight indicator: in-place on a TTY (`\r` overwrite, erased on completion), a plain line when output is piped - Session management: `/save`, `/list`, `/load ` (exact id only, no prefix matching), `/cfg [val]`, and auto-save to `~/.bantam/sessions/autosave.json` after every turn and on exit; session ids get `-1`, `-2`, ... suffixes on same-second collisions - Interactive mode (`/quit`, `/clear`, `/save`, `/list`, `/load `, `/cfg`, `/help`) and file input mode diff --git a/main.go b/main.go index 2a36a11..958e4fd 100644 --- a/main.go +++ b/main.go @@ -22,6 +22,7 @@ import ( "strconv" "strings" "time" + "unicode" "unicode/utf8" ) @@ -712,11 +713,27 @@ func cleanMessagesForLLM(msgs []Message) []Message { return out } +func filterText(s string) string { + var b strings.Builder + b.Grow(len(s)) + for _, r := range s { + if r == ' ' || r == '\t' || r == '\n' { + b.WriteRune(r) + } else if unicode.Is(unicode.Z, r) || unicode.IsControl(r) || unicode.Is(unicode.C, r) { + continue + } else if unicode.IsPrint(r) { + b.WriteRune(r) + } + } + return b.String() +} + func sanitizeMessages(msgs []Message) { for i := range msgs { if msgs[i].Role == "assistant" && len(msgs[i].ToolCalls) > 0 { for j := range msgs[i].ToolCalls { tc := &msgs[i].ToolCalls[j] + tc.Function.Arguments = filterText(tc.Function.Arguments) astr := tc.Function.Arguments var a map[string]any if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil { @@ -724,6 +741,8 @@ func sanitizeMessages(msgs []Message) { tc.Function.Arguments = string(fixed) } } + } else if msgs[i].Role == "tool" && msgs[i].Content != nil { + msgs[i].Content = strp(filterText(*msgs[i].Content)) } } } @@ -956,12 +975,13 @@ func parseStream(ctx context.Context, r io.Reader) (Message, Usage, error) { } func shell(ctx context.Context, cmd string, timeout int) string { + cmd = filterText(cmd) cmdCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second) defer cancel() c := exec.CommandContext(cmdCtx, "sh", "-c", cmd) c.WaitDelay = 100 * time.Millisecond out, err := c.CombinedOutput() - res := strings.TrimSpace(string(out)) + res := strings.TrimSpace(filterText(string(out))) if ctx.Err() != nil { return "[interrupted]\n\nexit: -1" } @@ -1018,6 +1038,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([] if u.Cached() > 0 { turnUsage.CachedTokens = u.Cached() } for j := range m.ToolCalls { tc := &m.ToolCalls[j] + tc.Function.Arguments = filterText(tc.Function.Arguments) astr := tc.Function.Arguments var a map[string]any if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil { @@ -1035,7 +1056,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([] if len(m.ToolCalls) == 0 { done = true; break } for _, tc := range m.ToolCalls { if err := ctx.Err(); err != nil { return msgs, turnUsage, err } - fn, astr := tc.Function.Name, tc.Function.Arguments + fn, astr := tc.Function.Name, filterText(tc.Function.Arguments) fmt.Println(c(fmt.Sprintf("[tool call: %s(%s)]", fn, astr), 33)) res, sty := "", 2 var a map[string]any @@ -1045,10 +1066,12 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([] switch fn { case "shell_exec": cmd, _ := a["command"].(string) + cmd = filterText(cmd) res = shell(ctx, cmd, cfg.ShellTimeout) if err := ctx.Err(); err != nil { return msgs, turnUsage, err } case "run_subagent": pr, _ := a["prompt"].(string) + pr = filterText(pr) if depth >= MAX_DEPTH { res, sty = fmt.Sprintf("[subagent depth limit (%d) reached, child not spawned]", MAX_DEPTH), 31 } else { @@ -1064,13 +1087,14 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([] } else { turnUsage.CompletionTokens += subu.CompletionTokens turnUsage.TotalTokens += subu.TotalTokens - res, sty = last(subr), 2 + res, sty = filterText(last(subr)), 2 } } default: res, sty = "Unknown tool: " + fn, 31 } } + res = filterText(res) fmt.Println(c("[tool result: "+fn+"]", 32) + "\n" + c(res, sty) + "\n") msgs = append(msgs, Message{Role: "tool", ToolCallID: tc.ID, Content: strp(res)}) } @@ -1422,7 +1446,7 @@ func main() { } u := strings.TrimSpace(string(data)) if strings.HasPrefix(u, "!") { - cmd := strings.TrimSpace(strings.TrimPrefix(u, "!")) + cmd := filterText(strings.TrimSpace(strings.TrimPrefix(u, "!"))) if cmd != "" { astr, _ := json.Marshal(map[string]string{"command": cmd}) fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33)) @@ -1551,7 +1575,7 @@ func main() { } continue case strings.HasPrefix(u, "!"): - cmd := strings.TrimSpace(strings.TrimPrefix(u, "!")) + cmd := filterText(strings.TrimSpace(strings.TrimPrefix(u, "!"))) if cmd != "" { astr, _ := json.Marshal(map[string]string{"command": cmd}) fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33)) diff --git a/main_test.go b/main_test.go index c309e32..5c6fd16 100644 --- a/main_test.go +++ b/main_test.go @@ -1851,6 +1851,95 @@ func TestFetchContextWindowFallbackConfig(t *testing.T) { } } +func TestFilterText(t *testing.T) { + cases := []struct { + name string + input string + expected string + }{ + { + name: "ASCII printable, spaces, tabs, newlines", + input: "Hello World!\t123\nLine 2 ~`@#$%", + expected: "Hello World!\t123\nLine 2 ~`@#$%", + }, + { + name: "CRLF normalization", + input: "line1\r\nline2\r\n", + expected: "line1\nline2\n", + }, + { + name: "Unicode printable letters, numbers, punctuation", + input: "こんにちは世界! Привет мир! 123 αβγ €$¥", + expected: "こんにちは世界! Привет мир! 123 αβγ €$¥", + }, + { + name: "Control characters stripped", + input: "null\x00bell\x07esc\x1b[31mred\x1b[0m\x7fdel", + expected: "nullbellesc[31mred[0mdel", + }, + { + name: "Zero-width and format characters stripped", + input: "hidden\u200Binjection\u200Cand\u200Djoiner\uFEFFbom\u202Ebidi\U000E0001tag", + expected: "hiddeninjectionandjoinerbombiditag", + }, + { + name: "Space-like unicode characters stripped", + input: "nbsp\u00A0space\u2000enquad\u2001emquad\u2009thin\u202Fnarrow\u3000ideo\u1680ogham\u2028lsep\u2029psep", + expected: "nbspspaceenquademquadthinnarrowideooghamlseppsep", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := filterText(tc.input) + if got != tc.expected { + t.Errorf("filterText(%q) = %q, expected %q", tc.input, got, tc.expected) + } + }) + } +} + +func TestSanitizeMessagesWithInvisibles(t *testing.T) { + msgs := []Message{ + { + Role: "assistant", + ToolCalls: []ToolCall{ + { + ID: "call_1", + Type: "function", + Function: struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + }{ + Name: "shell_exec", + Arguments: "{\"command\": \"cat\u200B \u00A0file.txt\"}", + }, + }, + }, + }, + { + Role: "tool", + ToolCallID: "call_1", + Content: strp("output\u200B\x00with\u00A0invisible\r\nexit: 0"), + }, + } + + sanitizeMessages(msgs) + + tcArgs := msgs[0].ToolCalls[0].Function.Arguments + if strings.Contains(tcArgs, "\u200B") || strings.Contains(tcArgs, "\u00A0") { + t.Errorf("Tool call arguments still contain invisible characters: %q", tcArgs) + } + + toolContent := *msgs[1].Content + if strings.Contains(toolContent, "\u200B") || strings.Contains(toolContent, "\x00") || strings.Contains(toolContent, "\u00A0") || strings.Contains(toolContent, "\r") { + t.Errorf("Tool content still contains invisible characters: %q", toolContent) + } + if !strings.Contains(toolContent, "outputwithinvisible\nexit: 0") { + t.Errorf("Tool content unexpected: %q", toolContent) + } +} + diff --git a/mb b/mb index ed7d410..9777707 100755 --- a/mb +++ b/mb @@ -21,9 +21,18 @@ sub sp { my $p = ''; $p =~ s/^\s+|\s+$//g; length($p) ? $p : $DEF_SP; } +sub filter_text { my $s = shift // ''; $s =~ s/[^\x20\t\n\p{L}\p{N}\p{P}\p{S}\p{M}]//g; $s } + sub T { my ($n, $d, $p) = @_; {type=>'function', function=>{name=>$n, description=>$d, parameters=>{type=>'object', properties=>$p, required=>[keys %$p]}}} } -sub sanitize_msgs { for my $m (@{$_[0]}) { if (ref $m eq 'HASH' && $m->{tool_calls}) { for my $tc (@{$m->{tool_calls}}) { my $a = eval { decode_json($tc->{function}{arguments} // '{}') }; $tc->{function}{arguments} = encode_json({invalid_raw => $tc->{function}{arguments} // ''}) if ref $a ne 'HASH'; } } } } +sub sanitize_msgs { for my $m (@{$_[0]}) { if (ref $m eq 'HASH') { + if ($m->{tool_calls}) { for my $tc (@{$m->{tool_calls}}) { + $tc->{function}{arguments} = filter_text($tc->{function}{arguments}); + my $a = eval { decode_json($tc->{function}{arguments} // '{}') }; + $tc->{function}{arguments} = encode_json({invalid_raw => $tc->{function}{arguments} // ''}) if ref $a ne 'HASH'; + } } + elsif ($m->{role} && $m->{role} eq 'tool' && defined $m->{content}) { $m->{content} = filter_text($m->{content}); } +} } } sub llm { my ($c, $msgs) = @_; # one non-streaming chat completion sanitize_msgs($msgs); @@ -43,10 +52,12 @@ sub llm { my ($c, $msgs) = @_; # one non-streaming chat completion die "API error: " . (length($rb) ? "$rb (HTTP $r->{status})" : ($r->{reason} || "HTTP $r->{status}")) . "\n"; } sub shell_exec { my ($cmd, $t) = @_; # run a command under a hard timeout + $cmd = filter_text($cmd); my $out = ''; eval { local $SIG{ALRM} = sub { alarm 0; die "timeout\n" }; alarm $t; $out = `$cmd 2>&1`; alarm 0; }; $out =~ s/\s+$//; utf8::decode($out); + $out = filter_text($out); $@ ? "$out\n[timeout after ${t}s]\nexit: -1" : "$out\nexit: " . ($? >> 8); } sub last_assistant { for my $m (reverse @{$_[0]}) { return $m->{content} if $m->{role} eq 'assistant' && defined $m->{content} && length $m->{content}; } '' } @@ -66,12 +77,14 @@ sub AL { my ($c, $msgs, $sp, $depth) = @_; # the agentic loop: LLM <-> tools unt last unless $tcs && @$tcs; for my $tc (@$tcs) { my $fn = $tc->{function}{name}; + $tc->{function}{arguments} = filter_text($tc->{function}{arguments}); my $a = eval { decode_json($tc->{function}{arguments} // '{}') }; my $res; if (ref $a ne 'HASH') { $tc->{function}{arguments} = encode_json({invalid_raw => $tc->{function}{arguments} // ''}); $res = "bad JSON args for $fn: " . ($tc->{function}{arguments} // ''); } - elsif ($fn eq 'shell_exec') { $res = shell_exec($a->{command} // '', $c->{shell_timeout}); } - elsif ($fn eq 'run_subagent') { $res = $depth >= 5 ? '[subagent depth limit (5) reached, child not spawned]' : last_assistant(AL($c, [{role=>'system', content=>"$sp\n\nImportant: this is a child agent"}, {role=>'user', content=>$a->{prompt} // ''}], $sp, $depth + 1)); } + elsif ($fn eq 'shell_exec') { $res = shell_exec(filter_text($a->{command} // ''), $c->{shell_timeout}); } + elsif ($fn eq 'run_subagent') { $res = $depth >= 5 ? '[subagent depth limit (5) reached, child not spawned]' : filter_text(last_assistant(AL($c, [{role=>'system', content=>"$sp\n\nImportant: this is a child agent"}, {role=>'user', content=>filter_text($a->{prompt} // '')}], $sp, $depth + 1))); } else { $res = "unknown tool: $fn"; } + $res = filter_text($res); print "[tool] $fn: $res\n"; push @$msgs, {role=>'tool', tool_call_id=>$tc->{id}, content=>$res}; }