106 lines
3.2 KiB
Go
106 lines
3.2 KiB
Go
package tools
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestSchemaRejectsMissingAndWrongType(t *testing.T) {
|
|
dir := t.TempDir()
|
|
reg := WorkspaceTools(dir)
|
|
missing := reg.Execute(context.Background(), "read_file", []byte(`{}`))
|
|
if !missing.IsError || !strings.Contains(missing.Content, "path") {
|
|
t.Fatalf("missing %#v", missing)
|
|
}
|
|
wrong := reg.Execute(context.Background(), "read_file", []byte(`{"path":1}`))
|
|
if !wrong.IsError || !strings.Contains(wrong.Content, "string") {
|
|
t.Fatalf("type %#v", wrong)
|
|
}
|
|
}
|
|
|
|
func TestPathEscapeAndSymlink(t *testing.T) {
|
|
dir := t.TempDir()
|
|
outside := t.TempDir()
|
|
secret := filepath.Join(outside, "secret.txt")
|
|
if err := os.WriteFile(secret, []byte("hidden"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
link := filepath.Join(dir, "leak")
|
|
if err := os.Symlink(secret, link); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
reg := WorkspaceTools(dir)
|
|
up := reg.Execute(context.Background(), "read_file", []byte(`{"path":"../secret.txt"}`))
|
|
if !up.IsError {
|
|
t.Fatalf("parent escape succeeded: %s", up.Content)
|
|
}
|
|
via := reg.Execute(context.Background(), "read_file", []byte(`{"path":"leak"}`))
|
|
if !via.IsError || strings.Contains(via.Content, "hidden") {
|
|
t.Fatalf("symlink escape: %#v", via)
|
|
}
|
|
absOutside := reg.Execute(context.Background(), "write_file", []byte(`{"path":"`+secret+`","content":"no"}`))
|
|
if !absOutside.IsError {
|
|
t.Fatal("absolute outside write succeeded")
|
|
}
|
|
}
|
|
|
|
func TestWriteReadEditGrep(t *testing.T) {
|
|
dir := t.TempDir()
|
|
reg := WorkspaceTools(dir)
|
|
wrote := reg.Execute(context.Background(), "write_file", []byte(`{"path":"docs/a.md","content":"alpha beta"}`))
|
|
if wrote.IsError {
|
|
t.Fatal(wrote.Content)
|
|
}
|
|
got := reg.Execute(context.Background(), "read_file", []byte(`{"path":"docs/a.md"}`))
|
|
if got.IsError || got.Content != "alpha beta" {
|
|
t.Fatalf("read %#v", got)
|
|
}
|
|
edited := reg.Execute(context.Background(), "edit_file", []byte(`{"path":"docs/a.md","old":"beta","new":"gamma"}`))
|
|
if edited.IsError {
|
|
t.Fatal(edited.Content)
|
|
}
|
|
found := reg.Execute(context.Background(), "grep", []byte(`{"pattern":"gamma"}`))
|
|
if found.IsError || !strings.Contains(found.Content, "docs/a.md:1:") {
|
|
t.Fatalf("grep %#v", found)
|
|
}
|
|
listed := reg.Execute(context.Background(), "list_dir", []byte(`{"path":"docs"}`))
|
|
if !strings.Contains(listed.Content, "a.md") {
|
|
t.Fatalf("list %#v", listed)
|
|
}
|
|
}
|
|
|
|
func TestBeforeHookBlocksAndAudits(t *testing.T) {
|
|
dir := t.TempDir()
|
|
reg := WorkspaceTools(dir)
|
|
var audit bytes.Buffer
|
|
reg.Audit = &audit
|
|
reg.Before = func(ctx context.Context, name string, args json.RawMessage) (bool, string) {
|
|
return true, "policy"
|
|
}
|
|
res := reg.Execute(context.Background(), "list_dir", []byte(`{}`))
|
|
if !res.IsError || res.Content != "policy" {
|
|
t.Fatalf("hook %#v", res)
|
|
}
|
|
if !strings.Contains(audit.String(), "blocked=true") {
|
|
t.Fatalf("audit %s", audit.String())
|
|
}
|
|
}
|
|
|
|
func TestAfterHookRewrites(t *testing.T) {
|
|
dir := t.TempDir()
|
|
reg := WorkspaceTools(dir)
|
|
reg.After = func(ctx context.Context, name string, args json.RawMessage, result Result) Result {
|
|
result.Content = "rewritten"
|
|
return result
|
|
}
|
|
res := reg.Execute(context.Background(), "list_dir", []byte(`{}`))
|
|
if res.Content != "rewritten" {
|
|
t.Fatalf("%#v", res)
|
|
}
|
|
}
|