190 lines
5.7 KiB
Go
190 lines
5.7 KiB
Go
package io
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
|
)
|
|
|
|
// httpToolForTest builds a tools set with opts and returns the http tool.
|
|
func httpToolForTest(t *testing.T, opts ...Option) *httpTool {
|
|
t.Helper()
|
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "http").(*httpTool)
|
|
}
|
|
|
|
func TestHTTPTool_Get(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"ok":true,"path":"`+r.URL.Path+`"}`)
|
|
}))
|
|
defer srv.Close()
|
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
|
|
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL + "/hello"})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
m := got.(map[string]any)
|
|
if m["status"] != http.StatusOK {
|
|
t.Fatalf("status = %v", m["status"])
|
|
}
|
|
body := m["body"].(string)
|
|
if !strings.Contains(body, `"ok":true`) || !strings.Contains(body, "/hello") {
|
|
t.Fatalf("body = %q", body)
|
|
}
|
|
if m["truncated"].(bool) {
|
|
t.Fatal("unexpected truncation")
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_Post(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
t.Fatalf("method = %s", r.Method)
|
|
}
|
|
body, _ := io.ReadAll(r.Body)
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
_, _ = w.Write(body)
|
|
}))
|
|
defer srv.Close()
|
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
|
|
got, err := ht.Execute(context.Background(), map[string]any{
|
|
"method": "POST",
|
|
"url": srv.URL,
|
|
"body": "name=cyrene",
|
|
"headers": map[string]any{"Content-Type": "application/x-www-form-urlencoded"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
if body := got.(map[string]any)["body"].(string); body != "name=cyrene" {
|
|
t.Fatalf("echoed body = %q", body)
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_UnsupportedMethod(t *testing.T) {
|
|
t.Parallel()
|
|
ht := httpToolForTest(t)
|
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "DELETE", "url": "http://example.com"})
|
|
if err == nil || !strings.Contains(err.Error(), "unsupported method") {
|
|
t.Fatalf("expected unsupported-method error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_Timeout(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
time.Sleep(200 * time.Millisecond)
|
|
_, _ = w.Write([]byte("late"))
|
|
}))
|
|
defer srv.Close()
|
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPTimeout(30*time.Millisecond))
|
|
|
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
|
if err == nil {
|
|
t.Fatal("expected timeout error")
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_RejectsBadContentType(t *testing.T) {
|
|
t.Parallel()
|
|
ht := httpToolForTest(t)
|
|
_, err := ht.Execute(context.Background(), map[string]any{
|
|
"method": "POST",
|
|
"url": "http://93.184.216.34/",
|
|
"body": "x",
|
|
"headers": map[string]any{"Content-Type": "not a real content type {{"},
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "invalid content-type") {
|
|
t.Fatalf("expected invalid content-type error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_ResponseTruncation(t *testing.T) {
|
|
t.Parallel()
|
|
payload := strings.Repeat("a", 5000)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = io.WriteString(w, payload)
|
|
}))
|
|
defer srv.Close()
|
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()), WithHTTPResponseLimit(1000))
|
|
|
|
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
m := got.(map[string]any)
|
|
if !m["truncated"].(bool) {
|
|
t.Fatal("expected truncated=true")
|
|
}
|
|
if body := m["body"].(string); len(body) != 1000 {
|
|
t.Fatalf("body len = %d", len(body))
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_SSRFDenied(t *testing.T) {
|
|
t.Parallel()
|
|
ht := httpToolForTest(t) // no whitelist
|
|
cases := []string{
|
|
"http://127.0.0.1:9/",
|
|
"http://localhost:9/",
|
|
"http://10.0.0.1/",
|
|
"http://172.16.0.1/",
|
|
"http://192.168.1.1/",
|
|
"http://169.254.1.1/",
|
|
"http://0.0.0.0/",
|
|
"http://[::1]:9/",
|
|
}
|
|
for _, u := range cases {
|
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": u})
|
|
if !errors.Is(err, ErrSSRF) {
|
|
t.Fatalf("url %q: expected ErrSSRF, got %v", u, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_SSRFAllowedViaWhitelist(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = io.WriteString(w, "ok")
|
|
}))
|
|
defer srv.Close()
|
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
|
|
|
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
|
if err != nil {
|
|
t.Fatalf("Execute (whitelisted): %v", err)
|
|
}
|
|
if body := got.(map[string]any)["body"].(string); body != "ok" {
|
|
t.Fatalf("body = %q", body)
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_UnsupportedScheme(t *testing.T) {
|
|
t.Parallel()
|
|
ht := httpToolForTest(t)
|
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": "ftp://example.com/x"})
|
|
if err == nil || !strings.Contains(err.Error(), "unsupported scheme") {
|
|
t.Fatalf("expected unsupported-scheme error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestHTTPTool_MissingURL(t *testing.T) {
|
|
t.Parallel()
|
|
ht := httpToolForTest(t)
|
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET"})
|
|
if err == nil || !strings.Contains(err.Error(), "url is required") {
|
|
t.Fatalf("expected url-required error, got %v", err)
|
|
}
|
|
}
|