127 lines
3.7 KiB
Go
127 lines
3.7 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"regexp"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/crypto/byteclaw/internal/llm"
|
|
)
|
|
|
|
// WebFetch downloads a URL with timeout and size limit. By default it refuses
|
|
// loopback/private/link-local targets (SSRF protection), checked at dial time
|
|
// so DNS rebinding and redirects are covered too. TLS verification is always on.
|
|
type WebFetch struct {
|
|
client *http.Client
|
|
maxBytes int64
|
|
}
|
|
|
|
var ErrBlockedAddr = errors.New("destination address is not allowed")
|
|
|
|
func NewWebFetch(timeout time.Duration, maxBytes int64, allowPrivate bool) *WebFetch {
|
|
d := &net.Dialer{Timeout: 10 * time.Second}
|
|
if !allowPrivate {
|
|
d.Control = func(_, address string, _ syscall.RawConn) error {
|
|
host, _, err := net.SplitHostPort(address)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ip := net.ParseIP(host)
|
|
if ip == nil || blockedIP(ip) {
|
|
return ErrBlockedAddr
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
tr := &http.Transport{DialContext: d.DialContext, TLSHandshakeTimeout: 10 * time.Second, ForceAttemptHTTP2: true}
|
|
return &WebFetch{maxBytes: maxBytes, client: &http.Client{
|
|
Timeout: timeout, Transport: tr,
|
|
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
|
if len(via) >= 5 {
|
|
return errors.New("too many redirects")
|
|
}
|
|
if req.URL.Scheme != "http" && req.URL.Scheme != "https" {
|
|
return errors.New("redirect to unsupported scheme")
|
|
}
|
|
return nil
|
|
},
|
|
}}
|
|
}
|
|
|
|
func blockedIP(ip net.IP) bool {
|
|
return ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() ||
|
|
ip.IsUnspecified() || ip.IsMulticast() || ip.IsInterfaceLocalMulticast() ||
|
|
ip.Equal(net.IPv4bcast) || (ip.To4() != nil && ip.To4()[0] == 100 && ip.To4()[1]&0xC0 == 64) // CGNAT 100.64/10
|
|
}
|
|
|
|
func (w *WebFetch) Spec() llm.ToolSpec {
|
|
return llm.ToolSpec{
|
|
Name: "web_fetch",
|
|
Description: "Fetch a public http(s) URL and return its text content (HTML is stripped, output truncated).",
|
|
Parameters: obj(map[string]any{"url": map[string]any{"type": "string"}}, "url"),
|
|
}
|
|
}
|
|
|
|
var (
|
|
reScript = regexp.MustCompile(`(?is)<(script|style|noscript)[^>]*>.*?</(script|style|noscript)>`)
|
|
reTag = regexp.MustCompile(`(?s)<[^>]+>`)
|
|
reSpace = regexp.MustCompile(`[ \t\r\f\v]+`)
|
|
reNL = regexp.MustCompile(`\n\s*\n+`)
|
|
)
|
|
|
|
func (w *WebFetch) Call(ctx context.Context, raw json.RawMessage) (string, error) {
|
|
var a struct {
|
|
URL string `json:"url"`
|
|
}
|
|
if err := json.Unmarshal(raw, &a); err != nil {
|
|
return "", err
|
|
}
|
|
u, err := url.Parse(a.URL)
|
|
if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
|
|
return "", fmt.Errorf("only absolute http(s) URLs are allowed")
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
req.Header.Set("User-Agent", "ByteClaw/0.1 (+self-hosted assistant)")
|
|
resp, err := w.client.Do(req)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(io.LimitReader(resp.Body, w.maxBytes+1))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
truncated := int64(len(body)) > w.maxBytes
|
|
if truncated {
|
|
body = body[:w.maxBytes]
|
|
}
|
|
text := string(body)
|
|
if strings.Contains(resp.Header.Get("Content-Type"), "html") {
|
|
text = reScript.ReplaceAllString(text, "")
|
|
text = reTag.ReplaceAllString(text, " ")
|
|
text = reSpace.ReplaceAllString(text, " ")
|
|
text = reNL.ReplaceAllString(text, "\n\n")
|
|
}
|
|
text = strings.TrimSpace(text)
|
|
const maxOut = 20000
|
|
if len(text) > maxOut {
|
|
text, truncated = text[:maxOut], true
|
|
}
|
|
out := fmt.Sprintf("HTTP %d %s\n\n%s", resp.StatusCode, resp.Header.Get("Content-Type"), text)
|
|
if truncated {
|
|
out += "\n\n[truncated]"
|
|
}
|
|
return out, nil
|
|
}
|