diff --git a/app/tools/url_policy.go b/app/tools/url_policy.go index 3fb19f98..8eebe2a9 100644 --- a/app/tools/url_policy.go +++ b/app/tools/url_policy.go @@ -10,7 +10,7 @@ import ( type directURLContextKey struct{} -var directURLPattern = regexp.MustCompile(`https?://[^\s<>"']+`) +var directURLPattern = regexp.MustCompile("https?://[^\\s<>\"'`]+") func WithAllowedDirectURLs(ctx context.Context, text string) context.Context { allowed := make(map[string]struct{}) @@ -40,9 +40,12 @@ func addAllowedDirectURLToMap(allowed map[string]struct{}, raw string) { func allowedDirectURL(ctx context.Context, raw string) bool { allowed, _ := ctx.Value(directURLContextKey{}).(map[string]struct{}) - raw = cleanDirectURL(raw) + cleaned := cleanDirectURL(raw) + if cleaned == "" || cleaned != raw { + return false + } - _, ok := allowed[raw] + _, ok := allowed[cleaned] return ok } diff --git a/app/tools/url_policy_test.go b/app/tools/url_policy_test.go new file mode 100644 index 00000000..1b80de9b --- /dev/null +++ b/app/tools/url_policy_test.go @@ -0,0 +1,21 @@ +//go:build windows || darwin + +package tools + +import "testing" + +func TestDirectURLsFromText_RejectsChangedToolArgument(t *testing.T) { + ctx := WithAllowedDirectURLs(t.Context(), "summarize https://attacker.example/x") + + if allowedDirectURL(ctx, "https://attacker.example/x!!!!") { + t.Fatal("expected changed tool argument to be rejected") + } +} + +func TestDirectURLsFromText_ExtractsMarkdownCodeSpanURL(t *testing.T) { + ctx := WithAllowedDirectURLs(t.Context(), "summarize `https://example.com/privacy`") + + if !allowedDirectURL(ctx, "https://example.com/privacy") { + t.Fatal("expected URL wrapped in backticks to be allowed") + } +} diff --git a/app/tools/web_fetch.go b/app/tools/web_fetch.go index ce91d1bd..bf11e63d 100644 --- a/app/tools/web_fetch.go +++ b/app/tools/web_fetch.go @@ -75,6 +75,9 @@ func (w *WebFetch) Execute(ctx context.Context, args map[string]any) (any, strin if err != nil { return nil, "", err } + for _, link := range result.Links { + addAllowedDirectURL(ctx, link) + } return result, "", nil }