1package tools
2
3import (
4 "context"
5 "crypto/tls"
6 "crypto/x509"
7 "encoding/json"
8 "net"
9 "net/http"
10 "net/http/httptest"
11 "net/netip"
12 "strings"
13 "sync"
14 "sync/atomic"
15 "testing"
16
17 "github.com/bnema/zut/packages/provider"
18)
19
20func testWebFetcher(t *testing.T, server *httptest.Server) *webFetcher {
21 t.Helper()
22 address := server.Listener.Addr().String()
23 return &webFetcher{
24 resolve: func(context.Context, string) ([]netip.Addr, error) {
25 return []netip.Addr{netip.MustParseAddr("8.8.8.8")}, nil
26 },
27 dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
28 return (&net.Dialer{}).DialContext(ctx, network, address)
29 },
30 }
31}
32
33func TestWebPageToolsNavigateOnlyStoredReferences(t *testing.T) {
34 var requests atomic.Int64
35 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
36 requests.Add(1)
37 w.Header().Set("Content-Type", "text/html; charset=utf-8")
38 if r.URL.Path == "/next" {
39 _, _ = w.Write([]byte("<html><body><p>Second page</p></body></html>"))
40 return
41 }
42 _, _ = w.Write([]byte(`<html><head><title>Example</title></head><body><script>ignore</script><main><h1>Heading</h1><p>Useful text</p><a href="/next">Next page</a></main></body></html>`))
43 }))
44 defer server.Close()
45
46 store := NewBrowsingStore()
47 ref := store.addSource("http://public.test/")
48 fetcher := testWebFetcher(t, server)
49 open := &WebOpenTool{store: store, fetcher: fetcher}
50 find := &WebFindTool{store: store}
51 click := &WebClickTool{store: store, fetcher: fetcher}
52
53 result, err := open.Execute(context.Background(), json.RawMessage(`{"ref_id":"`+ref+`"}`), nil)
54 if err != nil || result.IsError {
55 t.Fatalf("open result = %#v, err = %v", result, err)
56 }
57 text := result.Content[0].(provider.TextBlock).Text
58 if !strings.Contains(text, "Page reference: web-") || !strings.Contains(text, "Useful text") || strings.Contains(text, "ignore") || !strings.Contains(text, "[1] Next page") {
59 t.Fatalf("unexpected open output:\n%s", text)
60 }
61 pageRef := strings.Split(strings.Split(text, "Page reference: ")[1], "\n")[0]
62
63 result, err = find.Execute(context.Background(), json.RawMessage(`{"ref_id":"`+pageRef+`","pattern":"useful"}`), nil)
64 if err != nil || result.IsError || !strings.Contains(result.Content[0].(provider.TextBlock).Text, "Useful text") {
65 t.Fatalf("find result = %#v, err = %v", result, err)
66 }
67 if got := requests.Load(); got != 1 {
68 t.Fatalf("find made %d requests, want 1 total", got)
69 }
70
71 result, err = click.Execute(context.Background(), json.RawMessage(`{"ref_id":"`+pageRef+`","id":1}`), nil)
72 if err != nil || result.IsError || !strings.Contains(result.Content[0].(provider.TextBlock).Text, "Second page") {
73 t.Fatalf("click result = %#v, err = %v", result, err)
74 }
75 if got := requests.Load(); got != 2 {
76 t.Fatalf("requests = %d, want 2", got)
77 }
78}
79
80func TestWebPageSearchReturnsStoredSourceReference(t *testing.T) {
81 store := NewBrowsingStore()
82 tool := NewWebSearchToolWithStore(store)
83 tool.transport = webSearchRoundTripper(func(request *http.Request) (*http.Response, error) {
84 return webSearchHTMLResponse(request, webSearchFixture), nil
85 })
86 result, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"reference","max_results":1}`), nil)
87 if err != nil || result.IsError {
88 t.Fatalf("search result = %#v, err = %v", result, err)
89 }
90 text := result.Content[0].(provider.TextBlock).Text
91 if !strings.Contains(text, "(ref: web-") {
92 t.Fatalf("search did not expose a source reference: %s", text)
93 }
94 ref := strings.Split(strings.Split(text, "(ref: ")[1], ")")[0]
95 doc, ok := store.get(ref, false)
96 if !ok || !doc.source || doc.url != "https://example.com/one" {
97 t.Fatalf("stored source = %#v, found = %v", doc, ok)
98 }
99}
100
101func TestWebSearchDoesNotCommitSourcesAfterStoreRevocation(t *testing.T) {
102 store := NewBrowsingStore()
103 tool := NewWebSearchToolWithStore(store)
104 tool.transport = webSearchRoundTripper(func(request *http.Request) (*http.Response, error) {
105 store.Clear()
106 return webSearchHTMLResponse(request, webSearchFixture), nil
107 })
108 result, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"revoked"}`), nil)
109 if err != nil || !result.IsError || result.Content[0].(provider.TextBlock).Text != "web search: unavailable in this session" {
110 t.Fatalf("revoked search result = %#v, err = %v", result, err)
111 }
112 if _, ok := store.get("web-1", false); ok {
113 t.Fatal("revoked search retained a source reference")
114 }
115}
116
117func TestWebPageFetcherRedirectsWithFreshValidation(t *testing.T) {
118 var resolutions atomic.Int64
119 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
120 if r.URL.Path == "/redirect" {
121 http.Redirect(w, r, "/final", http.StatusFound)
122 return
123 }
124 w.Header().Set("Content-Type", "text/plain")
125 _, _ = w.Write([]byte("complete"))
126 }))
127 defer server.Close()
128 fetcher := testWebFetcher(t, server)
129 fetcher.resolve = func(context.Context, string) ([]netip.Addr, error) {
130 resolutions.Add(1)
131 return []netip.Addr{netip.MustParseAddr("8.8.8.8")}, nil
132 }
133 body, finalURL, media, message := fetcher.fetch(context.Background(), "http://public.test/redirect")
134 if message != "" || string(body) != "complete" || finalURL != "http://public.test/final" || media != "text/plain" {
135 t.Fatalf("redirect result body=%q url=%q media=%q message=%q", body, finalURL, media, message)
136 }
137 if resolutions.Load() != 2 {
138 t.Fatalf("resolutions = %d, want 2", resolutions.Load())
139 }
140}
141
142func TestWebPageFetcherPreservesHTTPSHostnameForSNI(t *testing.T) {
143 var serverName atomic.Value
144 server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
145 serverName.Store(r.TLS.ServerName)
146 w.Header().Set("Content-Type", "text/plain")
147 _, _ = w.Write([]byte("secure"))
148 }))
149 server.StartTLS()
150 defer server.Close()
151 fetcher := testWebFetcher(t, server)
152 roots := x509.NewCertPool()
153 roots.AddCert(server.Certificate())
154 fetcher.tlsConfig = &tls.Config{RootCAs: roots}
155 body, _, _, message := fetcher.fetch(context.Background(), "https://example.com/")
156 got, _ := serverName.Load().(string)
157 if message != "" || string(body) != "secure" || got != "example.com" {
158 t.Fatalf("HTTPS body=%q message=%q SNI=%q", body, message, got)
159 }
160}
161
162func TestWebPageFetcherPinsTheValidatedAddress(t *testing.T) {
163 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
164 w.Header().Set("Content-Type", "text/plain")
165 _, _ = w.Write([]byte("pinned"))
166 }))
167 defer server.Close()
168 var dialed atomic.Value
169 fetcher := testWebFetcher(t, server)
170 originalDial := fetcher.dial
171 fetcher.dial = func(ctx context.Context, network, address string) (net.Conn, error) {
172 dialed.Store(address)
173 return originalDial(ctx, network, address)
174 }
175 body, _, _, message := fetcher.fetch(context.Background(), "http://public.test/")
176 got, _ := dialed.Load().(string)
177 if message != "" || string(body) != "pinned" || got != "8.8.8.8:80" {
178 t.Fatalf("pinned fetch body=%q message=%q dialed=%q", body, message, got)
179 }
180}
181
182func TestWebPageFetcherRejectsPrivateAndMixedDestinationsBeforeDial(t *testing.T) {
183 var dials atomic.Int64
184 fetcher := &webFetcher{
185 resolve: func(context.Context, string) ([]netip.Addr, error) {
186 return []netip.Addr{netip.MustParseAddr("8.8.8.8"), netip.MustParseAddr("127.0.0.1")}, nil
187 },
188 dial: func(context.Context, string, string) (net.Conn, error) { dials.Add(1); return nil, context.Canceled },
189 }
190 _, _, _, message := fetcher.fetch(context.Background(), "http://public.test/")
191 if message != webPageDenied || dials.Load() != 0 {
192 t.Fatalf("mixed destination message=%q dials=%d", message, dials.Load())
193 }
194 _, _, _, message = fetcher.fetch(context.Background(), "http://127.0.0.1/")
195 if message != webPageDenied || dials.Load() != 0 {
196 t.Fatalf("literal private destination message=%q dials=%d", message, dials.Load())
197 }
198}
199
200func TestPublicWebIPRejectsSpecialIPv6Ranges(t *testing.T) {
201 for _, raw := range []string{"2001:db8::1", "2001:2::1", "64:ff9b:1::1", "fc00::1"} {
202 if publicWebIP(netip.MustParseAddr(raw)) {
203 t.Fatalf("special IPv6 address %s was allowed", raw)
204 }
205 }
206 if !publicWebIP(netip.MustParseAddr("2606:4700:4700::1111")) {
207 t.Fatal("public IPv6 address was rejected")
208 }
209}
210
211func TestWebPageParserBudgetAndCancellation(t *testing.T) {
212 deep := []byte("<html><body>" + strings.Repeat("<div>", webPageMaxDepth+1) + "text" + strings.Repeat("</div>", webPageMaxDepth+1) + "</body></html>")
213 if _, message := parseWebPage(deep, "http://public.test/", "text/html"); message != webPageUnparseable {
214 t.Fatalf("deep HTML message = %q, want %q", message, webPageUnparseable)
215 }
216 fetcher := &webFetcher{
217 resolve: func(context.Context, string) ([]netip.Addr, error) {
218 return []netip.Addr{netip.MustParseAddr("8.8.8.8")}, nil
219 },
220 dial: func(context.Context, string, string) (net.Conn, error) { return nil, context.Canceled },
221 }
222 ctx, cancel := context.WithCancel(context.Background())
223 cancel()
224 _, _, _, message := fetcher.fetch(ctx, "http://public.test/")
225 if message != webPageCancelled {
226 t.Fatalf("cancelled fetch message = %q, want %q", message, webPageCancelled)
227 }
228 fetcher.resolve = func(ctx context.Context, _ string) ([]netip.Addr, error) {
229 <-ctx.Done()
230 return nil, ctx.Err()
231 }
232 _, _, _, message = fetcher.fetch(ctx, "http://public.test/")
233 if message != webPageCancelled {
234 t.Fatalf("cancelled resolver message = %q, want %q", message, webPageCancelled)
235 }
236}
237
238func TestBrowsingStoreConcurrentAccess(t *testing.T) {
239 store := NewBrowsingStore()
240 var wait sync.WaitGroup
241 for worker := 0; worker < 16; worker++ {
242 wait.Add(1)
243 go func() {
244 defer wait.Done()
245 for iteration := 0; iteration < 100; iteration++ {
246 id := store.addSource("http://public.test/")
247 store.get(id, false)
248 if iteration%25 == 0 {
249 store.Clear()
250 }
251 }
252 }()
253 }
254 wait.Wait()
255}
256
257func TestWebOpenDirectURLUsesSameFetcherPolicy(t *testing.T) {
258 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
259 w.Header().Set("Content-Type", "text/html; charset=utf-8")
260 _, _ = w.Write([]byte(`<html><head><title>Direct</title></head><body><main><p>Direct page</p></main></body></html>`))
261 }))
262 defer server.Close()
263
264 store := NewBrowsingStore()
265 open := &WebOpenTool{store: store, fetcher: testWebFetcher(t, server)}
266
267 result, err := open.Execute(context.Background(), json.RawMessage(`{"url":"http://public.test/direct"}`), nil)
268 if err != nil || result.IsError {
269 t.Fatalf("direct open result = %#v, err = %v", result, err)
270 }
271 text := result.Content[0].(provider.TextBlock).Text
272 if !strings.Contains(text, "Direct page") || !strings.Contains(text, "Page reference: web-") {
273 t.Fatalf("unexpected direct open output:\n%s", text)
274 }
275
276 preview, err := open.Preview(context.Background(), json.RawMessage(`{"url":"http://public.test/direct"}`))
277 if err != nil || preview.IsError || !strings.Contains(preview.Content[0].(provider.TextBlock).Text, "No request has been made.") {
278 t.Fatalf("direct preview = %#v, err = %v", preview, err)
279 }
280
281 for _, raw := range []string{
282 `{}`,
283 `{"ref_id":"web-1","url":"http://public.test/"}`,
284 `{"url":"ftp://public.test/"}`,
285 `{"url":"http://public.test/#frag"}`,
286 `{"url":"not-a-url"}`,
287 `{"url":"http://public.test/","line":0}`,
288 } {
289 result, err := open.Execute(context.Background(), json.RawMessage(raw), nil)
290 if err != nil || !result.IsError || result.Content[0].(provider.TextBlock).Text != webPageInvalid {
291 t.Fatalf("direct open %s = %#v, err = %v", raw, result, err)
292 }
293 }
294}
295
296func TestWebPageStoreGenerationAndSchemaFailures(t *testing.T) {
297 store := NewBrowsingStore()
298 retained := store.addSource("http://public.test/")
299 generation := store.snapshotGeneration()
300 if id, committed := store.addPage(&webDocument{url: "http://public.test", lines: []string{strings.Repeat("x", webPageMaxStoreBytes+1)}}, generation); id != "" || committed {
301 t.Fatalf("oversized page commit = %q, %v", id, committed)
302 }
303 if _, ok := store.get(retained, false); !ok {
304 t.Fatal("oversized page evicted an existing document")
305 }
306 store.Clear()
307 if id, committed := store.addPage(&webDocument{url: "http://public.test", lines: []string{"text"}}, generation); id != "" || committed {
308 t.Fatalf("stale page commit = %q, %v", id, committed)
309 }
310 find := &WebFindTool{store: store}
311 result, err := find.Execute(context.Background(), json.RawMessage(`{"ref_id":"https://not-a-handle","pattern":"x"}`), nil)
312 if err != nil || !result.IsError || result.Content[0].(provider.TextBlock).Text != webPageInvalid {
313 t.Fatalf("arbitrary URL result = %#v, err = %v", result, err)
314 }
315}