mirror of
https://github.com/tiennm99/MTClaw.git
synced 2026-08-04 19:30:27 +00:00
tiennm99
fb61a57785
YAML config with env-indirected secrets, SQLite persistence, OpenAI provider, think/act/observe agent loop, fs/web_fetch/exec tools behind deny->allow->mode policy with approval flows, Telegram long-polling channel with gating and inline approvals, gateway with per-session serialization + instance lock + graceful shutdown, cron scheduler.
195 lines
7.7 KiB
Go
195 lines
7.7 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/netip"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/tiennm99/MTClaw/internal/agent"
|
|
)
|
|
|
|
func newWebFetchTool(t *testing.T) *webFetchTool {
|
|
t.Helper()
|
|
return &webFetchTool{client: newSafeHTTPClient(5 * time.Second), maxBytes: 1 << 20}
|
|
}
|
|
|
|
// TestIsBlockedAddr_Matrix is the SSRF address matrix from the phase 5 spec:
|
|
// every address a real request must never reach, plus a normal public IP
|
|
// that must still be allowed through (guarding against over-blocking).
|
|
func TestIsBlockedAddr_Matrix(t *testing.T) {
|
|
cases := []struct {
|
|
addr string
|
|
blocked bool
|
|
}{
|
|
{"127.0.0.1", true},
|
|
{"169.254.169.254", true}, // cloud metadata, link-local range
|
|
{"10.0.0.1", true},
|
|
{"0.0.0.0", true},
|
|
{"::1", true},
|
|
{"::ffff:127.0.0.1", true}, // IPv4-mapped IPv6
|
|
{"fe80::1", true}, // IPv6 link-local
|
|
{"fc00::1", true}, // IPv6 ULA / "private"
|
|
{"::", true}, // unspecified
|
|
{"172.16.0.5", true},
|
|
{"192.168.1.1", true},
|
|
{"100.64.0.1", true}, // CGNAT
|
|
{"224.0.0.1", true}, // multicast
|
|
{"ff02::1", true}, // IPv6 multicast
|
|
{"8.8.8.8", false},
|
|
{"93.184.216.34", false}, // example.com-ish public IP
|
|
}
|
|
for _, c := range cases {
|
|
addr, err := netip.ParseAddr(c.addr)
|
|
require.NoError(t, err, c.addr)
|
|
assert.Equal(t, c.blocked, isBlockedAddr(addr), "address %s", c.addr)
|
|
}
|
|
}
|
|
|
|
func TestWebFetch_RefusesLoopbackTarget(t *testing.T) {
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: "http://127.0.0.1:1/"}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "web_fetch:")
|
|
}
|
|
|
|
// TestWebFetch_HostnameResolvingToPrivateIPRefused covers "a hostname
|
|
// resolving to a private IP" from the phase 5 spec using "localhost",
|
|
// which every OS resolves locally (via /etc/hosts or its NSS equivalent)
|
|
// with no real DNS query, so the test needs no network access. The point
|
|
// is that Control receives the *resolved* address (127.0.0.1 or ::1), not
|
|
// the hostname text, so a pre-resolution string check on "localhost" would
|
|
// have been unnecessary and a check against the wrong thing would have
|
|
// missed this case entirely.
|
|
func TestWebFetch_HostnameResolvingToPrivateIPRefused(t *testing.T) {
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: "http://localhost:1/"}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "web_fetch:")
|
|
}
|
|
|
|
func TestWebFetch_RefusesUnsupportedScheme(t *testing.T) {
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: "ftp://example.com/file"}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "unsupported scheme")
|
|
}
|
|
|
|
func TestWebFetch_RefusesEmptyURL(t *testing.T) {
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: ""}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "must not be empty")
|
|
}
|
|
|
|
func TestWebFetch_InvalidArgsReturnsResultString(t *testing.T) {
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), json.RawMessage(`{bad`), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "invalid arguments")
|
|
}
|
|
|
|
// TestWebFetch_PublicURLSucceeds proves the guard does not over-block: a
|
|
// loopback httptest server dialed via 127.0.0.1 is of course itself
|
|
// blocked, so this test instead drives the full success path (HTML
|
|
// stripping, untrusted-content note) directly against an httptest server,
|
|
// while a separate opt-in end-to-end test covers a genuine public URL.
|
|
func TestWebFetch_FetchesAndStripsHTML(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/html")
|
|
_, _ = w.Write([]byte(`<html><head><style>.x{color:red}</style><script>alert(1)</script></head><body><h1>Hello</h1><p>World & friends</p></body></html>`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
// The httptest server binds to loopback, which controlRejectUnsafeAddr
|
|
// would otherwise refuse - build a client without the SSRF dialer for
|
|
// this test so it exercises the HTML-to-text and untrusted-content-note
|
|
// behavior instead of the address guard (covered separately above).
|
|
w := &webFetchTool{client: srv.Client(), maxBytes: 1 << 20}
|
|
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: srv.URL}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "Hello")
|
|
assert.Contains(t, out, "World & friends")
|
|
assert.NotContains(t, out, "alert(1)")
|
|
assert.NotContains(t, out, "color:red")
|
|
assert.NotContains(t, out, "<h1>")
|
|
assert.Contains(t, out, "untrusted third-party text")
|
|
}
|
|
|
|
func TestWebFetch_MaxBytesCapTruncates(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = w.Write([]byte(strings.Repeat("a", 100)))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
wf := &webFetchTool{client: srv.Client(), maxBytes: 10}
|
|
out, err := wf.run(context.Background(), mustArgs(t, webFetchArgs{URL: srv.URL}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "[truncated to 10 bytes]")
|
|
}
|
|
|
|
// TestWebFetch_RedirectToLocalhostRefused proves web_fetch follows a
|
|
// redirect (via CheckRedirect) and then still refuses to connect to the
|
|
// redirect target through the same Dialer.Control hook that guards the
|
|
// original request - so the redirect target's response body is never
|
|
// returned to the caller. httptest servers can only bind to loopback, so
|
|
// both hops here are loopback and therefore both individually blocked; that
|
|
// still exercises the real code path (CheckRedirect decides to follow, then
|
|
// Control fires again on the resulting dial) rather than a live public
|
|
// origin, which the phase 5 spec explicitly does not require this test to
|
|
// have. TestWebFetch_RefusesLoopbackTarget already shows Control blocks a
|
|
// direct connect through this same Transport, and net/http dials every hop
|
|
// - original or redirected - through the same Transport.DialContext, so
|
|
// there is no separate "first hop" code path that could bypass this.
|
|
func TestWebFetch_RedirectToLocalhostRefused(t *testing.T) {
|
|
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = w.Write([]byte("should never be reached"))
|
|
}))
|
|
defer target.Close()
|
|
|
|
redirectSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, target.URL, http.StatusFound)
|
|
}))
|
|
defer redirectSrv.Close()
|
|
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: redirectSrv.URL}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "web_fetch:")
|
|
assert.NotContains(t, out, "should never be reached")
|
|
}
|
|
|
|
func TestWebFetch_TooManyRedirectsRefused(t *testing.T) {
|
|
var target *httptest.Server
|
|
target = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, target.URL+"/next", http.StatusFound)
|
|
}))
|
|
defer target.Close()
|
|
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: target.URL}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "web_fetch:")
|
|
}
|
|
|
|
// TestWebFetch_PublicURLE2E is gated behind MTCLAW_E2E=1: it requires real
|
|
// internet access, which CI and the default local run must not depend on.
|
|
func TestWebFetch_PublicURLE2E(t *testing.T) {
|
|
if os.Getenv("MTCLAW_E2E") != "1" {
|
|
t.Skip("set MTCLAW_E2E=1 to run tests that require real internet access")
|
|
}
|
|
w := newWebFetchTool(t)
|
|
out, err := w.run(context.Background(), mustArgs(t, webFetchArgs{URL: "https://example.com"}), agent.Meta{})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, out, "HTTP 200")
|
|
}
|