feat(agentx): jail workspace tools and audit hooks
Keep read, write, edit, list, and grep inside the workspace root, including symlink targets, and let hooks block or rewrite a call.
This commit is contained in:
weichen
@@ -0,0 +1,198 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type readFile struct{ root string }
|
||||
|
||||
func (t *readFile) Name() string { return "read_file" }
|
||||
func (t *readFile) Description() string {
|
||||
return "Read a UTF-8 text file inside the workspace."
|
||||
}
|
||||
func (t *readFile) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","required":["path"],"properties":{"path":{"type":"string"}}}`)
|
||||
}
|
||||
func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, error) {
|
||||
rel, err := readArg(args, "path")
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
path, err := mustPath(t.root, rel)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
if len(b) > maxReadBytes {
|
||||
b = b[:maxReadBytes]
|
||||
}
|
||||
return Result{Content: string(b)}, nil
|
||||
}
|
||||
|
||||
type writeFile struct{ root string }
|
||||
|
||||
func (t *writeFile) Name() string { return "write_file" }
|
||||
func (t *writeFile) Description() string {
|
||||
return "Create or replace a UTF-8 text file inside the workspace."
|
||||
}
|
||||
func (t *writeFile) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","required":["path","content"],"properties":{"path":{"type":"string"},"content":{"type":"string"}}}`)
|
||||
}
|
||||
func (t *writeFile) Execute(_ context.Context, args json.RawMessage) (Result, error) {
|
||||
var obj struct {
|
||||
Path string `json:"path"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &obj); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
path, err := mustPath(t.root, obj.Path)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
if err := os.WriteFile(path, []byte(obj.Content), 0o644); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
return Result{Content: "wrote " + obj.Path}, nil
|
||||
}
|
||||
|
||||
type editFile struct{ root string }
|
||||
|
||||
func (t *editFile) Name() string { return "edit_file" }
|
||||
func (t *editFile) Description() string {
|
||||
return "Replace one unique exact text span in a workspace file."
|
||||
}
|
||||
func (t *editFile) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","required":["path","old","new"],"properties":{"path":{"type":"string"},"old":{"type":"string"},"new":{"type":"string"}}}`)
|
||||
}
|
||||
func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, error) {
|
||||
var obj struct {
|
||||
Path string `json:"path"`
|
||||
Old string `json:"old"`
|
||||
New string `json:"new"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &obj); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if obj.Old == "" {
|
||||
return Result{Content: "old text is empty", IsError: true}, nil
|
||||
}
|
||||
path, err := mustPath(t.root, obj.Path)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
text := string(b)
|
||||
n := strings.Count(text, obj.Old)
|
||||
if n != 1 {
|
||||
return Result{Content: fmt.Sprintf("expected one match, found %d", n), IsError: true}, nil
|
||||
}
|
||||
updated := strings.Replace(text, obj.Old, obj.New, 1)
|
||||
if err := os.WriteFile(path, []byte(updated), 0o644); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
return Result{Content: "edited " + obj.Path}, nil
|
||||
}
|
||||
|
||||
type listDir struct{ root string }
|
||||
|
||||
func (t *listDir) Name() string { return "list_dir" }
|
||||
func (t *listDir) Description() string {
|
||||
return "List names in a workspace directory."
|
||||
}
|
||||
func (t *listDir) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)
|
||||
}
|
||||
func (t *listDir) Execute(_ context.Context, args json.RawMessage) (Result, error) {
|
||||
rel, _ := readArg(args, "path")
|
||||
if rel == "" {
|
||||
rel = "."
|
||||
}
|
||||
path, err := mustPath(t.root, rel)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
entries, err := os.ReadDir(path)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
names := make([]string, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
if e.IsDir() {
|
||||
name += "/"
|
||||
}
|
||||
names = append(names, name)
|
||||
}
|
||||
return Result{Content: strings.Join(names, "\n")}, nil
|
||||
}
|
||||
|
||||
type grepTool struct{ root string }
|
||||
|
||||
func (t *grepTool) Name() string { return "grep" }
|
||||
func (t *grepTool) Description() string {
|
||||
return "Search workspace text files for a literal substring."
|
||||
}
|
||||
func (t *grepTool) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","required":["pattern"],"properties":{"pattern":{"type":"string"},"path":{"type":"string"}}}`)
|
||||
}
|
||||
func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, error) {
|
||||
var obj struct {
|
||||
Pattern string `json:"pattern"`
|
||||
Path string `json:"path"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &obj); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if obj.Path == "" {
|
||||
obj.Path = "."
|
||||
}
|
||||
start, err := mustPath(t.root, obj.Path)
|
||||
if err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
var hits []string
|
||||
err = filepath.WalkDir(start, func(path string, d os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil || d.IsDir() {
|
||||
return nil
|
||||
}
|
||||
if len(hits) >= 50 {
|
||||
return fmt.Errorf("stop")
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil || len(b) > maxReadBytes {
|
||||
return nil
|
||||
}
|
||||
rel, _ := filepath.Rel(t.root, path)
|
||||
for i, line := range strings.Split(string(b), "\n") {
|
||||
if strings.Contains(line, obj.Pattern) {
|
||||
hits = append(hits, fmt.Sprintf("%s:%d:%s", filepath.ToSlash(rel), i+1, line))
|
||||
if len(hits) >= 50 {
|
||||
return fmt.Errorf("stop")
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil && err.Error() != "stop" {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
if len(hits) == 0 {
|
||||
return Result{Content: "no matches"}, nil
|
||||
}
|
||||
return Result{Content: strings.Join(hits, "\n")}, nil
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Resolve joins rel to root and rejects paths that escape the workspace,
|
||||
// including symlink targets outside root. Missing files are allowed when their
|
||||
// existing ancestors stay inside root.
|
||||
func Resolve(root, rel string) (string, error) {
|
||||
rootAbs, err := filepath.Abs(root)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
rootAbs = filepath.Clean(rootAbs)
|
||||
resolvedRoot, err := filepath.EvalSymlinks(rootAbs)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
rel = strings.TrimSpace(rel)
|
||||
if rel == "" {
|
||||
rel = "."
|
||||
}
|
||||
if filepath.IsAbs(rel) {
|
||||
inside, err := filepath.Rel(rootAbs, filepath.Clean(rel))
|
||||
if err != nil || escapes(inside) {
|
||||
return "", fmt.Errorf("path escapes workspace")
|
||||
}
|
||||
rel = inside
|
||||
}
|
||||
rel = filepath.Clean(rel)
|
||||
if escapes(rel) {
|
||||
return "", fmt.Errorf("path escapes workspace")
|
||||
}
|
||||
if rel == "." {
|
||||
return resolvedRoot, nil
|
||||
}
|
||||
current := resolvedRoot
|
||||
parts := strings.Split(rel, string(os.PathSeparator))
|
||||
for i, part := range parts {
|
||||
if part == "" || part == "." {
|
||||
continue
|
||||
}
|
||||
if part == ".." {
|
||||
return "", fmt.Errorf("path escapes workspace")
|
||||
}
|
||||
next := filepath.Join(current, part)
|
||||
info, err := os.Lstat(next)
|
||||
if os.IsNotExist(err) {
|
||||
tail := filepath.Join(parts[i:]...)
|
||||
return filepath.Join(current, tail), nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 {
|
||||
target, err := filepath.EvalSymlinks(next)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("path escapes workspace")
|
||||
}
|
||||
relTarget, err := filepath.Rel(resolvedRoot, target)
|
||||
if err != nil || escapes(relTarget) {
|
||||
return "", fmt.Errorf("path escapes workspace")
|
||||
}
|
||||
current = target
|
||||
continue
|
||||
}
|
||||
current = next
|
||||
}
|
||||
relFinal, err := filepath.Rel(resolvedRoot, current)
|
||||
if err != nil || escapes(relFinal) {
|
||||
return "", fmt.Errorf("path escapes workspace")
|
||||
}
|
||||
return current, nil
|
||||
}
|
||||
|
||||
func escapes(rel string) bool {
|
||||
if rel == ".." {
|
||||
return true
|
||||
}
|
||||
return strings.HasPrefix(rel, ".."+string(os.PathSeparator))
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Result is the tool payload returned to the model.
|
||||
type Result struct {
|
||||
Content string
|
||||
IsError bool
|
||||
}
|
||||
|
||||
// Tool is one executable function with a JSON schema.
|
||||
type Tool interface {
|
||||
Name() string
|
||||
Description() string
|
||||
Parameters() json.RawMessage
|
||||
Execute(ctx context.Context, args json.RawMessage) (Result, error)
|
||||
}
|
||||
|
||||
// BeforeHook can block a call after schema validation.
|
||||
type BeforeHook func(ctx context.Context, name string, args json.RawMessage) (block bool, reason string)
|
||||
|
||||
// AfterHook can rewrite a tool result before it is returned.
|
||||
type AfterHook func(ctx context.Context, name string, args json.RawMessage, result Result) Result
|
||||
|
||||
// Registry executes jailed tools and appends an audit line per call.
|
||||
type Registry struct {
|
||||
root string
|
||||
tools map[string]Tool
|
||||
order []string
|
||||
Before BeforeHook
|
||||
After AfterHook
|
||||
Audit io.Writer
|
||||
}
|
||||
|
||||
func NewRegistry(root string) *Registry {
|
||||
return &Registry{root: root, tools: map[string]Tool{}}
|
||||
}
|
||||
|
||||
func (r *Registry) Root() string { return r.root }
|
||||
|
||||
func (r *Registry) Register(t Tool) {
|
||||
if _, ok := r.tools[t.Name()]; !ok {
|
||||
r.order = append(r.order, t.Name())
|
||||
}
|
||||
r.tools[t.Name()] = t
|
||||
}
|
||||
|
||||
func (r *Registry) Decls() []Decl {
|
||||
out := make([]Decl, 0, len(r.order))
|
||||
for _, name := range r.order {
|
||||
t := r.tools[name]
|
||||
out = append(out, Decl{Name: t.Name(), Description: t.Description(), Parameters: t.Parameters()})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Decl is the schema view of a tool.
|
||||
type Decl struct {
|
||||
Name string
|
||||
Description string
|
||||
Parameters json.RawMessage
|
||||
}
|
||||
|
||||
// Execute validates arguments, runs hooks, and audits the outcome.
|
||||
func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessage) Result {
|
||||
t, ok := r.tools[name]
|
||||
if !ok {
|
||||
res := Result{Content: "unknown tool " + name, IsError: true}
|
||||
r.audit(name, false, true, "unknown")
|
||||
return res
|
||||
}
|
||||
if len(args) == 0 {
|
||||
args = []byte(`{}`)
|
||||
}
|
||||
if err := validateArgs(t.Parameters(), args); err != nil {
|
||||
res := Result{Content: err.Error(), IsError: true}
|
||||
r.audit(name, false, true, "schema")
|
||||
return res
|
||||
}
|
||||
if r.Before != nil {
|
||||
if block, reason := r.Before(ctx, name, args); block {
|
||||
if reason == "" {
|
||||
reason = "blocked"
|
||||
}
|
||||
res := Result{Content: reason, IsError: true}
|
||||
r.audit(name, true, true, reason)
|
||||
return res
|
||||
}
|
||||
}
|
||||
res, err := t.Execute(ctx, args)
|
||||
if err != nil {
|
||||
res = Result{Content: err.Error(), IsError: true}
|
||||
}
|
||||
if r.After != nil {
|
||||
res = r.After(ctx, name, args, res)
|
||||
}
|
||||
r.audit(name, false, res.IsError, "")
|
||||
return res
|
||||
}
|
||||
|
||||
func (r *Registry) audit(name string, blocked, isError bool, note string) {
|
||||
if r.Audit == nil {
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(r.Audit, "%s\t%s\tblocked=%t\terror=%t\t%s\n", time.Now().UTC().Format(time.RFC3339Nano), name, blocked, isError, note)
|
||||
}
|
||||
|
||||
func validateArgs(schema, args json.RawMessage) error {
|
||||
var spec struct {
|
||||
Type string `json:"type"`
|
||||
Required []string `json:"required"`
|
||||
Properties map[string]json.RawMessage `json:"properties"`
|
||||
}
|
||||
if len(schema) > 0 {
|
||||
if err := json.Unmarshal(schema, &spec); err != nil {
|
||||
return fmt.Errorf("invalid tool schema")
|
||||
}
|
||||
}
|
||||
var obj map[string]any
|
||||
if err := json.Unmarshal(args, &obj); err != nil {
|
||||
return fmt.Errorf("arguments must be a JSON object")
|
||||
}
|
||||
if spec.Type != "" && spec.Type != "object" {
|
||||
return fmt.Errorf("unsupported schema type %s", spec.Type)
|
||||
}
|
||||
for _, req := range spec.Required {
|
||||
v, ok := obj[req]
|
||||
if !ok || v == nil {
|
||||
return fmt.Errorf("missing required argument %s", req)
|
||||
}
|
||||
if s, ok := v.(string); ok && strings.TrimSpace(s) == "" {
|
||||
return fmt.Errorf("missing required argument %s", req)
|
||||
}
|
||||
}
|
||||
for key, raw := range spec.Properties {
|
||||
v, ok := obj[key]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
var prop struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &prop); err != nil || prop.Type == "" {
|
||||
continue
|
||||
}
|
||||
if !matchesType(prop.Type, v) {
|
||||
return fmt.Errorf("argument %s must be %s", key, prop.Type)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func matchesType(kind string, v any) bool {
|
||||
switch kind {
|
||||
case "string":
|
||||
_, ok := v.(string)
|
||||
return ok
|
||||
case "number":
|
||||
_, ok := v.(float64)
|
||||
return ok
|
||||
case "integer":
|
||||
f, ok := v.(float64)
|
||||
return ok && f == float64(int64(f))
|
||||
case "boolean":
|
||||
_, ok := v.(bool)
|
||||
return ok
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// WorkspaceTools registers the document workspace tool set.
|
||||
func WorkspaceTools(root string) *Registry {
|
||||
r := NewRegistry(root)
|
||||
r.Register(&readFile{root: root})
|
||||
r.Register(&writeFile{root: root})
|
||||
r.Register(&editFile{root: root})
|
||||
r.Register(&listDir{root: root})
|
||||
r.Register(&grepTool{root: root})
|
||||
return r
|
||||
}
|
||||
|
||||
func mustPath(root, rel string) (string, error) {
|
||||
return Resolve(root, rel)
|
||||
}
|
||||
|
||||
func readArg(args json.RawMessage, key string) (string, error) {
|
||||
var obj map[string]string
|
||||
if err := json.Unmarshal(args, &obj); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return obj[key], nil
|
||||
}
|
||||
|
||||
const maxReadBytes = 200 << 10
|
||||
@@ -0,0 +1,105 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user