zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-09-08 02:08:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): bound crawler response bodies to 4 MiB (#18677)

Browse Source
This commit is contained in:
Lem0nTea2002
2026-08-25 13:59:13 +08:00
committed by GitHub
parent 24abfcd43d
commit bb488db9d4
2 changed files with 39 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
internal/agent/tool/crawler.go
View File

@@ -35,6 +35,8 @@ const crawlerToolName = "web_crawler"
const crawlerToolDescription = "Crawls a web page and returns its extracted text content and links."
const maxCrawlerResponseBytes = 4 << 20
// crawlerArgs is the JSON shape the model sends in. query is the
// Python-compatible argument name; url is accepted for older Go
// callers. max_depth and max_pages are accepted for API symmetry with
@@ -205,11 +207,15 @@ func (c *CrawlerTool) InvokableRun(ctx context.Context, argumentsInJSON string,
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
body, err := io.ReadAll(io.LimitReader(resp.Body, maxCrawlerResponseBytes+1))
if err != nil {
return crawlerStubResult(crawlerResult{URL: targetURL, Status: resp.StatusCode, Error: "read body: " + err.Error()}),
fmt.Errorf("crawler: read body: %w", err)
}
if len(body) > maxCrawlerResponseBytes {
err = errors.New("crawler: response too large")
return crawlerStubResult(crawlerResult{URL: targetURL, Status: resp.StatusCode, Error: err.Error()}), err
}
page, err := extractPage(body)
if err != nil {

32
internal/agent/tool/crawler_test.go
View File

@@ -103,6 +103,38 @@ func TestCrawler_FetchesAndExtractsText(t *testing.T) {
}
}
func TestCrawler_RejectsOversizedResponse(t *testing.T) {
t.Parallel()
ctx := t.Context()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(strings.Repeat("x", maxCrawlerResponseBytes+1)))
}))
defer srv.Close()
loopbackResolver := func(rawURL string) (string, net.IP, error) {
u, err := url.Parse(rawURL)
if err != nil {
return "", nil, err
}
host := u.Hostname()
return host, net.ParseIP(host), nil
}
c := NewCrawlerTool().WithResolver(loopbackResolver)
out, err := c.InvokableRun(ctx, `{"query":`+jsonString(srv.URL)+`}`)
if err == nil || !strings.Contains(err.Error(), "response too large") {
t.Fatalf("err = %v, want response too large", err)
}
var got crawlerResult
if jerr := json.Unmarshal([]byte(out), &got); jerr != nil {
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
}
if !strings.Contains(got.Error, "response too large") {
t.Fatalf("_ERROR = %q, want response too large", got.Error)
}
}
func TestCrawler_RejectsMaxDepthGreaterThanZero(t *testing.T) {
t.Parallel()
ctx := t.Context()