feat(tools/io): 内置 IO/网络工具(file(路径安全)/http(SSRF防护)/web-search/web-fetch/knowledge)+ 大小截断
This commit is contained in:
1 parent
488cf9c79c
commit
3ce9489f8e
17 files changed
+2223
No files matched your search
@@ -0,0 +1,70 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
// argString returns the string value of args[key], or "" when absent or of
|
||||
// another type.
|
||||
func argString(args map[string]any, key string) string {
|
||||
if args == nil {
|
||||
return ""
|
||||
}
|
||||
s, _ := args[key].(string)
|
||||
return s
|
||||
}
|
||||
|
||||
// argInt returns the integer value of args[key] and whether it was a number.
|
||||
func argInt(args map[string]any, key string) (int, bool) {
|
||||
if args == nil {
|
||||
return 0, false
|
||||
}
|
||||
return toInt(args[key])
|
||||
}
|
||||
|
||||
// argStringMap coerces a JSON object argument to a string map, accepting both
|
||||
// map[string]any and map[string]string inputs, dropping non-string values.
|
||||
func argStringMap(args map[string]any, key string) map[string]string {
|
||||
out := map[string]string{}
|
||||
if args == nil {
|
||||
return out
|
||||
}
|
||||
switch v := args[key].(type) {
|
||||
case map[string]any:
|
||||
for k, val := range v {
|
||||
if s, ok := val.(string); ok {
|
||||
out[k] = s
|
||||
}
|
||||
}
|
||||
case map[string]string:
|
||||
for k, val := range v {
|
||||
out[k] = val
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// toInt converts a decoded-JSON number (float64 / json.Number) or a Go integer
|
||||
// to an int, reporting whether the conversion succeeded.
|
||||
func toInt(v any) (int, bool) {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x, true
|
||||
case int32:
|
||||
return int(x), true
|
||||
case int64:
|
||||
return int(x), true
|
||||
case float32:
|
||||
return int(x), true
|
||||
case float64:
|
||||
return int(x), true
|
||||
case json.Number:
|
||||
n, err := x.Int64()
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return int(n), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
// Package io implements the built-in I/O and network tools for Cyrene v2
|
||||
// (DevDocs/design/09-插件生态.md). It exposes file, http, web-search, web-fetch
|
||||
// and knowledge as pkg/plugins/sdk.Tool values, wired to the process-internal
|
||||
// config, logger and runtime/memory retriever.
|
||||
//
|
||||
// Build the full configured tool set with New(cfg, opts...). Tools() and
|
||||
// Register(reg) are convenience wrappers that use platform defaults; production
|
||||
// wiring should call New(cfg, opts...) so the tools honor the real DataDir,
|
||||
// search endpoint and knowledge retriever.
|
||||
//
|
||||
// Security constraints
|
||||
//
|
||||
// - file: every path is cleaned and validated to stay inside the configured
|
||||
// data directory or an explicit whitelist; ".." traversal and symlink escape
|
||||
// are rejected; reads and writes are size-capped.
|
||||
// - http: GET/POST only, timeouts, response truncation and SSRF protection
|
||||
// (loopback/private/link-local addresses and localhost are refused unless
|
||||
// the host is on an explicit allowlist).
|
||||
// - web-search: requires a configured search endpoint, otherwise a clear error.
|
||||
// - web-fetch: size- and content-type-limited readable text extraction.
|
||||
// - knowledge: RAG lookup over the runtime memory retriever; a clear message
|
||||
// when no knowledge base is configured.
|
||||
//
|
||||
// All tools are safe for concurrent use, as required by sdk.Tool.
|
||||
package io
|
||||
@@ -0,0 +1,311 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
stdio "io"
|
||||
"io/fs"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// Compile-time assertion that fileTool implements sdk.Tool.
|
||||
var _ sdk.Tool = (*fileTool)(nil)
|
||||
|
||||
// fileTool reads, writes, lists and searches files, confined to the configured
|
||||
// data directory plus any explicit whitelist. Paths are cleaned and re-checked
|
||||
// against those roots, so ".." traversal and symlink escapes are rejected.
|
||||
type fileTool struct {
|
||||
roots []string // absolute, cleaned allowed directories
|
||||
maxRead int64
|
||||
maxWrite int64
|
||||
maxList int
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
func newFileTool(t *tools) *fileTool {
|
||||
roots := make([]string, 0, 1+len(t.fileWhitelist))
|
||||
if t.cfg != nil && t.cfg.DataDir != "" {
|
||||
roots = append(roots, absoluteClean(t.cfg.DataDir))
|
||||
}
|
||||
for _, w := range t.fileWhitelist {
|
||||
roots = append(roots, absoluteClean(w))
|
||||
}
|
||||
return &fileTool{
|
||||
roots: roots,
|
||||
maxRead: t.maxReadBytes,
|
||||
maxWrite: t.maxWriteBytes,
|
||||
maxList: t.maxList,
|
||||
log: t.log,
|
||||
}
|
||||
}
|
||||
|
||||
// Name implements sdk.Tool.
|
||||
func (f *fileTool) Name() string { return "file" }
|
||||
|
||||
// Description implements sdk.Tool.
|
||||
func (f *fileTool) Description() string {
|
||||
return "Read, write, list or search files, restricted to the configured data directory and explicit whitelist."
|
||||
}
|
||||
|
||||
// Parameters implements sdk.Tool.
|
||||
func (f *fileTool) Parameters() (map[string]any, error) {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"action": map[string]any{"type": "string", "enum": []any{"read", "write", "list", "search"}, "description": "Operation to perform."},
|
||||
"path": map[string]any{"type": "string", "description": "Absolute or data-directory-relative file path."},
|
||||
"content": map[string]any{"type": "string", "description": "Body to write."},
|
||||
"query": map[string]any{"type": "string", "description": "Substring filter for search."},
|
||||
},
|
||||
"required": []any{"action", "path"},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Execute implements sdk.Tool.
|
||||
func (f *fileTool) Execute(ctx context.Context, args any) (any, error) {
|
||||
m, ok := args.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("io: file: args must be an object")
|
||||
}
|
||||
switch action := argString(m, "action"); action {
|
||||
case "read":
|
||||
return f.read(ctx, argString(m, "path"))
|
||||
case "write":
|
||||
return f.write(ctx, argString(m, "path"), argString(m, "content"))
|
||||
case "list":
|
||||
return f.list(ctx, argString(m, "path"))
|
||||
case "search":
|
||||
return f.search(ctx, argString(m, "path"), argString(m, "query"))
|
||||
default:
|
||||
return nil, fmt.Errorf("io: file: unknown action %q", action)
|
||||
}
|
||||
}
|
||||
|
||||
// read returns a file's content, capped at maxRead bytes.
|
||||
func (f *fileTool) read(ctx context.Context, p string) (any, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
real, err := f.resolveExisting(p)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
info, err := os.Stat(real)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: stat: %w", err)
|
||||
}
|
||||
if info.IsDir() {
|
||||
return nil, fmt.Errorf("io: file: %q is a directory", real)
|
||||
}
|
||||
if info.Size() > f.maxRead {
|
||||
return nil, fmt.Errorf("io: file: %q exceeds the %d-byte read limit", real, f.maxRead)
|
||||
}
|
||||
fh, err := os.Open(real)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: open: %w", err)
|
||||
}
|
||||
defer fh.Close()
|
||||
data, err := stdio.ReadAll(stdio.LimitReader(fh, f.maxRead+1))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: read: %w", err)
|
||||
}
|
||||
if int64(len(data)) > f.maxRead {
|
||||
return nil, fmt.Errorf("io: file: %q exceeds the %d-byte read limit", real, f.maxRead)
|
||||
}
|
||||
return map[string]any{"path": real, "size": len(data), "content": string(data)}, nil
|
||||
}
|
||||
|
||||
// write creates or overwrites a file, capped at maxWrite bytes and confined to
|
||||
// the allowed roots (parent dirs are created inside the root only).
|
||||
func (f *fileTool) write(ctx context.Context, p, content string) (any, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
clean, err := f.resolve(p)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if int64(len(content)) > f.maxWrite {
|
||||
return nil, fmt.Errorf("io: file: content exceeds the %d-byte write limit", f.maxWrite)
|
||||
}
|
||||
parent := filepath.Dir(clean)
|
||||
if err := os.MkdirAll(parent, dirPerm); err != nil {
|
||||
return nil, fmt.Errorf("io: file: mkdir: %w", err)
|
||||
}
|
||||
realParent, err := filepath.EvalSymlinks(parent)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: resolve parent: %w", err)
|
||||
}
|
||||
if !f.within(realParent) {
|
||||
return nil, fmt.Errorf("io: file: parent directory escapes the allowed directories")
|
||||
}
|
||||
if info, err := os.Lstat(clean); err == nil && info.Mode()&os.ModeSymlink != 0 {
|
||||
real, err := filepath.EvalSymlinks(clean)
|
||||
if err != nil || !f.within(real) {
|
||||
return nil, fmt.Errorf("io: file: %q escapes the allowed directories", p)
|
||||
}
|
||||
}
|
||||
if err := os.WriteFile(clean, []byte(content), filePerm); err != nil {
|
||||
return nil, fmt.Errorf("io: file: write: %w", err)
|
||||
}
|
||||
return map[string]any{"path": clean, "size": len(content)}, nil
|
||||
}
|
||||
|
||||
// list returns the entries of a directory, capped at maxList.
|
||||
func (f *fileTool) list(ctx context.Context, p string) (any, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
real, err := f.resolveExisting(p)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
info, err := os.Stat(real)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: stat: %w", err)
|
||||
}
|
||||
if !info.IsDir() {
|
||||
return nil, fmt.Errorf("io: file: %q is not a directory", real)
|
||||
}
|
||||
entries, err := os.ReadDir(real)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: readdir: %w", err)
|
||||
}
|
||||
if len(entries) > f.maxList {
|
||||
entries = entries[:f.maxList]
|
||||
}
|
||||
out := make([]map[string]any, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
ei, err := e.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, map[string]any{"name": e.Name(), "size": ei.Size(), "dir": ei.IsDir()})
|
||||
}
|
||||
return map[string]any{"path": real, "entries": out, "count": len(out)}, nil
|
||||
}
|
||||
|
||||
// search walks a base directory (default the primary root) and returns files
|
||||
// whose names contain query, capped at maxList results.
|
||||
func (f *fileTool) search(ctx context.Context, base, query string) (any, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(query) == "" {
|
||||
return nil, fmt.Errorf("io: file: search query is required")
|
||||
}
|
||||
if strings.TrimSpace(base) == "" {
|
||||
if len(f.roots) == 0 {
|
||||
return nil, fmt.Errorf("io: file: no base directory configured")
|
||||
}
|
||||
base = f.roots[0]
|
||||
}
|
||||
real, err := f.resolveExisting(base)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
info, err := os.Stat(real)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: file: stat: %w", err)
|
||||
}
|
||||
if !info.IsDir() {
|
||||
return nil, fmt.Errorf("io: file: %q is not a directory", real)
|
||||
}
|
||||
q := strings.ToLower(query)
|
||||
hits := make([]map[string]any, 0)
|
||||
walkErr := filepath.WalkDir(real, func(pp string, d fs.DirEntry, err error) error {
|
||||
if err != nil || d.IsDir() {
|
||||
return nil
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(d.Name()), q) {
|
||||
return nil
|
||||
}
|
||||
di, err := d.Info()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
rel, rerr := filepath.Rel(real, pp)
|
||||
if rerr != nil {
|
||||
rel = pp
|
||||
}
|
||||
hits = append(hits, map[string]any{"path": rel, "size": di.Size()})
|
||||
if len(hits) >= f.maxList {
|
||||
return fs.SkipAll
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if walkErr != nil && !errors.Is(walkErr, fs.SkipAll) {
|
||||
return nil, fmt.Errorf("io: file: walk: %w", walkErr)
|
||||
}
|
||||
return map[string]any{"base": real, "query": query, "hits": hits, "count": len(hits)}, nil
|
||||
}
|
||||
|
||||
// resolve cleans p and verifies it stays within one of the allowed roots. It
|
||||
// rejects empty paths and any path that escapes via ".." or an absolute target
|
||||
// outside the roots.
|
||||
func (f *fileTool) resolve(p string) (string, error) {
|
||||
if strings.TrimSpace(p) == "" {
|
||||
return "", fmt.Errorf("io: file: path is required")
|
||||
}
|
||||
if len(f.roots) == 0 {
|
||||
return "", fmt.Errorf("io: file: no base directory configured")
|
||||
}
|
||||
clean := filepath.Clean(p)
|
||||
if !filepath.IsAbs(clean) {
|
||||
clean = filepath.Join(f.roots[0], clean)
|
||||
}
|
||||
clean = filepath.Clean(clean)
|
||||
if !f.within(clean) {
|
||||
return "", fmt.Errorf("io: file: %q is outside the allowed directories", p)
|
||||
}
|
||||
return clean, nil
|
||||
}
|
||||
|
||||
// resolveExisting resolves p and verifies the symlink-resolved real path stays
|
||||
// within the roots, guarding against escaping through a symlink. It errors when
|
||||
// the target does not exist.
|
||||
func (f *fileTool) resolveExisting(p string) (string, error) {
|
||||
clean, err := f.resolve(p)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
real, err := filepath.EvalSymlinks(clean)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("io: file: resolve %q: %w", p, err)
|
||||
}
|
||||
if !f.within(real) {
|
||||
f.log.Warn("io file: rejected path escaping allowed directories", "path", p, "resolved", real)
|
||||
return "", fmt.Errorf("io: file: %q escapes the allowed directories", p)
|
||||
}
|
||||
return real, nil
|
||||
}
|
||||
|
||||
// within reports whether p is inside (or equal to) one of the allowed roots.
|
||||
func (f *fileTool) within(p string) bool {
|
||||
for _, root := range f.roots {
|
||||
rel, err := filepath.Rel(root, p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if rel == "." || (rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// absoluteClean absolutizes and cleans p, returning the best-effort absolute
|
||||
// form used as an allowed root.
|
||||
func absoluteClean(p string) string {
|
||||
a, err := filepath.Abs(p)
|
||||
if err != nil {
|
||||
return filepath.Clean(p)
|
||||
}
|
||||
return filepath.Clean(a)
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// fileToolForTest builds a tools set over cfg with opts and returns the file tool.
|
||||
func fileToolForTest(t *testing.T, cfg *config.Config, opts ...Option) *fileTool {
|
||||
t.Helper()
|
||||
return findTool(t, buildTools(t, cfg, opts...), "file").(*fileTool)
|
||||
}
|
||||
|
||||
func TestFileTool_WriteAndRead(t *testing.T) {
|
||||
t.Parallel()
|
||||
dir := t.TempDir()
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||
|
||||
if _, err := ft.Execute(context.Background(), map[string]any{"action": "write", "path": "sub/a.txt", "content": "hello world"}); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
got, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "sub/a.txt"})
|
||||
if err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["content"] != "hello world" {
|
||||
t.Fatalf("content = %q", m["content"])
|
||||
}
|
||||
if m["size"] != len("hello world") {
|
||||
t.Fatalf("size = %v", m["size"])
|
||||
}
|
||||
if m["path"] != filepath.Join(dir, "sub", "a.txt") {
|
||||
t.Fatalf("path = %q", m["path"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_RejectsTraversal(t *testing.T) {
|
||||
t.Parallel()
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: t.TempDir()})
|
||||
cases := []string{
|
||||
"../escape.txt",
|
||||
"../../../../etc/passwd",
|
||||
"sub/../../escape.txt",
|
||||
}
|
||||
for _, p := range cases {
|
||||
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": p})
|
||||
if err == nil {
|
||||
t.Fatalf("expected traversal rejection for %q", p)
|
||||
}
|
||||
}
|
||||
// An absolute path outside the data dir is also rejected.
|
||||
outside := filepath.Join(t.TempDir(), "outside.txt")
|
||||
_, err := ft.Execute(context.Background(), map[string]any{"action": "write", "path": outside, "content": "x"})
|
||||
if err == nil {
|
||||
t.Fatal("expected rejection for absolute path outside the data dir")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_ReadSizeLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
dir := t.TempDir()
|
||||
// Write with the default limit, then read with a tiny cap.
|
||||
if _, err := fileToolForTest(t, &config.Config{DataDir: dir}).
|
||||
Execute(context.Background(), map[string]any{"action": "write", "path": "big.txt", "content": "1234567890"}); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: dir}, WithReadLimit(4))
|
||||
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "big.txt"})
|
||||
if err == nil || !strings.Contains(err.Error(), "read limit") {
|
||||
t.Fatalf("expected read-limit error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_WriteSizeLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: t.TempDir()}, WithWriteLimit(4))
|
||||
_, err := ft.Execute(context.Background(), map[string]any{"action": "write", "path": "x.txt", "content": "too long"})
|
||||
if err == nil || !strings.Contains(err.Error(), "write limit") {
|
||||
t.Fatalf("expected write-limit error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_List(t *testing.T) {
|
||||
t.Parallel()
|
||||
dir := t.TempDir()
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||
for _, name := range []string{"a.txt", "b.txt", "sub"} {
|
||||
p := filepath.Join(dir, name)
|
||||
if name == "sub" {
|
||||
if err := os.MkdirAll(p, dirPerm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
} else if err := os.WriteFile(p, []byte("x"), filePerm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
got, err := ft.Execute(context.Background(), map[string]any{"action": "list", "path": "."})
|
||||
if err != nil {
|
||||
t.Fatalf("list: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["count"].(int) != 3 {
|
||||
t.Fatalf("count = %v", m["count"])
|
||||
}
|
||||
entries := m["entries"].([]map[string]any)
|
||||
found := map[string]bool{}
|
||||
for _, e := range entries {
|
||||
found[e["name"].(string)] = true
|
||||
}
|
||||
for _, want := range []string{"a.txt", "b.txt", "sub"} {
|
||||
if !found[want] {
|
||||
t.Errorf("missing entry %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_Search(t *testing.T) {
|
||||
t.Parallel()
|
||||
dir := t.TempDir()
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||
for _, name := range []string{"report_2026.txt", "notes.md", "report_old.txt"} {
|
||||
if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), filePerm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
got, err := ft.Execute(context.Background(), map[string]any{"action": "search", "query": "report"})
|
||||
if err != nil {
|
||||
t.Fatalf("search: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["count"].(int) != 2 {
|
||||
t.Fatalf("count = %v", m["count"])
|
||||
}
|
||||
if _, err := ft.Execute(context.Background(), map[string]any{"action": "search", "query": ""}); err == nil {
|
||||
t.Fatal("expected error for empty search query")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_NoBaseDirectory(t *testing.T) {
|
||||
t.Parallel()
|
||||
ft := fileToolForTest(t, &config.Config{}) // empty DataDir, no whitelist
|
||||
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "a.txt"})
|
||||
if err == nil || !strings.Contains(err.Error(), "no base directory") {
|
||||
t.Fatalf("expected no-base-directory error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_PathWhitelist(t *testing.T) {
|
||||
t.Parallel()
|
||||
dir := t.TempDir()
|
||||
extra := t.TempDir()
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: dir}, WithFilePathWhitelist(extra))
|
||||
if err := os.WriteFile(filepath.Join(extra, "w.txt"), []byte("hi"), filePerm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": filepath.Join(extra, "w.txt")})
|
||||
if err != nil {
|
||||
t.Fatalf("read whitelisted file: %v", err)
|
||||
}
|
||||
if got.(map[string]any)["content"] != "hi" {
|
||||
t.Fatalf("content = %q", got.(map[string]any)["content"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileTool_RejectsSymlinkEscape(t *testing.T) {
|
||||
t.Parallel()
|
||||
dir := t.TempDir()
|
||||
secret := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(secret, "secret.txt"), []byte("hidden"), filePerm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
link := filepath.Join(dir, "link.txt")
|
||||
if err := os.Symlink(filepath.Join(secret, "secret.txt"), link); err != nil {
|
||||
t.Skipf("symlink unsupported: %v", err)
|
||||
}
|
||||
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "link.txt"})
|
||||
if err == nil {
|
||||
t.Fatal("expected symlink escape to be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// silentLogger is the base logger used by test fixtures so that expected
|
||||
// security-denial WARN records are not dumped into the test output.
|
||||
func silentLogger() *slog.Logger {
|
||||
return logger.New(logger.Options{Out: io.Discard})
|
||||
}
|
||||
|
||||
// buildTools builds a tool set over cfg with a silent logger, applying opts on
|
||||
// top (caller options win over the default discard logger).
|
||||
func buildTools(t *testing.T, cfg *config.Config, opts ...Option) []sdk.Tool {
|
||||
t.Helper()
|
||||
all := append([]Option{WithLogger(silentLogger())}, opts...)
|
||||
set, err := New(cfg, all...)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
return set
|
||||
}
|
||||
|
||||
// findTool returns the tool with the given name from a set, failing the test
|
||||
// when it is absent.
|
||||
func findTool(t *testing.T, set []sdk.Tool, name string) sdk.Tool {
|
||||
t.Helper()
|
||||
for _, tool := range set {
|
||||
if tool.Name() == name {
|
||||
return tool
|
||||
}
|
||||
}
|
||||
t.Fatalf("tool %q not found in set", name)
|
||||
return nil
|
||||
}
|
||||
|
||||
// serverHost returns the bare host (no port, no brackets) of a test server, the
|
||||
// value to put on the SSRF allowlist so the tools may reach it.
|
||||
func serverHost(srv *httptest.Server) string {
|
||||
u, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return u.Hostname()
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
stdio "io"
|
||||
"log/slog"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// Compile-time assertion that httpTool implements sdk.Tool.
|
||||
var _ sdk.Tool = (*httpTool)(nil)
|
||||
|
||||
// httpTool performs a single GET or POST request with headers, body, a timeout
|
||||
// and SSRF protection, truncating the response to a size cap.
|
||||
type httpTool struct {
|
||||
client *http.Client
|
||||
allowedHosts []string
|
||||
maxBytes int64
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
func newHTTPTool(t *tools) *httpTool {
|
||||
return &httpTool{
|
||||
client: t.client,
|
||||
allowedHosts: t.allowedHosts,
|
||||
maxBytes: t.maxHTTPBytes,
|
||||
log: t.log,
|
||||
}
|
||||
}
|
||||
|
||||
// Name implements sdk.Tool.
|
||||
func (h *httpTool) Name() string { return "http" }
|
||||
|
||||
// Description implements sdk.Tool.
|
||||
func (h *httpTool) Description() string {
|
||||
return "Perform an HTTP GET or POST request with headers, body, timeout and SSRF protection."
|
||||
}
|
||||
|
||||
// Parameters implements sdk.Tool.
|
||||
func (h *httpTool) Parameters() (map[string]any, error) {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"method": map[string]any{"type": "string", "enum": []any{"GET", "POST"}, "description": "HTTP method."},
|
||||
"url": map[string]any{"type": "string", "description": "Absolute http(s) URL."},
|
||||
"headers": map[string]any{"type": "object", "description": "Request headers."},
|
||||
"body": map[string]any{"type": "string", "description": "Request body for POST."},
|
||||
"timeout_ms": map[string]any{"type": "integer", "description": "Per-request timeout in milliseconds."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Response size cap in bytes."},
|
||||
},
|
||||
"required": []any{"url"},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Execute implements sdk.Tool.
|
||||
func (h *httpTool) Execute(ctx context.Context, args any) (any, error) {
|
||||
m, ok := args.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("io: http: args must be an object")
|
||||
}
|
||||
method := strings.ToUpper(argString(m, "method"))
|
||||
if method == "" {
|
||||
method = http.MethodGet
|
||||
}
|
||||
if method != http.MethodGet && method != http.MethodPost {
|
||||
return nil, fmt.Errorf("io: http: unsupported method %q", method)
|
||||
}
|
||||
rawURL := argString(m, "url")
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return nil, fmt.Errorf("io: http: url is required")
|
||||
}
|
||||
u, err := validateRequestURL(ctx, rawURL, h.allowedHosts)
|
||||
if err != nil {
|
||||
h.log.Warn("io http: blocked request", "url", rawURL, "err", err)
|
||||
return nil, err
|
||||
}
|
||||
body := argString(m, "body")
|
||||
var reqBody stdio.Reader
|
||||
if body != "" {
|
||||
reqBody = strings.NewReader(body)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, u.String(), reqBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: http: build request: %w", err)
|
||||
}
|
||||
for k, v := range argStringMap(m, "headers") {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
ct := req.Header.Get("Content-Type")
|
||||
if ct != "" {
|
||||
if _, err := validateContentType(ct); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if body != "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
if ms, ok := argInt(m, "timeout_ms"); ok && ms > 0 {
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Duration(ms)*time.Millisecond)
|
||||
defer cancel()
|
||||
req = req.WithContext(ctx)
|
||||
}
|
||||
if h.client == nil {
|
||||
return nil, fmt.Errorf("io: http: no underlying client")
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: http: do: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
data, truncated, err := readLimited(resp.Body, h.maxBytes)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: http: read: %w", err)
|
||||
}
|
||||
ctOut, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||
return map[string]any{
|
||||
"status": resp.StatusCode,
|
||||
"content_type": ctOut,
|
||||
"body": string(data),
|
||||
"truncated": truncated,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// httpToolForTest builds a tools set with opts and returns the http tool.
|
||||
func httpToolForTest(t *testing.T, opts ...Option) *httpTool {
|
||||
t.Helper()
|
||||
return findTool(t, buildTools(t, &config.Config{}, opts...), "http").(*httpTool)
|
||||
}
|
||||
|
||||
func TestHTTPTool_Get(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"ok":true,"path":"`+r.URL.Path+`"}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
|
||||
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL + "/hello"})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["status"] != http.StatusOK {
|
||||
t.Fatalf("status = %v", m["status"])
|
||||
}
|
||||
body := m["body"].(string)
|
||||
if !strings.Contains(body, `"ok":true`) || !strings.Contains(body, "/hello") {
|
||||
t.Fatalf("body = %q", body)
|
||||
}
|
||||
if m["truncated"].(bool) {
|
||||
t.Fatal("unexpected truncation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_Post(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Fatalf("method = %s", r.Method)
|
||||
}
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
w.Header().Set("Content-Type", "text/plain")
|
||||
_, _ = w.Write(body)
|
||||
}))
|
||||
defer srv.Close()
|
||||
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
|
||||
got, err := ht.Execute(context.Background(), map[string]any{
|
||||
"method": "POST",
|
||||
"url": srv.URL,
|
||||
"body": "name=cyrene",
|
||||
"headers": map[string]any{"Content-Type": "application/x-www-form-urlencoded"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
if body := got.(map[string]any)["body"].(string); body != "name=cyrene" {
|
||||
t.Fatalf("echoed body = %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_UnsupportedMethod(t *testing.T) {
|
||||
t.Parallel()
|
||||
ht := httpToolForTest(t)
|
||||
_, err := ht.Execute(context.Background(), map[string]any{"method": "DELETE", "url": "http://example.com"})
|
||||
if err == nil || !strings.Contains(err.Error(), "unsupported method") {
|
||||
t.Fatalf("expected unsupported-method error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_Timeout(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
_, _ = w.Write([]byte("late"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPTimeout(30*time.Millisecond))
|
||||
|
||||
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
||||
if err == nil {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_RejectsBadContentType(t *testing.T) {
|
||||
t.Parallel()
|
||||
ht := httpToolForTest(t)
|
||||
_, err := ht.Execute(context.Background(), map[string]any{
|
||||
"method": "POST",
|
||||
"url": "http://93.184.216.34/",
|
||||
"body": "x",
|
||||
"headers": map[string]any{"Content-Type": "not a real content type {{"},
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "invalid content-type") {
|
||||
t.Fatalf("expected invalid content-type error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_ResponseTruncation(t *testing.T) {
|
||||
t.Parallel()
|
||||
payload := strings.Repeat("a", 5000)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, payload)
|
||||
}))
|
||||
defer srv.Close()
|
||||
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()), WithHTTPResponseLimit(1000))
|
||||
|
||||
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if !m["truncated"].(bool) {
|
||||
t.Fatal("expected truncated=true")
|
||||
}
|
||||
if body := m["body"].(string); len(body) != 1000 {
|
||||
t.Fatalf("body len = %d", len(body))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_SSRFDenied(t *testing.T) {
|
||||
t.Parallel()
|
||||
ht := httpToolForTest(t) // no whitelist
|
||||
cases := []string{
|
||||
"http://127.0.0.1:9/",
|
||||
"http://localhost:9/",
|
||||
"http://10.0.0.1/",
|
||||
"http://172.16.0.1/",
|
||||
"http://192.168.1.1/",
|
||||
"http://169.254.1.1/",
|
||||
"http://0.0.0.0/",
|
||||
"http://[::1]:9/",
|
||||
}
|
||||
for _, u := range cases {
|
||||
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": u})
|
||||
if !errors.Is(err, ErrSSRF) {
|
||||
t.Fatalf("url %q: expected ErrSSRF, got %v", u, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_SSRFAllowedViaWhitelist(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, "ok")
|
||||
}))
|
||||
defer srv.Close()
|
||||
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
|
||||
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute (whitelisted): %v", err)
|
||||
}
|
||||
if body := got.(map[string]any)["body"].(string); body != "ok" {
|
||||
t.Fatalf("body = %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_UnsupportedScheme(t *testing.T) {
|
||||
t.Parallel()
|
||||
ht := httpToolForTest(t)
|
||||
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": "ftp://example.com/x"})
|
||||
if err == nil || !strings.Contains(err.Error(), "unsupported scheme") {
|
||||
t.Fatalf("expected unsupported-scheme error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPTool_MissingURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
ht := httpToolForTest(t)
|
||||
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET"})
|
||||
if err == nil || !strings.Contains(err.Error(), "url is required") {
|
||||
t.Fatalf("expected url-required error, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// KnowledgeRetriever is the subset of the runtime memory retriever the knowledge
|
||||
// tool depends on. *memory.Retriever implements it; nil means no knowledge base
|
||||
// is configured.
|
||||
type KnowledgeRetriever interface {
|
||||
Retrieve(ctx context.Context, personaID, query string, topK int) ([]memory.Result, error)
|
||||
}
|
||||
|
||||
// Compile-time assertion that knowledgeTool implements sdk.Tool.
|
||||
var _ sdk.Tool = (*knowledgeTool)(nil)
|
||||
|
||||
// knowledgeTool answers a question from the configured knowledge base via a
|
||||
// memory retriever, returning ranked results or a clear message when absent.
|
||||
type knowledgeTool struct {
|
||||
retriever KnowledgeRetriever
|
||||
maxResults int
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
func newKnowledgeTool(t *tools) *knowledgeTool {
|
||||
return &knowledgeTool{
|
||||
retriever: t.retriever,
|
||||
maxResults: t.maxResults,
|
||||
log: t.log,
|
||||
}
|
||||
}
|
||||
|
||||
// Name implements sdk.Tool.
|
||||
func (k *knowledgeTool) Name() string { return "knowledge" }
|
||||
|
||||
// Description implements sdk.Tool.
|
||||
func (k *knowledgeTool) Description() string {
|
||||
return "Query the knowledge base and long-term memory store for context relevant to a question."
|
||||
}
|
||||
|
||||
// Parameters implements sdk.Tool.
|
||||
func (k *knowledgeTool) Parameters() (map[string]any, error) {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Query topic."},
|
||||
"persona": map[string]any{"type": "string", "description": "Persona scope id (optional)."},
|
||||
"top_k": map[string]any{"type": "integer", "description": "Max results."},
|
||||
},
|
||||
"required": []any{"query"},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Execute implements sdk.Tool.
|
||||
func (k *knowledgeTool) Execute(ctx context.Context, args any) (any, error) {
|
||||
if k.retriever == nil {
|
||||
return nil, fmt.Errorf("io: knowledge: no knowledge base configured")
|
||||
}
|
||||
m, ok := args.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("io: knowledge: args must be an object")
|
||||
}
|
||||
query := argString(m, "query")
|
||||
if strings.TrimSpace(query) == "" {
|
||||
return nil, fmt.Errorf("io: knowledge: query is required")
|
||||
}
|
||||
persona := argString(m, "persona")
|
||||
topK := k.maxResults
|
||||
if v, ok := argInt(m, "top_k"); ok && v > 0 {
|
||||
topK = v
|
||||
}
|
||||
results, err := k.retriever.Retrieve(ctx, persona, query, topK)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: knowledge: retrieve: %w", err)
|
||||
}
|
||||
out := make([]map[string]any, 0, len(results))
|
||||
for _, r := range results {
|
||||
out = append(out, map[string]any{
|
||||
"id": r.Item.ID,
|
||||
"kind": string(r.Item.Kind),
|
||||
"scope": string(r.Item.Scope),
|
||||
"content": r.Item.Content,
|
||||
"score": r.Score,
|
||||
})
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return map[string]any{
|
||||
"query": query,
|
||||
"count": 0,
|
||||
"note": "no relevant knowledge found",
|
||||
"results": []map[string]any{},
|
||||
}, nil
|
||||
}
|
||||
return map[string]any{"query": query, "count": len(out), "results": out}, nil
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
|
||||
)
|
||||
|
||||
// stubRetriever is a fake KnowledgeRetriever for testing knowledge responses.
|
||||
type stubRetriever struct {
|
||||
results []memory.Result
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stubRetriever) Retrieve(ctx context.Context, personaID, query string, topK int) ([]memory.Result, error) {
|
||||
_ = ctx
|
||||
return s.results, s.err
|
||||
}
|
||||
|
||||
// knowledgeToolForTest builds a tools set with opts and returns the knowledge
|
||||
// tool.
|
||||
func knowledgeToolForTest(t *testing.T, opts ...Option) *knowledgeTool {
|
||||
t.Helper()
|
||||
return findTool(t, buildTools(t, &config.Config{}, opts...), "knowledge").(*knowledgeTool)
|
||||
}
|
||||
|
||||
func TestKnowledge_NoKnowledgeBase(t *testing.T) {
|
||||
t.Parallel()
|
||||
kt := knowledgeToolForTest(t) // no retriever wired
|
||||
_, err := kt.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||
if err == nil || !strings.Contains(err.Error(), "no knowledge base") {
|
||||
t.Fatalf("expected no-knowledge-base error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledge_ReturnsResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
sr := &stubRetriever{results: []memory.Result{
|
||||
{Item: memory.MemoryItem{ID: "1", Content: "Cyrene 是主人", Kind: memory.KindWorldEvent, Scope: memory.ScopeProject}, Score: 0.9},
|
||||
}}
|
||||
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(sr))
|
||||
|
||||
got, err := kt.Execute(context.Background(), map[string]any{"query": "cyrene", "persona": "p1", "top_k": 5})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["count"].(int) != 1 {
|
||||
t.Fatalf("count = %v", m["count"])
|
||||
}
|
||||
results := m["results"].([]map[string]any)
|
||||
if results[0]["content"] != "Cyrene 是主人" {
|
||||
t.Fatalf("content = %v", results[0]["content"])
|
||||
}
|
||||
if results[0]["kind"] != "world_event" {
|
||||
t.Fatalf("kind = %v", results[0]["kind"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledge_EmptyResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(&stubRetriever{}))
|
||||
got, err := kt.Execute(context.Background(), map[string]any{"query": "nothing"})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["count"].(int) != 0 {
|
||||
t.Fatalf("count = %v", m["count"])
|
||||
}
|
||||
if m["note"] != "no relevant knowledge found" {
|
||||
t.Fatalf("note = %v", m["note"])
|
||||
}
|
||||
if _, ok := m["results"].([]map[string]any); !ok {
|
||||
t.Fatalf("results absent or wrong type: %T", m["results"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledge_RetrieverError(t *testing.T) {
|
||||
t.Parallel()
|
||||
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(&stubRetriever{err: errors.New("boom")}))
|
||||
_, err := kt.Execute(context.Background(), map[string]any{"query": "x"})
|
||||
if err == nil || !strings.Contains(err.Error(), "retrieve") {
|
||||
t.Fatalf("expected retrieve error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledge_MissingQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(&stubRetriever{}))
|
||||
_, err := kt.Execute(context.Background(), map[string]any{})
|
||||
if err == nil || !strings.Contains(err.Error(), "query is required") {
|
||||
t.Fatalf("expected query-required error, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrSSRF is returned when a requested URL targets a loopback, private or
|
||||
// link-local address or a localhost hostname that is not on the allowlist.
|
||||
// Use errors.Is to test for it.
|
||||
var ErrSSRF = errors.New("io: disallowed url")
|
||||
|
||||
// validateRequestURL parses rawURL and refuses SSRF-prone targets: it returns an
|
||||
// error wrapping ErrSSRF for loopback/private/link-local/unspecified addresses
|
||||
// and for localhost hostnames, unless the host (or hostname) is on the explicit
|
||||
// allowlist. A hostname that resolves to any denied address is also refused.
|
||||
func validateRequestURL(ctx context.Context, rawURL string, allowed []string) (*url.URL, error) {
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return nil, fmt.Errorf("io: url is required")
|
||||
}
|
||||
u, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: parse url: %w", err)
|
||||
}
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return nil, fmt.Errorf("io: unsupported scheme %q", u.Scheme)
|
||||
}
|
||||
host := hostOnly(u.Host)
|
||||
if host == "" {
|
||||
return nil, fmt.Errorf("io: url has no host")
|
||||
}
|
||||
if isHostAllowed(host, allowed) {
|
||||
return u, nil
|
||||
}
|
||||
if ip := net.ParseIP(host); ip != nil {
|
||||
if isDeniedIP(ip) {
|
||||
return nil, fmt.Errorf("%w: %s", ErrSSRF, host)
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
addrs, err := net.DefaultResolver.LookupHost(ctx, host)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: resolve %q: %w", host, err)
|
||||
}
|
||||
for _, a := range addrs {
|
||||
if ip := net.ParseIP(a); ip != nil && isDeniedIP(ip) {
|
||||
return nil, fmt.Errorf("%w: %q resolves to %q", ErrSSRF, host, a)
|
||||
}
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
// isDeniedIP reports whether ip is a loopback, private, link-local, multicast or
|
||||
// unspecified address — the targets SSRF protection refuses by default.
|
||||
func isDeniedIP(ip net.IP) bool {
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
return ip.IsLoopback() ||
|
||||
ip.IsPrivate() ||
|
||||
ip.IsLinkLocalUnicast() ||
|
||||
ip.IsLinkLocalMulticast() ||
|
||||
ip.IsUnspecified()
|
||||
}
|
||||
|
||||
// isHostAllowed reports whether the host (with any port stripped) matches an
|
||||
// entry on the allowlist.
|
||||
func isHostAllowed(host string, allowed []string) bool {
|
||||
h := strings.ToLower(hostOnly(host))
|
||||
if h == "" {
|
||||
return false
|
||||
}
|
||||
for _, a := range allowed {
|
||||
if strings.EqualFold(hostOnly(a), h) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// hostOnly strips the port (and any IPv6 brackets) from a host:port string.
|
||||
func hostOnly(hostport string) string {
|
||||
h := strings.TrimSpace(hostport)
|
||||
if h == "" {
|
||||
return ""
|
||||
}
|
||||
if strings.HasPrefix(h, "[") {
|
||||
if j := strings.Index(h, "]"); j >= 0 {
|
||||
return h[1:j]
|
||||
}
|
||||
return h
|
||||
}
|
||||
if i := strings.LastIndex(h, ":"); i >= 0 {
|
||||
return h[:i]
|
||||
}
|
||||
return h
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/platform"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// source tags every log record emitted by this package (project rule §10).
|
||||
const source = "plugins/tools/io"
|
||||
|
||||
// Default size, timeout and result limits. Named constants only — no bare
|
||||
// magic numbers (project rule §10).
|
||||
const (
|
||||
defaultMaxReadBytes = 5 << 20 // 5 MiB file read cap
|
||||
defaultMaxWriteBytes = 5 << 20 // 5 MiB file write cap
|
||||
defaultMaxHTTPBytes = 1 << 20 // 1 MiB http response cap
|
||||
defaultMaxFetchBytes = 2 << 20 // 2 MiB web-fetch cap
|
||||
defaultHTTPTimeout = 15 * time.Second
|
||||
defaultMaxResults = 10
|
||||
defaultMaxListEntries = 200
|
||||
maxRedirects = 3
|
||||
filePerm = 0o644
|
||||
dirPerm = 0o755
|
||||
)
|
||||
|
||||
// tools is the constructed, shared state for the whole tool set. A single HTTP
|
||||
// client is shared by http/web-search/web-fetch so the connection pool is
|
||||
// unified (project rule §10); it is never a package-level global.
|
||||
type tools struct {
|
||||
cfg *config.Config
|
||||
client *http.Client
|
||||
allowedHosts []string
|
||||
searchURL string
|
||||
retriever KnowledgeRetriever
|
||||
fileWhitelist []string
|
||||
maxReadBytes int64
|
||||
maxWriteBytes int64
|
||||
maxHTTPBytes int64
|
||||
maxFetchBytes int64
|
||||
httpTimeout time.Duration
|
||||
maxResults int
|
||||
maxList int
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// Option customises the tool set built by New.
|
||||
type Option func(*tools)
|
||||
|
||||
// WithHTTPTimeout sets the per-request timeout. Non-positive values are ignored.
|
||||
func WithHTTPTimeout(d time.Duration) Option {
|
||||
return func(t *tools) {
|
||||
if d > 0 {
|
||||
t.httpTimeout = d
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPClient injects the client shared by the network tools. When nil (the
|
||||
// default) New builds one with a unified timeout and redirect SSRF guard.
|
||||
func WithHTTPClient(c *http.Client) Option {
|
||||
return func(t *tools) {
|
||||
if c != nil {
|
||||
t.client = c
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithAllowedHosts adds hosts to the SSRF allowlist. Hosts here bypass the
|
||||
// loopback/private deny rules (they are the "host whitelist" side of the SSRF
|
||||
// policy), letting the agent reach an internal search endpoint or the runtime.
|
||||
func WithAllowedHosts(hosts ...string) Option {
|
||||
return func(t *tools) { t.allowedHosts = append(t.allowedHosts, hosts...) }
|
||||
}
|
||||
|
||||
// WithSearchEndpoint sets the web-search endpoint (a SearXNG-compatible URL).
|
||||
// WithSearchEndpoint("") clears it so web-search reports a clear error.
|
||||
func WithSearchEndpoint(u string) Option {
|
||||
return func(t *tools) { t.searchURL = u }
|
||||
}
|
||||
|
||||
// WithKnowledgeRetriever wires the knowledge tool to a runtime memory retriever.
|
||||
// WithKnowledgeRetriever(nil) makes knowledge report "no knowledge base".
|
||||
func WithKnowledgeRetriever(r KnowledgeRetriever) Option {
|
||||
return func(t *tools) { t.retriever = r }
|
||||
}
|
||||
|
||||
// WithFilePathWhitelist adds extra directories the file tool may access, in
|
||||
// addition to the configured DataDir.
|
||||
func WithFilePathWhitelist(paths ...string) Option {
|
||||
return func(t *tools) { t.fileWhitelist = append(t.fileWhitelist, paths...) }
|
||||
}
|
||||
|
||||
// WithReadLimit overrides the file read cap (bytes). Non-positive values are
|
||||
// ignored.
|
||||
func WithReadLimit(n int64) Option {
|
||||
return func(t *tools) {
|
||||
if n > 0 {
|
||||
t.maxReadBytes = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithWriteLimit overrides the file write cap (bytes). Non-positive values are
|
||||
// ignored.
|
||||
func WithWriteLimit(n int64) Option {
|
||||
return func(t *tools) {
|
||||
if n > 0 {
|
||||
t.maxWriteBytes = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPResponseLimit overrides the http response cap (bytes). Non-positive
|
||||
// values are ignored.
|
||||
func WithHTTPResponseLimit(n int64) Option {
|
||||
return func(t *tools) {
|
||||
if n > 0 {
|
||||
t.maxHTTPBytes = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithMaxResults caps the number of results returned by web-search and
|
||||
// knowledge. Non-positive values are ignored.
|
||||
func WithMaxResults(n int) Option {
|
||||
return func(t *tools) {
|
||||
if n > 0 {
|
||||
t.maxResults = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithLogger injects the base logger. When nil the package builds one tagged
|
||||
// with source "plugins/tools/io".
|
||||
func WithLogger(l *slog.Logger) Option {
|
||||
return func(t *tools) {
|
||||
if l != nil {
|
||||
t.log = l
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// New builds the full configured tool set over cfg and opts. It returns an error
|
||||
// only when cfg is nil. A nil retriever or empty search endpoint is allowed: the
|
||||
// corresponding tool reports a clear runtime error rather than failing startup.
|
||||
func New(cfg *config.Config, opts ...Option) ([]sdk.Tool, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("io: new: nil config")
|
||||
}
|
||||
t := &tools{
|
||||
cfg: cfg,
|
||||
maxReadBytes: defaultMaxReadBytes,
|
||||
maxWriteBytes: defaultMaxWriteBytes,
|
||||
maxHTTPBytes: defaultMaxHTTPBytes,
|
||||
maxFetchBytes: defaultMaxFetchBytes,
|
||||
httpTimeout: defaultHTTPTimeout,
|
||||
maxResults: defaultMaxResults,
|
||||
maxList: defaultMaxListEntries,
|
||||
log: logger.WithSource(logger.New(logger.Options{}), source),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
if opt != nil {
|
||||
opt(t)
|
||||
}
|
||||
}
|
||||
if t.client == nil {
|
||||
t.client = newHTTPClient(t)
|
||||
}
|
||||
return t.build(), nil
|
||||
}
|
||||
|
||||
// Tools returns the tool set built from platform defaults. Production wiring
|
||||
// should prefer New(cfg, opts...) so the tools honor the real configuration.
|
||||
func Tools() []sdk.Tool {
|
||||
set, err := New(&config.Config{DataDir: platform.DefaultDataDir()})
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return set
|
||||
}
|
||||
|
||||
// Register registers the default tool set (from Tools) with reg. It returns
|
||||
// sdk.ErrNilRegistrar when reg is nil and stops at the first registration error.
|
||||
func Register(reg sdk.Registrar) error {
|
||||
if reg == nil {
|
||||
return sdk.ErrNilRegistrar
|
||||
}
|
||||
for _, t := range Tools() {
|
||||
if err := sdk.RegisterTool(reg, t); err != nil {
|
||||
return fmt.Errorf("io: register %q: %w", t.Name(), err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// build assembles the tools sharing this set's client, limits and logger.
|
||||
func (t *tools) build() []sdk.Tool {
|
||||
return []sdk.Tool{
|
||||
newFileTool(t),
|
||||
newHTTPTool(t),
|
||||
newWebSearchTool(t),
|
||||
newWebFetchTool(t),
|
||||
newKnowledgeTool(t),
|
||||
}
|
||||
}
|
||||
|
||||
// newHTTPClient builds a client with a unified timeout and an SSRF-guarded
|
||||
// redirect policy that refuses to follow into private/loopback targets.
|
||||
func newHTTPClient(t *tools) *http.Client {
|
||||
return &http.Client{
|
||||
Timeout: t.httpTimeout,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= maxRedirects {
|
||||
return fmt.Errorf("io: too many redirects")
|
||||
}
|
||||
if req.URL == nil {
|
||||
return ErrSSRF
|
||||
}
|
||||
if _, err := validateRequestURL(req.Context(), req.URL.String(), t.allowedHosts); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
func TestTools_ReturnsAllFiveTools(t *testing.T) {
|
||||
t.Parallel()
|
||||
set, err := New(&config.Config{DataDir: t.TempDir()})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
if len(set) != 5 {
|
||||
t.Fatalf("expected 5 tools, got %d", len(set))
|
||||
}
|
||||
want := []string{"file", "http", "web-search", "web-fetch", "knowledge"}
|
||||
seen := map[string]bool{}
|
||||
for _, tool := range set {
|
||||
seen[tool.Name()] = true
|
||||
if tool.Description() == "" {
|
||||
t.Errorf("tool %q has an empty description", tool.Name())
|
||||
}
|
||||
params, err := tool.Parameters()
|
||||
if err != nil {
|
||||
t.Errorf("tool %q Parameters error: %v", tool.Name(), err)
|
||||
}
|
||||
if params == nil {
|
||||
t.Errorf("tool %q has nil parameters", tool.Name())
|
||||
}
|
||||
}
|
||||
for _, name := range want {
|
||||
if !seen[name] {
|
||||
t.Errorf("missing tool %q", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNew_NilConfigErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := New(nil); err == nil {
|
||||
t.Fatal("expected error for nil config")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegister_RegistersEveryTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
reg := registry.New()
|
||||
if err := Register(reg); err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
for _, name := range []string{"file", "http", "web-search", "web-fetch", "knowledge"} {
|
||||
c, err := reg.Get(name)
|
||||
if err != nil {
|
||||
t.Fatalf("Get(%q): %v", name, err)
|
||||
}
|
||||
if c.Kind != registry.KindTool {
|
||||
t.Errorf("kind of %q = %q, want tool", name, c.Kind)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegister_NilRegistrar(t *testing.T) {
|
||||
t.Parallel()
|
||||
if err := Register(nil); !errors.Is(err, sdk.ErrNilRegistrar) {
|
||||
t.Fatalf("expected sdk.ErrNilRegistrar, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegister_DuplicateErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
reg := registry.New()
|
||||
if err := Register(reg); err != nil {
|
||||
t.Fatalf("first Register: %v", err)
|
||||
}
|
||||
if err := Register(reg); !errors.Is(err, registry.ErrDuplicate) {
|
||||
t.Fatalf("expected registry.ErrDuplicate, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTools_ConcurrentDescribe(t *testing.T) {
|
||||
t.Parallel()
|
||||
set, err := New(&config.Config{DataDir: t.TempDir()})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
// Exercise the read-only contract concurrently to catch any shared-state
|
||||
// mutation (the dispatcher may call Parameters from many goroutines).
|
||||
for _, tool := range set {
|
||||
tool := tool
|
||||
t.Run(tool.Name(), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
for i := 0; i < 10; i++ {
|
||||
if _, err := tool.Parameters(); err != nil {
|
||||
t.Fatalf("Parameters: %v", err)
|
||||
}
|
||||
if tool.Name() == "" || tool.Description() == "" {
|
||||
t.Fatalf("empty identity for %v", tool)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// readLimited reads up to max bytes from r, returning the payload and whether it
|
||||
// was truncated. A non-positive max falls back to the http response cap.
|
||||
func readLimited(r io.Reader, max int64) ([]byte, bool, error) {
|
||||
if max <= 0 {
|
||||
max = defaultMaxHTTPBytes
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(r, max+1))
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
truncated := int64(len(data)) > max
|
||||
if truncated {
|
||||
data = data[:max]
|
||||
}
|
||||
return data, truncated, nil
|
||||
}
|
||||
|
||||
// validateContentType parses a caller-supplied Content-Type, returning the
|
||||
// normalized media type. An empty value is allowed and yields "".
|
||||
func validateContentType(ct string) (string, error) {
|
||||
if strings.TrimSpace(ct) == "" {
|
||||
return "", nil
|
||||
}
|
||||
mt, _, err := mime.ParseMediaType(ct)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("io: invalid content-type %q: %w", ct, err)
|
||||
}
|
||||
return mt, nil
|
||||
}
|
||||
|
||||
// isTextContentType reports whether a media type is acceptable to fetch as text.
|
||||
// An empty (unset) media type is accepted so a server that omits Content-Type is
|
||||
// treated as text; binary media types are refused.
|
||||
func isTextContentType(mt string) bool {
|
||||
mt = strings.ToLower(strings.TrimSpace(mt))
|
||||
if mt == "" {
|
||||
return true
|
||||
}
|
||||
if strings.HasPrefix(mt, "text/") {
|
||||
return true
|
||||
}
|
||||
switch mt {
|
||||
case "application/json", "application/ld+json", "application/xml",
|
||||
"application/xhtml+xml", "application/rss+xml", "application/atom+xml",
|
||||
"application/x-www-form-urlencoded", "application/javascript",
|
||||
"application/x-javascript", "image/svg+xml", "application/graphql":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// isHTML reports whether a media type should be converted from HTML to text.
|
||||
func isHTML(mt string) bool {
|
||||
mt = strings.ToLower(strings.TrimSpace(mt))
|
||||
return mt == "text/html" || mt == "application/xhtml+xml"
|
||||
}
|
||||
|
||||
// htmlToText strips script/style/comment blocks and tags from an HTML document,
|
||||
// decoding a few common entities and collapsing whitespace for readable text.
|
||||
func htmlToText(src string) string {
|
||||
s := dropBlocks(src)
|
||||
var out strings.Builder
|
||||
inTag := false
|
||||
for i := 0; i < len(s); i++ {
|
||||
switch c := s[i]; {
|
||||
case c == '<':
|
||||
inTag = true
|
||||
case c == '>':
|
||||
inTag = false
|
||||
case !inTag:
|
||||
out.WriteByte(c)
|
||||
}
|
||||
}
|
||||
return normalizeSpace(htmlUnescape(out.String()))
|
||||
}
|
||||
|
||||
// dropBlocks removes <script>, <style> and HTML comment regions.
|
||||
func dropBlocks(s string) string {
|
||||
for _, tag := range []string{"script", "style"} {
|
||||
for {
|
||||
start := strings.Index(s, "<"+tag)
|
||||
if start < 0 {
|
||||
break
|
||||
}
|
||||
end := strings.Index(s, "</"+tag+">")
|
||||
if end < 0 {
|
||||
s = s[:start]
|
||||
break
|
||||
}
|
||||
s = s[:start] + s[end+len(tag)+3:]
|
||||
}
|
||||
}
|
||||
for {
|
||||
start := strings.Index(s, "<!--")
|
||||
if start < 0 {
|
||||
break
|
||||
}
|
||||
end := strings.Index(s, "-->")
|
||||
if end < 0 {
|
||||
s = s[:start]
|
||||
break
|
||||
}
|
||||
s = s[:start] + s[end+3:]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// htmlUnescape replaces the small set of entities common in extracted text.
|
||||
func htmlUnescape(s string) string {
|
||||
return strings.NewReplacer(
|
||||
"&", "&",
|
||||
"<", "<",
|
||||
">", ">",
|
||||
""", "\"",
|
||||
"'", "'",
|
||||
" ", " ",
|
||||
).Replace(s)
|
||||
}
|
||||
|
||||
// normalizeSpace collapses runs of whitespace to single spaces and trims ends.
|
||||
func normalizeSpace(s string) string {
|
||||
return strings.Join(strings.Fields(s), " ")
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// Compile-time assertion that webFetchTool implements sdk.Tool.
|
||||
var _ sdk.Tool = (*webFetchTool)(nil)
|
||||
|
||||
// webFetchTool fetches a URL over HTTP and returns its readable text, limited by
|
||||
// size and content type, extracting text from HTML when appropriate.
|
||||
type webFetchTool struct {
|
||||
client *http.Client
|
||||
allowedHosts []string
|
||||
maxBytes int64
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
func newWebFetchTool(t *tools) *webFetchTool {
|
||||
return &webFetchTool{
|
||||
client: t.client,
|
||||
allowedHosts: t.allowedHosts,
|
||||
maxBytes: t.maxFetchBytes,
|
||||
log: t.log,
|
||||
}
|
||||
}
|
||||
|
||||
// Name implements sdk.Tool.
|
||||
func (w *webFetchTool) Name() string { return "web-fetch" }
|
||||
|
||||
// Description implements sdk.Tool.
|
||||
func (w *webFetchTool) Description() string {
|
||||
return "Fetch a URL and return its readable text, limited by size and content type."
|
||||
}
|
||||
|
||||
// Parameters implements sdk.Tool.
|
||||
func (w *webFetchTool) Parameters() (map[string]any, error) {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"url": map[string]any{"type": "string", "description": "Absolute http(s) URL."},
|
||||
"max_bytes": map[string]any{"type": "integer", "description": "Response size cap in bytes."},
|
||||
},
|
||||
"required": []any{"url"},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Execute implements sdk.Tool.
|
||||
func (w *webFetchTool) Execute(ctx context.Context, args any) (any, error) {
|
||||
m, ok := args.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("io: web-fetch: args must be an object")
|
||||
}
|
||||
rawURL := argString(m, "url")
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return nil, fmt.Errorf("io: web-fetch: url is required")
|
||||
}
|
||||
if _, err := validateRequestURL(ctx, rawURL, w.allowedHosts); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if w.client == nil {
|
||||
return nil, fmt.Errorf("io: web-fetch: no underlying client")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: web-fetch: build request: %w", err)
|
||||
}
|
||||
resp, err := w.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: web-fetch: do: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("io: web-fetch: url returned status %d", resp.StatusCode)
|
||||
}
|
||||
mt := "text/plain"
|
||||
if rawCT := resp.Header.Get("Content-Type"); rawCT != "" {
|
||||
var err error
|
||||
mt, _, err = mime.ParseMediaType(rawCT)
|
||||
if err != nil || !isTextContentType(mt) {
|
||||
return nil, fmt.Errorf("io: web-fetch: unexpected content type %q", rawCT)
|
||||
}
|
||||
}
|
||||
max := w.maxBytes
|
||||
if v, ok := argInt(m, "max_bytes"); ok && v > 0 {
|
||||
max = int64(v)
|
||||
}
|
||||
data, truncated, err := readLimited(resp.Body, max)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: web-fetch: read: %w", err)
|
||||
}
|
||||
text := string(data)
|
||||
if isHTML(mt) {
|
||||
text = htmlToText(text)
|
||||
}
|
||||
return map[string]any{
|
||||
"url": rawURL,
|
||||
"content_type": mt,
|
||||
"text": text,
|
||||
"truncated": truncated,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// webFetchToolForTest builds a tools set with opts and returns the web-fetch
|
||||
// tool.
|
||||
func webFetchToolForTest(t *testing.T, opts ...Option) *webFetchTool {
|
||||
t.Helper()
|
||||
return findTool(t, buildTools(t, &config.Config{}, opts...), "web-fetch").(*webFetchTool)
|
||||
}
|
||||
|
||||
func TestWebFetch_TextExtraction(t *testing.T) {
|
||||
t.Parallel()
|
||||
doc := `<!doctype html><html><head><title>T</title>` +
|
||||
`<style>.x{color:red}</style><script>alert(1)</script></head>` +
|
||||
`<body><h1>Hello</h1><p>World & more</p><!-- hidden --></body></html>`
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
_, _ = io.WriteString(w, doc)
|
||||
}))
|
||||
defer srv.Close()
|
||||
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
|
||||
got, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
text := m["text"].(string)
|
||||
if !strings.Contains(text, "Hello") || !strings.Contains(text, "World & more") {
|
||||
t.Fatalf("text = %q", text)
|
||||
}
|
||||
if strings.Contains(text, "<") || strings.Contains(text, "alert") {
|
||||
t.Fatalf("text still has tags/script: %q", text)
|
||||
}
|
||||
if m["content_type"] != "text/html" {
|
||||
t.Fatalf("content_type = %q", m["content_type"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetch_RejectsBinary(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = io.WriteString(w, "\x89PNG")
|
||||
}))
|
||||
defer srv.Close()
|
||||
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
|
||||
_, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
||||
if err == nil || !strings.Contains(err.Error(), "unexpected content type") {
|
||||
t.Fatalf("expected content-type error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetch_SSRFDenied(t *testing.T) {
|
||||
t.Parallel()
|
||||
wf := webFetchToolForTest(t)
|
||||
_, err := wf.Execute(context.Background(), map[string]any{"url": "http://localhost:9/"})
|
||||
if !errors.Is(err, ErrSSRF) {
|
||||
t.Fatalf("expected ErrSSRF, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetch_SizeLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/plain")
|
||||
_, _ = io.WriteString(w, strings.Repeat("z", 200))
|
||||
}))
|
||||
defer srv.Close()
|
||||
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
|
||||
got, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL, "max_bytes": 10})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if !m["truncated"].(bool) {
|
||||
t.Fatal("expected truncated=true")
|
||||
}
|
||||
if len(m["text"].(string)) != 10 {
|
||||
t.Fatalf("text len = %d", len(m["text"].(string)))
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetch_Non200Status(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "bad", http.StatusNotFound)
|
||||
}))
|
||||
defer srv.Close()
|
||||
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||
_, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
||||
if err == nil || !strings.Contains(err.Error(), "status 404") {
|
||||
t.Fatalf("expected status-404 error, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||
)
|
||||
|
||||
// Compile-time assertion that webSearchTool implements sdk.Tool.
|
||||
var _ sdk.Tool = (*webSearchTool)(nil)
|
||||
|
||||
// webSearchTool queries a configured SearXNG-compatible (or custom) search
|
||||
// endpoint and returns ranked results, capped in size and count.
|
||||
type webSearchTool struct {
|
||||
client *http.Client
|
||||
endpoint string
|
||||
allowedHosts []string
|
||||
maxBytes int64
|
||||
maxResults int
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
func newWebSearchTool(t *tools) *webSearchTool {
|
||||
return &webSearchTool{
|
||||
client: t.client,
|
||||
endpoint: t.searchURL,
|
||||
allowedHosts: t.allowedHosts,
|
||||
maxBytes: t.maxHTTPBytes,
|
||||
maxResults: t.maxResults,
|
||||
log: t.log,
|
||||
}
|
||||
}
|
||||
|
||||
// Name implements sdk.Tool.
|
||||
func (s *webSearchTool) Name() string { return "web-search" }
|
||||
|
||||
// Description implements sdk.Tool.
|
||||
func (s *webSearchTool) Description() string {
|
||||
return "Search the web via a configured SearXNG or custom search endpoint and return ranked results."
|
||||
}
|
||||
|
||||
// Parameters implements sdk.Tool.
|
||||
func (s *webSearchTool) Parameters() (map[string]any, error) {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Search query."},
|
||||
"count": map[string]any{"type": "integer", "description": "Max results to return."},
|
||||
},
|
||||
"required": []any{"query"},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Execute implements sdk.Tool.
|
||||
func (s *webSearchTool) Execute(ctx context.Context, args any) (any, error) {
|
||||
m, ok := args.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("io: web-search: args must be an object")
|
||||
}
|
||||
query := argString(m, "query")
|
||||
if strings.TrimSpace(query) == "" {
|
||||
return nil, fmt.Errorf("io: web-search: query is required")
|
||||
}
|
||||
if strings.TrimSpace(s.endpoint) == "" {
|
||||
return nil, fmt.Errorf("io: web-search: no search endpoint configured (set a SearXNG or custom URL)")
|
||||
}
|
||||
target, err := buildSearchURL(s.endpoint, query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, err := validateRequestURL(ctx, target, s.allowedHosts); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if s.client == nil {
|
||||
return nil, fmt.Errorf("io: web-search: no underlying client")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: web-search: build request: %w", err)
|
||||
}
|
||||
resp, err := s.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: web-search: do: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("io: web-search: endpoint returned status %d", resp.StatusCode)
|
||||
}
|
||||
data, truncated, err := readLimited(resp.Body, s.maxBytes)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("io: web-search: read: %w", err)
|
||||
}
|
||||
if truncated {
|
||||
return nil, fmt.Errorf("io: web-search: response exceeds the %d-byte limit", s.maxBytes)
|
||||
}
|
||||
results, err := parseSearchResults(data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
n := s.maxResults
|
||||
if v, ok := argInt(m, "count"); ok && v > 0 {
|
||||
n = v
|
||||
}
|
||||
if len(results) > n {
|
||||
results = results[:n]
|
||||
}
|
||||
return map[string]any{
|
||||
"query": query,
|
||||
"count": len(results),
|
||||
"results": results,
|
||||
"truncated": truncated,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildSearchURL expands an endpoint against a query. It supports a "{query}"
|
||||
// placeholder, otherwise appends q= and format=json for a SearXNG JSON API.
|
||||
func buildSearchURL(endpoint, query string) (string, error) {
|
||||
if strings.TrimSpace(endpoint) == "" {
|
||||
return "", fmt.Errorf("io: web-search: no search endpoint configured")
|
||||
}
|
||||
q := url.QueryEscape(query)
|
||||
if strings.Contains(endpoint, "{query}") {
|
||||
return strings.ReplaceAll(endpoint, "{query}", q), nil
|
||||
}
|
||||
sep := "?"
|
||||
if strings.Contains(endpoint, "?") {
|
||||
sep = "&"
|
||||
}
|
||||
return endpoint + sep + "q=" + q + "&format=json", nil
|
||||
}
|
||||
|
||||
// searchResult is one entry of a SearXNG-compatible results payload.
|
||||
type searchResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
Snippet string `json:"snippet"`
|
||||
}
|
||||
|
||||
// parseSearchResults decodes a {"results":[...]} payload, falling back to the
|
||||
// snippet field when content is absent.
|
||||
func parseSearchResults(data []byte) ([]searchResult, error) {
|
||||
var payload struct {
|
||||
Results []searchResult `json:"results"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &payload); err != nil {
|
||||
return nil, fmt.Errorf("io: web-search: decode results: %w", err)
|
||||
}
|
||||
for i := range payload.Results {
|
||||
if payload.Results[i].Content == "" {
|
||||
payload.Results[i].Content = payload.Results[i].Snippet
|
||||
}
|
||||
}
|
||||
return payload.Results, nil
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// webSearchToolForTest builds a tools set with opts and returns the web-search
|
||||
// tool.
|
||||
func webSearchToolForTest(t *testing.T, opts ...Option) *webSearchTool {
|
||||
t.Helper()
|
||||
return findTool(t, buildTools(t, &config.Config{}, opts...), "web-search").(*webSearchTool)
|
||||
}
|
||||
|
||||
func TestWebSearch_NoEndpoint(t *testing.T) {
|
||||
t.Parallel()
|
||||
ws := webSearchToolForTest(t) // no endpoint configured
|
||||
_, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||
if err == nil || !strings.Contains(err.Error(), "no search endpoint") {
|
||||
t.Fatalf("expected no-endpoint error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSearch_Query(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Query().Get("q") != "cyrene" {
|
||||
t.Errorf("query param q = %q", r.URL.Query().Get("q"))
|
||||
}
|
||||
if r.URL.Query().Get("format") != "json" {
|
||||
t.Errorf("format = %q", r.URL.Query().Get("format"))
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"results":[
|
||||
{"title":"A","url":"http://a.example","content":"alpha"},
|
||||
{"title":"B","url":"http://b.example","snippet":"beta"}
|
||||
]}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
ws := webSearchToolForTest(t,
|
||||
WithSearchEndpoint(srv.URL),
|
||||
WithAllowedHosts(serverHost(srv)),
|
||||
WithHTTPClient(srv.Client()),
|
||||
)
|
||||
|
||||
got, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene", "count": 2})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
m := got.(map[string]any)
|
||||
if m["count"].(int) != 2 {
|
||||
t.Fatalf("count = %v", m["count"])
|
||||
}
|
||||
results := m["results"].([]searchResult)
|
||||
if results[0].Title != "A" || results[0].Content != "alpha" {
|
||||
t.Fatalf("result[0] = %+v", results[0])
|
||||
}
|
||||
if results[1].Content != "beta" {
|
||||
t.Fatalf("result[1] snippet fallback = %+v", results[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSearch_SizeLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
big := strings.Repeat("a", 64)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, big)
|
||||
}))
|
||||
defer srv.Close()
|
||||
ws := webSearchToolForTest(t,
|
||||
WithSearchEndpoint(srv.URL),
|
||||
WithAllowedHosts(serverHost(srv)),
|
||||
WithHTTPClient(srv.Client()),
|
||||
WithHTTPResponseLimit(16),
|
||||
)
|
||||
_, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||
if err == nil || !strings.Contains(err.Error(), "exceeds the 16-byte limit") {
|
||||
t.Fatalf("expected size-limit error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSearch_CountCapped(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"results":[{"title":"A","url":"http://a","content":"x"}]}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
ws := webSearchToolForTest(t,
|
||||
WithSearchEndpoint(srv.URL),
|
||||
WithAllowedHosts(serverHost(srv)),
|
||||
WithHTTPClient(srv.Client()),
|
||||
WithMaxResults(0), // ignored; default 10
|
||||
WithMaxResults(1),
|
||||
)
|
||||
got, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
if m := got.(map[string]any); m["count"].(int) != 1 {
|
||||
t.Fatalf("count = %v", m["count"])
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user