Agent-Native/runtime/tools/fetch.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"html"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"time"
)

// FetchConfig 声明 fetch_url 的运行条件；未注册（组合根未提供）则工具不存在。
type FetchConfig struct {
	// Client 可选注入（测试替身）；缺省构造 20s 超时客户端。
	// Transport 为 nil 或 *http.Transport 时，Clone 后替换 DialContext：
	// 拨号时再次解析，只连接通过 SSRF 校验的 IP。
	// 自定义 RoundTripper 无法替换拨号，只做请求前预校验；实现不得在
	// RoundTrip 里重新解析到未校验的地址。
	Client *http.Client
	// MaxBytes 响应体上限（默认 2 MiB）。
	MaxBytes int
	// AllowedHosts 非空时仅允许这些主机（域名或 host:port）；空 = 公网放开，
	// 但私网/回环/链路本地地址在任何模式下都被拒绝（SSRF 硬边界）。
	AllowedHosts []string
	// Resolve 可选注入主机解析（测试替身）；缺省 net.LookupIP。
	Resolve func(host string) ([]net.IP, error)
}

const (
	fetchTimeout  = 20 * time.Second
	fetchMaxBytes = 2 << 20
)

func (c *FetchConfig) resolver() func(string) ([]net.IP, error) {
	if c.Resolve != nil {
		return c.Resolve
	}
	return net.LookupIP
}

// RegisterFetch 注册 fetch_url 工具（cfg 必须非 nil，由组合根显式提供）。
func RegisterFetch(r *Registry, cfg FetchConfig) {
	r.Register(&fetchTool{cfg: cfg})
}

type fetchTool struct {
	cfg FetchConfig
}

func (t *fetchTool) Name() string { return "fetch_url" }
func (t *fetchTool) Description() string {
	return "Fetch an http(s) URL and return readable text (HTML stripped, up to 2 MiB, 20s timeout). Private/loopback network targets are always rejected."
}
func (t *fetchTool) Parameters() json.RawMessage {
	return json.RawMessage(`{"type":"object","required":["url"],"properties":{"url":{"type":"string"}}}`)
}

func (t *fetchTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
	var obj struct {
		URL string `json:"url"`
	}
	if err := json.Unmarshal(args, &obj); err != nil {
		return Result{}, err
	}
	raw := strings.TrimSpace(obj.URL)
	if raw == "" {
		return Result{Content: "url is empty", IsError: true}, nil
	}
	client := t.cfg.Client
	if client == nil {
		client = &http.Client{Timeout: fetchTimeout}
	}
	allowed := make(map[string]bool, len(t.cfg.AllowedHosts))
	for _, h := range t.cfg.AllowedHosts {
		allowed[strings.ToLower(h)] = true
	}
	maxBytes := t.cfg.MaxBytes
	if maxBytes <= 0 {
		maxBytes = fetchMaxBytes
	}

	reqCtx, cancel := context.WithTimeout(ctx, fetchTimeout)
	defer cancel()
	req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, raw, nil)
	if err != nil {
		return Result{Content: err.Error(), IsError: true}, nil
	}
	req.Header.Set("User-Agent", "agent-native-fetch/1.0")
	var finalURL string
	resolve := t.cfg.resolver()
	hostRedirect := client.CheckRedirect
	transport, _ := pinFetchTransport(client.Transport, resolve)
	client2 := &http.Client{
		Transport: transport,
		Jar:       client.Jar,
		Timeout:   client.Timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= 5 {
				return fmt.Errorf("stopped after 5 redirects")
			}
			if err := validateFetchTarget(req.URL, allowed, resolve); err != nil {
				return err
			}
			if hostRedirect != nil {
				return hostRedirect(req, via)
			}
			return nil
		},
	}
	if err := validateFetchTarget(req.URL, allowed, resolve); err != nil {
		return Result{Content: err.Error(), IsError: true}, nil
	}
	resp, err := client2.Do(req)
	if err != nil {
		return Result{Content: err.Error(), IsError: true}, nil
	}
	defer resp.Body.Close()
	finalURL = resp.Request.URL.String()
	body, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxBytes)+1))
	if err != nil {
		return Result{Content: err.Error(), IsError: true}, nil
	}
	if len(body) > maxBytes {
		body = body[:maxBytes]
	}
	ct := resp.Header.Get("Content-Type")
	text, ok := fetchReadableText(ct, string(body))
	if !ok {
		return Result{Content: fmt.Sprintf("unsupported content type %q at %s (only text/html/json/xml/yaml/csv are extracted)", ct, finalURL), IsError: true}, nil
	}
	return Result{Content: "url: " + finalURL + "\ncontent-type: " + ct + "\n\n" + text}, nil
}

// pinFetchTransport 在默认或 *http.Transport 上安装「解析并拨号已校验 IP」的 DialContext。
// 自定义 RoundTripper 原样返回，调用方保留预校验。
func pinFetchTransport(base http.RoundTripper, resolve func(string) ([]net.IP, error)) (http.RoundTripper, bool) {
	var tr *http.Transport
	switch t := base.(type) {
	case nil:
		tr = http.DefaultTransport.(*http.Transport).Clone()
	case *http.Transport:
		tr = t.Clone()
	default:
		return base, false
	}
	// 不走环境变量里的代理：代理拨号会绕开「只连接已校验 IP」。
	tr.Proxy = nil
	dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
	tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			return nil, err
		}
		ips, err := resolve(host)
		if err != nil {
			return nil, fmt.Errorf("resolve %q: %w", host, err)
		}
		var chosen net.IP
		for _, ip := range ips {
			if ipIsPrivate(ip) {
				return nil, fmt.Errorf("target %q resolves to private/loopback address %s (SSRF guard)", host, ip)
			}
			if chosen == nil {
				chosen = ip
			}
		}
		if chosen == nil {
			return nil, fmt.Errorf("resolve %q: no addresses", host)
		}
		return dialer.DialContext(ctx, network, net.JoinHostPort(chosen.String(), port))
	}
	return tr, true
}

// validateFetchTarget 校验 scheme、主机策略与 SSRF 硬边界（私网/回环/链路本地）。
func validateFetchTarget(u *url.URL, allowed map[string]bool, resolve func(string) ([]net.IP, error)) error {
	if u.Scheme != "http" && u.Scheme != "https" {
		return fmt.Errorf("scheme %q not allowed (http/https only)", u.Scheme)
	}
	host := u.Hostname()
	if host == "" {
		return fmt.Errorf("empty host")
	}
	if len(allowed) > 0 && !allowed[strings.ToLower(host)] {
		return fmt.Errorf("host %q not in allowed hosts", host)
	}
	// 解析全部地址：任一命中私网段即拒绝（域名可能同时解析到公私网地址）。
	addrs, err := resolve(host)
	if err != nil {
		return fmt.Errorf("resolve %q: %w", host, err)
	}
	for _, ip := range addrs {
		if ipIsPrivate(ip) {
			return fmt.Errorf("target %q resolves to private/loopback address %s (SSRF guard)", host, ip)
		}
	}
	return nil
}

var fetchPrivateNets = mustParseCIDRs(
	"0.0.0.0/8", "10.0.0.0/8", "127.0.0.0/8", "169.254.0.0/16",
	"172.16.0.0/12", "192.168.0.0/16", "100.64.0.0/10",
	"::1/128", "fc00::/7", "fe80::/10",
)

func mustParseCIDRs(specs ...string) []*net.IPNet {
	var out []*net.IPNet
	for _, s := range specs {
		_, n, err := net.ParseCIDR(s)
		if err != nil {
			panic("fetch: bad cidr " + s)
		}
		out = append(out, n)
	}
	return out
}

func ipIsPrivate(ip net.IP) bool {
	// IPv4（含 Go 内部的 ::ffff:a.b.c.d 映射形式）先归一化成 4 字节，
	// 否则 ::ffff:0:0/96 这类映射网段会误吞所有 IPv4 地址。
	if v4 := ip.To4(); v4 != nil {
		ip = v4
	}
	if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() {
		return true
	}
	for _, n := range fetchPrivateNets {
		if n.Contains(ip) {
			return true
		}
	}
	return false
}

// fetchReadableText 按内容类型抽取可读文本；HTML 剥脚本样式与标签。
func fetchReadableText(contentType, body string) (string, bool) {
	ct := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0]))
	if ct == "" {
		ct = "text/plain" // 无声明按纯文本处理
	}
	switch {
	case ct == "text/html":
		return htmlToText(body), true
	case strings.HasPrefix(ct, "text/"),
		ct == "application/json",
		ct == "application/xml",
		strings.HasSuffix(ct, "+json"),
		strings.HasSuffix(ct, "+xml"),
		ct == "application/yaml" || ct == "text/yaml",
		ct == "application/csv" || ct == "text/csv":
		return body, true
	default:
		return "", false
	}
}

var (
	fetchScriptRe = regexp.MustCompile(`(?is)<(script|style|noscript)[^>]*>.*?</(script|style|noscript)>`)
	fetchTagRe    = regexp.MustCompile(`(?s)<[^>]*>`)
	fetchSpaceRe  = regexp.MustCompile(`\n{3,}`)
)

func htmlToText(body string) string {
	s := fetchScriptRe.ReplaceAllString(body, "")
	s = fetchTagRe.ReplaceAllString(s, "\n")
	s = html.UnescapeString(s)
	s = strings.ReplaceAll(s, "\r", "")
	s = fetchSpaceRe.ReplaceAllString(s, "\n\n")
	return strings.TrimSpace(s)
}