IPI protection

This commit is contained in:
Luxferre
2026-08-18 09:24:42 +03:00
parent e9dd44d97b
commit b8686c82ab
4 changed files with 137 additions and 9 deletions
+3 -1
View File
@@ -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. 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: 3. If there are pending `tool_calls` in the assistant response:
- For each tool call, output a trace log (`[tool call: name(args)]`). - 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`). - 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]`). - Output a trace log of the result (`[tool result: name]`).
- Append tool result message (`role: "tool"`, `tool_call_id`, `content`: result string) to `messages`. - Append tool result message (`role: "tool"`, `tool_call_id`, `content`: result string) to `messages`.
- Loop back to step 1. - Loop back to step 1.
@@ -173,6 +174,7 @@ MicroBantam (`mb`) is a compressed Perl 5 reference implementation of the same a
### Features ### 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 - 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 - 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 <id>` (exact id only, no prefix matching), `/cfg <param> [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 - Session management: `/save`, `/list`, `/load <id>` (exact id only, no prefix matching), `/cfg <param> [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 <id>`, `/cfg`, `/help`) and file input mode - Interactive mode (`/quit`, `/clear`, `/save`, `/list`, `/load <id>`, `/cfg`, `/help`) and file input mode
+29 -5
View File
@@ -22,6 +22,7 @@ import (
"strconv" "strconv"
"strings" "strings"
"time" "time"
"unicode"
"unicode/utf8" "unicode/utf8"
) )
@@ -712,11 +713,27 @@ func cleanMessagesForLLM(msgs []Message) []Message {
return out 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) { func sanitizeMessages(msgs []Message) {
for i := range msgs { for i := range msgs {
if msgs[i].Role == "assistant" && len(msgs[i].ToolCalls) > 0 { if msgs[i].Role == "assistant" && len(msgs[i].ToolCalls) > 0 {
for j := range msgs[i].ToolCalls { for j := range msgs[i].ToolCalls {
tc := &msgs[i].ToolCalls[j] tc := &msgs[i].ToolCalls[j]
tc.Function.Arguments = filterText(tc.Function.Arguments)
astr := tc.Function.Arguments astr := tc.Function.Arguments
var a map[string]any var a map[string]any
if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil { if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil {
@@ -724,6 +741,8 @@ func sanitizeMessages(msgs []Message) {
tc.Function.Arguments = string(fixed) 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 { func shell(ctx context.Context, cmd string, timeout int) string {
cmd = filterText(cmd)
cmdCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second) cmdCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
defer cancel() defer cancel()
c := exec.CommandContext(cmdCtx, "sh", "-c", cmd) c := exec.CommandContext(cmdCtx, "sh", "-c", cmd)
c.WaitDelay = 100 * time.Millisecond c.WaitDelay = 100 * time.Millisecond
out, err := c.CombinedOutput() out, err := c.CombinedOutput()
res := strings.TrimSpace(string(out)) res := strings.TrimSpace(filterText(string(out)))
if ctx.Err() != nil { if ctx.Err() != nil {
return "[interrupted]\n\nexit: -1" 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() } if u.Cached() > 0 { turnUsage.CachedTokens = u.Cached() }
for j := range m.ToolCalls { for j := range m.ToolCalls {
tc := &m.ToolCalls[j] tc := &m.ToolCalls[j]
tc.Function.Arguments = filterText(tc.Function.Arguments)
astr := tc.Function.Arguments astr := tc.Function.Arguments
var a map[string]any var a map[string]any
if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil { 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 } if len(m.ToolCalls) == 0 { done = true; break }
for _, tc := range m.ToolCalls { for _, tc := range m.ToolCalls {
if err := ctx.Err(); err != nil { return msgs, turnUsage, err } 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)) fmt.Println(c(fmt.Sprintf("[tool call: %s(%s)]", fn, astr), 33))
res, sty := "", 2 res, sty := "", 2
var a map[string]any var a map[string]any
@@ -1045,10 +1066,12 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
switch fn { switch fn {
case "shell_exec": case "shell_exec":
cmd, _ := a["command"].(string) cmd, _ := a["command"].(string)
cmd = filterText(cmd)
res = shell(ctx, cmd, cfg.ShellTimeout) res = shell(ctx, cmd, cfg.ShellTimeout)
if err := ctx.Err(); err != nil { return msgs, turnUsage, err } if err := ctx.Err(); err != nil { return msgs, turnUsage, err }
case "run_subagent": case "run_subagent":
pr, _ := a["prompt"].(string) pr, _ := a["prompt"].(string)
pr = filterText(pr)
if depth >= MAX_DEPTH { if depth >= MAX_DEPTH {
res, sty = fmt.Sprintf("[subagent depth limit (%d) reached, child not spawned]", MAX_DEPTH), 31 res, sty = fmt.Sprintf("[subagent depth limit (%d) reached, child not spawned]", MAX_DEPTH), 31
} else { } else {
@@ -1064,13 +1087,14 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
} else { } else {
turnUsage.CompletionTokens += subu.CompletionTokens turnUsage.CompletionTokens += subu.CompletionTokens
turnUsage.TotalTokens += subu.TotalTokens turnUsage.TotalTokens += subu.TotalTokens
res, sty = last(subr), 2 res, sty = filterText(last(subr)), 2
} }
} }
default: default:
res, sty = "Unknown tool: " + fn, 31 res, sty = "Unknown tool: " + fn, 31
} }
} }
res = filterText(res)
fmt.Println(c("[tool result: "+fn+"]", 32) + "\n" + c(res, sty) + "\n") fmt.Println(c("[tool result: "+fn+"]", 32) + "\n" + c(res, sty) + "\n")
msgs = append(msgs, Message{Role: "tool", ToolCallID: tc.ID, Content: strp(res)}) msgs = append(msgs, Message{Role: "tool", ToolCallID: tc.ID, Content: strp(res)})
} }
@@ -1422,7 +1446,7 @@ func main() {
} }
u := strings.TrimSpace(string(data)) u := strings.TrimSpace(string(data))
if strings.HasPrefix(u, "!") { if strings.HasPrefix(u, "!") {
cmd := strings.TrimSpace(strings.TrimPrefix(u, "!")) cmd := filterText(strings.TrimSpace(strings.TrimPrefix(u, "!")))
if cmd != "" { if cmd != "" {
astr, _ := json.Marshal(map[string]string{"command": cmd}) astr, _ := json.Marshal(map[string]string{"command": cmd})
fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33)) fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
@@ -1551,7 +1575,7 @@ func main() {
} }
continue continue
case strings.HasPrefix(u, "!"): case strings.HasPrefix(u, "!"):
cmd := strings.TrimSpace(strings.TrimPrefix(u, "!")) cmd := filterText(strings.TrimSpace(strings.TrimPrefix(u, "!")))
if cmd != "" { if cmd != "" {
astr, _ := json.Marshal(map[string]string{"command": cmd}) astr, _ := json.Marshal(map[string]string{"command": cmd})
fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33)) fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
+89
View File
@@ -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)
}
}
+16 -3
View File
@@ -21,9 +21,18 @@ sub sp { my $p = '';
$p =~ s/^\s+|\s+$//g; $p =~ s/^\s+|\s+$//g;
length($p) ? $p : $DEF_SP; } 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 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 sub llm { my ($c, $msgs) = @_; # one non-streaming chat completion
sanitize_msgs($msgs); 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"; } 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 sub shell_exec { my ($cmd, $t) = @_; # run a command under a hard timeout
$cmd = filter_text($cmd);
my $out = ''; my $out = '';
eval { local $SIG{ALRM} = sub { alarm 0; die "timeout\n" }; alarm $t; $out = `$cmd 2>&1`; alarm 0; }; eval { local $SIG{ALRM} = sub { alarm 0; die "timeout\n" }; alarm $t; $out = `$cmd 2>&1`; alarm 0; };
$out =~ s/\s+$//; $out =~ s/\s+$//;
utf8::decode($out); utf8::decode($out);
$out = filter_text($out);
$@ ? "$out\n[timeout after ${t}s]\nexit: -1" : "$out\nexit: " . ($? >> 8); } $@ ? "$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}; } '' } 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; last unless $tcs && @$tcs;
for my $tc (@$tcs) { for my $tc (@$tcs) {
my $fn = $tc->{function}{name}; my $fn = $tc->{function}{name};
$tc->{function}{arguments} = filter_text($tc->{function}{arguments});
my $a = eval { decode_json($tc->{function}{arguments} // '{}') }; my $a = eval { decode_json($tc->{function}{arguments} // '{}') };
my $res; 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} // ''); } 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 '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]' : 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 '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"; } else { $res = "unknown tool: $fn"; }
$res = filter_text($res);
print "[tool] $fn: $res\n"; print "[tool] $fn: $res\n";
push @$msgs, {role=>'tool', tool_call_id=>$tc->{id}, content=>$res}; push @$msgs, {role=>'tool', tool_call_id=>$tc->{id}, content=>$res};
} }