Fix lint errors and test failures

This commit is contained in:
chengpeng-wang 2026-03-03 04:29:39 -05:00
parent 09d4f3ffb2
commit 094c6cf134
2 changed files with 77 additions and 58 deletions

View file

@ -594,7 +594,8 @@ func blockPrivateTarget(ctx context.Context, parsedURL *url.URL) error {
if ip == nil {
continue
}
if ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() {
if ip.IsLoopback() || ip.IsPrivate() ||
ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() {
return fmt.Errorf("requests to private/internal addresses are not allowed")
}
}
@ -620,7 +621,7 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
return ErrorResult("missing domain in URL")
}
if err := blockPrivateTarget(ctx, parsedURL); err != nil {
if err = blockPrivateTarget(ctx, parsedURL); err != nil {
return ErrorResult(err.Error())
}

View file

@ -5,6 +5,7 @@ import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
@ -16,23 +17,49 @@ import (
const testFetchLimit = int64(10 * 1024 * 1024)
// mockRoundTripper intercepts HTTP requests and returns a fake response,
// allowing tests to bypass actual network calls and SSRF checks.
// Tests must use a non-private URL (e.g. RFC 5737 range 192.0.2.0/24)
// so the SSRF check passes before the transport is invoked.
type mockRoundTripper struct {
fn func(r *http.Request) *http.Response
}
func (m *mockRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
return m.fn(r), nil
}
func newMockFetchTool(
maxChars int,
fetchLimitBytes int64,
fn func(*http.Request) *http.Response,
) (*WebFetchTool, error) {
tool, err := NewWebFetchTool(maxChars, fetchLimitBytes)
if err != nil {
return nil, err
}
tool.client = &http.Client{Transport: &mockRoundTripper{fn: fn}}
return tool, nil
}
// TestWebTool_WebFetch_Success verifies successful URL fetching
func TestWebTool_WebFetch_Success(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusOK)
w.Write([]byte("<html><body><h1>Test Page</h1><p>Content here</p></body></html>"))
}))
defer server.Close()
tool, err := NewWebFetchTool(50000, testFetchLimit)
tool, err := newMockFetchTool(50000, testFetchLimit, func(r *http.Request) *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/html"}},
Body: io.NopCloser(
strings.NewReader("<html><body><h1>Test Page</h1><p>Content here</p></body></html>"),
),
}
})
if err != nil {
t.Fatalf("Failed to create web fetch tool: %v", err)
}
ctx := context.Background()
args := map[string]any{
"url": server.URL,
"url": "http://192.0.2.1/test", // RFC 5737 documentation IP, passes SSRF check
}
result := tool.Execute(ctx, args)
@ -58,21 +85,20 @@ func TestWebTool_WebFetch_JSON(t *testing.T) {
testData := map[string]string{"key": "value", "number": "123"}
expectedJSON, _ := json.MarshalIndent(testData, "", " ")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
w.Write(expectedJSON)
}))
defer server.Close()
tool, err := NewWebFetchTool(50000, testFetchLimit)
tool, err := newMockFetchTool(50000, testFetchLimit, func(r *http.Request) *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(bytes.NewReader(expectedJSON)),
}
})
if err != nil {
logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()})
}
ctx := context.Background()
args := map[string]any{
"url": server.URL,
"url": "http://192.0.2.1/test", // RFC 5737 documentation IP, passes SSRF check
}
result := tool.Execute(ctx, args)
@ -165,21 +191,20 @@ func TestWebTool_WebFetch_MissingURL(t *testing.T) {
func TestWebTool_WebFetch_Truncation(t *testing.T) {
longContent := strings.Repeat("x", 20000)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.WriteHeader(http.StatusOK)
w.Write([]byte(longContent))
}))
defer server.Close()
tool, err := NewWebFetchTool(1000, testFetchLimit) // Limit to 1000 chars
tool, err := newMockFetchTool(1000, testFetchLimit, func(r *http.Request) *http.Response { // Limit to 1000 chars
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/plain"}},
Body: io.NopCloser(strings.NewReader(longContent)),
}
})
if err != nil {
logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()})
}
ctx := context.Background()
args := map[string]any{
"url": server.URL,
"url": "http://192.0.2.1/test", // RFC 5737 documentation IP, passes SSRF check
}
result := tool.Execute(ctx, args)
@ -205,29 +230,25 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) {
}
func TestWebFetchTool_PayloadTooLarge(t *testing.T) {
// Create a mock HTTP server
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusOK)
// Generate a payload intentionally larger than our limit.
// Limit: 10 * 1024 * 1024 (10MB). We generate 10MB + 100 bytes of the letter 'A'.
largeData := bytes.Repeat([]byte("A"), int(testFetchLimit)+100)
// Generate a payload intentionally larger than our limit.
// Limit: 10 * 1024 * 1024 (10MB). We generate 10MB + 100 bytes of the letter 'A'.
largeData := bytes.Repeat([]byte("A"), int(testFetchLimit)+100)
w.Write(largeData)
}))
// Ensure the server is shut down at the end of the test
defer ts.Close()
// Initialize the tool
tool, err := NewWebFetchTool(50000, testFetchLimit)
// Initialize the tool with a mock transport returning the oversized payload.
tool, err := newMockFetchTool(50000, testFetchLimit, func(r *http.Request) *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/html"}},
Body: io.NopCloser(bytes.NewReader(largeData)),
}
})
if err != nil {
logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()})
}
// Prepare the arguments pointing to the URL of our local mock server
// Prepare the arguments with an RFC 5737 documentation IP that passes SSRF check.
args := map[string]any{
"url": ts.URL,
"url": "http://192.0.2.1/test",
}
// Execute the tool
@ -286,25 +307,22 @@ func TestWebTool_WebSearch_MissingQuery(t *testing.T) {
// TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction
func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusOK)
w.Write(
[]byte(
`<html><body><script>alert('test');</script><style>body{color:red;}</style><h1>Title</h1><p>Content</p></body></html>`,
),
)
}))
defer server.Close()
htmlContent := `<html><body><script>alert('test');</script><style>body{color:red;}</style><h1>Title</h1><p>Content</p></body></html>`
tool, err := NewWebFetchTool(50000, testFetchLimit)
tool, err := newMockFetchTool(50000, testFetchLimit, func(r *http.Request) *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/html"}},
Body: io.NopCloser(strings.NewReader(htmlContent)),
}
})
if err != nil {
logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()})
}
ctx := context.Background()
args := map[string]any{
"url": server.URL,
"url": "http://192.0.2.1/test", // RFC 5737 documentation IP, passes SSRF check
}
result := tool.Execute(ctx, args)