109 lines
3.5 KiB
Go
109 lines
3.5 KiB
Go
package io
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
|
)
|
|
|
|
// webFetchToolForTest builds a tools set with opts and returns the web-fetch
|
|
// tool.
|
|
func webFetchToolForTest(t *testing.T, opts ...Option) *webFetchTool {
|
|
t.Helper()
|
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "web-fetch").(*webFetchTool)
|
|
}
|
|
|
|
func TestWebFetch_TextExtraction(t *testing.T) {
|
|
t.Parallel()
|
|
doc := `<!doctype html><html><head><title>T</title>` +
|
|
`<style>.x{color:red}</style><script>alert(1)</script></head>` +
|
|
`<body><h1>Hello</h1><p>World & more</p><!-- hidden --></body></html>`
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
|
_, _ = io.WriteString(w, doc)
|
|
}))
|
|
defer srv.Close()
|
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
|
|
got, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
m := got.(map[string]any)
|
|
text := m["text"].(string)
|
|
if !strings.Contains(text, "Hello") || !strings.Contains(text, "World & more") {
|
|
t.Fatalf("text = %q", text)
|
|
}
|
|
if strings.Contains(text, "<") || strings.Contains(text, "alert") {
|
|
t.Fatalf("text still has tags/script: %q", text)
|
|
}
|
|
if m["content_type"] != "text/html" {
|
|
t.Fatalf("content_type = %q", m["content_type"])
|
|
}
|
|
}
|
|
|
|
func TestWebFetch_RejectsBinary(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = io.WriteString(w, "\x89PNG")
|
|
}))
|
|
defer srv.Close()
|
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
|
|
_, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
|
if err == nil || !strings.Contains(err.Error(), "unexpected content type") {
|
|
t.Fatalf("expected content-type error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestWebFetch_SSRFDenied(t *testing.T) {
|
|
t.Parallel()
|
|
wf := webFetchToolForTest(t)
|
|
_, err := wf.Execute(context.Background(), map[string]any{"url": "http://localhost:9/"})
|
|
if !errors.Is(err, ErrSSRF) {
|
|
t.Fatalf("expected ErrSSRF, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestWebFetch_SizeLimit(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
_, _ = io.WriteString(w, strings.Repeat("z", 200))
|
|
}))
|
|
defer srv.Close()
|
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
|
|
got, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL, "max_bytes": 10})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
m := got.(map[string]any)
|
|
if !m["truncated"].(bool) {
|
|
t.Fatal("expected truncated=true")
|
|
}
|
|
if len(m["text"].(string)) != 10 {
|
|
t.Fatalf("text len = %d", len(m["text"].(string)))
|
|
}
|
|
}
|
|
|
|
func TestWebFetch_Non200Status(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "bad", http.StatusNotFound)
|
|
}))
|
|
defer srv.Close()
|
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
_, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
|
if err == nil || !strings.Contains(err.Error(), "status 404") {
|
|
t.Fatalf("expected status-404 error, got %v", err)
|
|
}
|
|
}
|