feat(tools/io): 内置 IO/网络工具（file(路径安全)/http(SSRF防护)/web-search/web-fetch/knowledge）+ 大小截断
This commit is contained in:
Aska committed
1 parent
488cf9c79c
commit
3ce9489f8e
17 files changed
+2223
No files matched your search
@@ -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
|
||||
},
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user