Agent-Native/runtime/tools/fetch_test.go

package tools

import (
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// fakeResolve 固定域名 → 地址（无 DNS 依赖的封闭测试）。
func fakeResolve(mapping map[string][]net.IP) func(string) ([]net.IP, error) {
	return func(host string) ([]net.IP, error) {
		if ips, ok := mapping[host]; ok {
			return ips, nil
		}
		// IP 字面量直接解析。
		if ip := net.ParseIP(host); ip != nil {
			return []net.IP{ip}, nil
		}
		return nil, &net.DNSError{Err: "no such host", Name: host}
	}
}

func fetchReg(cfg FetchConfig) *Registry {
	reg := NewRegistry("/")
	RegisterFetch(reg, cfg)
	return reg
}

func TestFetchReadableTextUnit(t *testing.T) {
	out, ok := fetchReadableText("text/html; charset=utf-8", "<html><script>bad()</script><h1>Hi</h1><p>A &amp; B</p></html>")
	if !ok || !strings.Contains(out, "Hi") || !strings.Contains(out, "A & B") || strings.Contains(out, "bad()") {
		t.Fatalf("html strip wrong: %q ok=%v", out, ok)
	}
	out, ok = fetchReadableText("application/json", `{"x":1}`)
	if !ok || out != `{"x":1}` {
		t.Fatalf("json passthrough wrong: %q", out)
	}
	_, ok = fetchReadableText("image/png", "binary")
	if ok {
		t.Fatal("binary content type must be rejected")
	}
	out, ok = fetchReadableText("", "no content type")
	if !ok || out != "no content type" {
		t.Fatalf("empty content type should default to text, got %q %v", out, ok)
	}
}

func TestFetchURLSSRDGuards(t *testing.T) {
	reg := fetchReg(FetchConfig{Resolve: fakeResolve(nil)})

	for _, target := range []string{
		`{"url":"http://127.0.0.1/x"}`,
		`{"url":"http://10.0.0.5/x"}`,
		`{"url":"http://192.168.1.1/x"}`,
		`{"url":"http://169.254.169.254/latest/meta-data"}`,
		`{"url":"http://[::1]/x"}`,
		`{"url":"file:///etc/passwd"}`,
	} {
		res := runTool(t, reg, "fetch_url", target)
		if !res.IsError {
			t.Fatalf("SSRF/scheme guard must reject %s, got %+v", target, res)
		}
	}

	// 域名解析到私网同样拒绝。
	reg2 := fetchReg(FetchConfig{Resolve: fakeResolve(map[string][]net.IP{
		"internal.example": {net.ParseIP("192.168.0.10")},
	})})
	res := runTool(t, reg2, "fetch_url", `{"url":"http://internal.example/x"}`)
	if !res.IsError || !strings.Contains(res.Content, "SSRF guard") {
		t.Fatalf("private resolution must be rejected, got %+v", res)
	}
}

func TestFetchURLAllowedHostsPolicy(t *testing.T) {
	reg := fetchReg(FetchConfig{
		AllowedHosts: []string{"api.example.com"},
		Resolve:      fakeResolve(map[string][]net.IP{"api.example.com": {net.ParseIP("93.184.216.34")}}),
	})
	res := runTool(t, reg, "fetch_url", `{"url":"http://other.example.com/x"}`)
	if !res.IsError || !strings.Contains(res.Content, "not in allowed hosts") {
		t.Fatalf("host policy must reject, got %+v", res)
	}
	// 允许的主机：校验必须放行（真实网络下可能成功取回，也可能因出口
	// 受限连接失败——两者都不算策略拦截）。
	res = runTool(t, reg, "fetch_url", `{"url":"http://api.example.com/x"}`)
	if strings.Contains(res.Content, "not in allowed") || strings.Contains(res.Content, "SSRF") {
		t.Fatalf("allowed host should pass validation, got %+v", res)
	}
}

func TestFetchURLEndToEndWithStubTransport(t *testing.T) {
	var handler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		_, _ = w.Write([]byte("fetched body"))
	}
	reg := fetchReg(FetchConfig{
		Client:  &http.Client{Transport: stubTransport{handler: handler}},
		Resolve: fakeResolve(map[string][]net.IP{"ok.example": {net.ParseIP("93.184.216.34")}}),
	})
	res := runTool(t, reg, "fetch_url", `{"url":"http://ok.example/x"}`)
	if res.IsError {
		t.Fatal(res.Content)
	}
	if !strings.Contains(res.Content, "fetched body") || !strings.Contains(res.Content, "url: http://ok.example/x") {
		t.Fatalf("stub fetch wrong: %+v", res)
	}
}

type stubTransport struct{ handler http.HandlerFunc }

func (s stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	rec := httptest.NewRecorder()
	s.handler(rec, req)
	resp := rec.Result()
	resp.Request = req
	return resp, nil
}

func TestFetchURLRedirectToPrivateRejected(t *testing.T) {
	reg := fetchReg(FetchConfig{
		Client: &http.Client{Transport: redirectTransport{to: "http://127.0.0.1/evil"}},
		Resolve: fakeResolve(map[string][]net.IP{
			"public.example": {net.ParseIP("93.184.216.34")},
			"127.0.0.1":      {net.ParseIP("127.0.0.1")},
		}),
	})
	res := runTool(t, reg, "fetch_url", `{"url":"http://public.example/redirect"}`)
	if !res.IsError {
		t.Fatalf("redirect to private target must be rejected, got %+v", res)
	}
}

type redirectTransport struct{ to string }

func (s redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	rec := httptest.NewRecorder()
	rec.Header().Set("Location", s.to)
	rec.Code = http.StatusFound
	resp := rec.Result()
	resp.Request = req
	return resp, nil
}

func TestFetchResolveChangesToLoopbackRejected(t *testing.T) {
	var n int
	reg := fetchReg(FetchConfig{
		Resolve: func(host string) ([]net.IP, error) {
			n++
			if n == 1 {
				return []net.IP{net.ParseIP("93.184.216.34")}, nil
			}
			return []net.IP{net.ParseIP("127.0.0.1")}, nil
		},
	})
	res := runTool(t, reg, "fetch_url", `{"url":"http://flip.example/x"}`)
	if !res.IsError || !strings.Contains(res.Content, "private/loopback") {
		t.Fatalf("second resolve to loopback must be rejected, got %+v", res)
	}
	if n < 2 {
		t.Fatalf("dial should resolve again, got %d calls", n)
	}
}

func TestFetchHostCheckRedirectStillRejectsPrivate(t *testing.T) {
	reg := fetchReg(FetchConfig{
		Client: &http.Client{
			Transport: redirectTransport{to: "http://127.0.0.1/evil"},
			CheckRedirect: func(req *http.Request, via []*http.Request) error {
				return nil
			},
		},
		Resolve: fakeResolve(map[string][]net.IP{
			"public.example": {net.ParseIP("93.184.216.34")},
			"127.0.0.1":      {net.ParseIP("127.0.0.1")},
		}),
	})
	res := runTool(t, reg, "fetch_url", `{"url":"http://public.example/redirect"}`)
	if !res.IsError || !strings.Contains(res.Content, "private/loopback") {
		t.Fatalf("host CheckRedirect must not bypass SSRF guard, got %+v", res)
	}
}

func TestFetchNotRegisteredByDefault(t *testing.T) {
	reg := WorkspaceTools(t.TempDir())
	for _, d := range reg.Decls() {
		if d.Name == "fetch_url" {
			t.Fatal("fetch_url must not register without explicit FetchConfig")
		}
	}
}