From 8dd4c9a417740ddd7064c17485d908f9f3ada848 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 12 Aug 2026 15:01:41 +0100 Subject: [PATCH] fix: more mermaid protection --- .../src/components/DocumentArtifactView.tsx | 13 ++- .../src/components/MermaidRenderer.test.tsx | 54 +++++++++++ frontend/src/components/MermaidRenderer.tsx | 51 ++++++---- .../src/components/mermaidRenderQueue.test.ts | 95 +++++++++++++++++++ .../src/components/mermaidSecurity.test.ts | 35 +++++++ frontend/src/components/mermaidSecurity.ts | 46 +++++++++ 6 files changed, 266 insertions(+), 28 deletions(-) create mode 100644 frontend/src/components/MermaidRenderer.test.tsx create mode 100644 frontend/src/components/mermaidRenderQueue.test.ts create mode 100644 frontend/src/components/mermaidSecurity.test.ts create mode 100644 frontend/src/components/mermaidSecurity.ts diff --git a/frontend/src/components/DocumentArtifactView.tsx b/frontend/src/components/DocumentArtifactView.tsx index 61db9692..1050c3e2 100644 --- a/frontend/src/components/DocumentArtifactView.tsx +++ b/frontend/src/components/DocumentArtifactView.tsx @@ -1,4 +1,3 @@ -import mermaid from 'mermaid'; import { Download, FileText, History, RotateCcw } from 'lucide-react'; import { useEffect, useMemo, useRef, useState, type ReactNode } from 'react'; import { useSelector } from 'react-redux'; @@ -7,6 +6,7 @@ import userService from '../api/services/userService'; import { useDarkTheme } from '../hooks'; import { selectToken } from '../preferences/preferenceSlice'; import MarkdownPreview from './MarkdownPreview'; +import { renderMermaidDiagram } from './mermaidSecurity'; import Spinner from './Spinner'; import { buildPreviewDocument, @@ -67,13 +67,12 @@ function FramePreview({ let cancelled = false; setRendered(null); setRenderError(false); - mermaid.initialize({ - startOnLoad: false, + renderMermaidDiagram({ + id: `artifact-mermaid-${Date.now()}`, + code: source, + isDarkTheme, securityLevel: 'sandbox', - theme: isDarkTheme ? 'dark' : 'default', - }); - mermaid - .render(`artifact-mermaid-${Date.now()}`, source) + }) .then(({ svg }) => { if (!cancelled) setRendered(svg); }) diff --git a/frontend/src/components/MermaidRenderer.test.tsx b/frontend/src/components/MermaidRenderer.test.tsx new file mode 100644 index 00000000..409235c0 --- /dev/null +++ b/frontend/src/components/MermaidRenderer.test.tsx @@ -0,0 +1,54 @@ +import { act } from 'react'; +import { createRoot, type Root } from 'react-dom/client'; + +Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + +const renderMermaidDiagramMock = vi.hoisted(() => vi.fn()); + +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})); +vi.mock('react-redux', () => ({ useSelector: () => 'idle' })); +vi.mock('../hooks', () => ({ useDarkTheme: () => [false] })); +vi.mock('./CopyButton', () => ({ default: () => null })); +vi.mock('./mermaidSecurity', () => ({ + renderMermaidDiagram: renderMermaidDiagramMock, +})); + +import MermaidRenderer from './MermaidRenderer'; + +describe('MermaidRenderer', () => { + let container: HTMLDivElement; + let root: Root; + + beforeEach(() => { + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + renderMermaidDiagramMock.mockReset(); + }); + + afterEach(async () => { + await act(async () => root.unmount()); + container.remove(); + }); + + it('keeps its visible host when Mermaid removes its temporary render ID', async () => { + renderMermaidDiagramMock.mockImplementation( + async ({ id }: { id: string }) => { + // Mermaid removes an existing node with the supplied render ID. + document.getElementById(id)?.remove(); + return { svg: '' }; + }, + ); + + await act(async () => { + root.render( + B'} isLoading={false} />, + ); + }); + + expect(container.querySelector('pre.mermaid')).not.toBeNull(); + expect(container.querySelector('svg[data-rendered="true"]')).not.toBeNull(); + }); +}); diff --git a/frontend/src/components/MermaidRenderer.tsx b/frontend/src/components/MermaidRenderer.tsx index 6a9a1ef8..ff86d882 100644 --- a/frontend/src/components/MermaidRenderer.tsx +++ b/frontend/src/components/MermaidRenderer.tsx @@ -1,4 +1,3 @@ -import mermaid from 'mermaid'; import React, { useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; @@ -11,6 +10,7 @@ import { import { selectStatus } from '../conversation/conversationSlice'; import { useDarkTheme } from '../hooks'; import CopyButton from './CopyButton'; +import { renderMermaidDiagram } from './mermaidSecurity'; import { Button } from './ui/button'; import { MermaidRendererProps } from './types'; @@ -23,12 +23,14 @@ const MermaidRenderer: React.FC = ({ const diagramId = useRef( `mermaid-${Date.now()}-${Math.random().toString(36).substring(2)}`, ); + const renderSequence = useRef(0); const status = useSelector(selectStatus); const [error, setError] = useState(null); const [showCode, setShowCode] = useState(false); const [showDownloadMenu, setShowDownloadMenu] = useState(false); const downloadMenuRef = useRef(null); const containerRef = useRef(null); + const diagramRef = useRef(null); const [hoverPosition, setHoverPosition] = useState<{ x: number; y: number; @@ -83,26 +85,33 @@ const MermaidRenderer: React.FC = ({ return `${hoverPosition.x * 100}% ${hoverPosition.y * 100}%`; }; - useEffect(() => { - const renderDiagram = async () => { - mermaid.initialize({ - startOnLoad: true, - theme: isDarkTheme ? 'dark' : 'default', - securityLevel: 'loose', - suppressErrorRendering: true, - }); + const isCurrentlyLoading = + isLoading !== undefined ? isLoading : status === 'loading'; - const isCurrentlyLoading = - isLoading !== undefined ? isLoading : status === 'loading'; + useEffect(() => { + let cancelled = false; + const renderDiagram = async () => { if (!isCurrentlyLoading && code) { try { - const element = document.getElementById(diagramId.current); + setError(null); + const { svg, bindFunctions } = await renderMermaidDiagram({ + // Mermaid owns and removes this temporary ID while rendering. It + // must never match the visible host's ID. + id: `${diagramId.current}-render-${++renderSequence.current}`, + code, + isDarkTheme, + securityLevel: 'strict', + }); + if (cancelled) return; + + const element = diagramRef.current; if (element) { - element.removeAttribute('data-processed'); - await mermaid.parse(code); //syntax check - mermaid.contentLoaded(); + // Mermaid strict mode sanitizes this generated SVG before it is returned. + element.innerHTML = svg; + bindFunctions?.(element); } } catch (err) { + if (cancelled) return; console.error('Error rendering mermaid diagram:', err); setError( `Failed to render diagram: ${err instanceof Error ? err.message : String(err)}`, @@ -112,7 +121,10 @@ const MermaidRenderer: React.FC = ({ }; renderDiagram(); - }, [code, isDarkTheme, isLoading]); + return () => { + cancelled = true; + }; + }, [code, isDarkTheme, isCurrentlyLoading]); useEffect(() => { const handleClickOutside = (event: MouseEvent) => { @@ -258,8 +270,6 @@ const MermaidRenderer: React.FC = ({ { label: 'Download as MMD', action: downloadMmd }, ]; - const isCurrentlyLoading = - isLoading !== undefined ? isLoading : status === 'loading'; const showDiagramOptions = !isCurrentlyLoading && !error; const errorRender = !isCurrentlyLoading && error; @@ -395,6 +405,7 @@ const MermaidRenderer: React.FC = ({ )} 
 = ({
                 display: 'flex',
                 justifyContent: 'center',
               }}
-            >
-              {code}
-
+ /> {showCode && ( diff --git a/frontend/src/components/mermaidRenderQueue.test.ts b/frontend/src/components/mermaidRenderQueue.test.ts new file mode 100644 index 00000000..80b40fe6 --- /dev/null +++ b/frontend/src/components/mermaidRenderQueue.test.ts @@ -0,0 +1,95 @@ +const mermaidMock = vi.hoisted(() => ({ + initialize: vi.fn(), + render: vi.fn(), +})); + +vi.mock('mermaid', () => ({ default: mermaidMock })); + +import { renderMermaidDiagram } from './mermaidSecurity'; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describe('renderMermaidDiagram', () => { + beforeEach(() => { + mermaidMock.initialize.mockReset(); + mermaidMock.render.mockReset(); + }); + + it('serializes initialization with rendering across security levels', async () => { + const firstRender = deferred<{ svg: string }>(); + mermaidMock.render + .mockImplementationOnce(() => firstRender.promise) + .mockResolvedValueOnce({ svg: '' }); + + const conversation = renderMermaidDiagram({ + id: 'conversation-diagram', + code: 'flowchart LR\nA --> B', + isDarkTheme: false, + securityLevel: 'strict', + }); + const artifact = renderMermaidDiagram({ + id: 'artifact-diagram', + code: 'flowchart LR\nB --> C', + isDarkTheme: true, + securityLevel: 'sandbox', + }); + + await vi.waitFor(() => expect(mermaidMock.render).toHaveBeenCalledTimes(1)); + expect(mermaidMock.initialize).toHaveBeenLastCalledWith( + expect.objectContaining({ securityLevel: 'strict' }), + ); + + firstRender.resolve({ svg: '' }); + await conversation; + await artifact; + + expect(mermaidMock.initialize.mock.calls.map(([config]) => config)).toEqual( + [ + expect.objectContaining({ securityLevel: 'strict' }), + expect.objectContaining({ securityLevel: 'sandbox' }), + ], + ); + expect(mermaidMock.render.mock.calls).toEqual([ + ['conversation-diagram', 'flowchart LR\nA --> B'], + ['artifact-diagram', 'flowchart LR\nB --> C'], + ]); + }); + + it('renders from the code argument again when only the theme changes', async () => { + mermaidMock.render.mockResolvedValue({ svg: '' }); + const code = 'flowchart LR\nA --> B'; + + await renderMermaidDiagram({ + id: 'theme-diagram', + code, + isDarkTheme: false, + securityLevel: 'strict', + }); + await renderMermaidDiagram({ + id: 'theme-diagram', + code, + isDarkTheme: true, + securityLevel: 'strict', + }); + + expect(mermaidMock.render).toHaveBeenNthCalledWith( + 1, + 'theme-diagram', + code, + ); + expect(mermaidMock.render).toHaveBeenNthCalledWith( + 2, + 'theme-diagram', + code, + ); + expect( + mermaidMock.initialize.mock.calls.map(([config]) => config.theme), + ).toEqual(['default', 'dark']); + }); +}); diff --git a/frontend/src/components/mermaidSecurity.test.ts b/frontend/src/components/mermaidSecurity.test.ts new file mode 100644 index 00000000..312d4a71 --- /dev/null +++ b/frontend/src/components/mermaidSecurity.test.ts @@ -0,0 +1,35 @@ +import mermaid from 'mermaid'; + +import { initializeSecureMermaid } from './mermaidSecurity'; + +describe('initializeSecureMermaid', () => { + beforeEach(() => { + document.body.replaceChildren(); + }); + + it('ignores security-level downgrade attempts in init directives', async () => { + initializeSecureMermaid(false); + const storedPayload = [ + '%%{init: {"securityLevel":"loose"}}%%', + 'flowchart LR', + ' A[Click me]', + ' click A "javascript:alert(document.domain)"', + ].join('\n'); + + const parsed = await mermaid.parse(storedPayload); + + expect(parsed).toMatchObject({ diagramType: 'flowchart-v2' }); + expect(mermaid.mermaidAPI.getConfig().securityLevel).toBe('strict'); + }); + + it('continues to accept ordinary diagrams', async () => { + initializeSecureMermaid(true); + + const parsed = await mermaid.parse( + 'flowchart LR\n A[Start] --> B[Finish]', + ); + + expect(parsed).toMatchObject({ diagramType: 'flowchart-v2' }); + expect(mermaid.mermaidAPI.getConfig().theme).toBe('dark'); + }); +}); diff --git a/frontend/src/components/mermaidSecurity.ts b/frontend/src/components/mermaidSecurity.ts new file mode 100644 index 00000000..57f098b2 --- /dev/null +++ b/frontend/src/components/mermaidSecurity.ts @@ -0,0 +1,46 @@ +import mermaid, { type RenderResult } from 'mermaid'; + +export type MermaidSecurityLevel = 'strict' | 'sandbox'; + +type MermaidRenderOptions = { + id: string; + code: string; + isDarkTheme: boolean; + securityLevel: MermaidSecurityLevel; +}; + +// Mermaid's own render queue does not snapshot its process-global config. +// Keep initialization and its matching render in one application-level task. +let renderQueue: Promise = Promise.resolve(); + +/** Configure Mermaid immediately before one queued render. */ +export function initializeSecureMermaid( + isDarkTheme: boolean, + securityLevel: MermaidSecurityLevel = 'strict', +): void { + mermaid.initialize({ + startOnLoad: false, + theme: isDarkTheme ? 'dark' : 'default', + securityLevel, + suppressErrorRendering: true, + }); +} + +/** Render untrusted Mermaid source with configuration isolated from other views. */ +export function renderMermaidDiagram({ + id, + code, + isDarkTheme, + securityLevel, +}: MermaidRenderOptions): Promise { + const render = () => { + initializeSecureMermaid(isDarkTheme, securityLevel); + return mermaid.render(id, code); + }; + const result = renderQueue.then(render, render); + renderQueue = result.then( + () => undefined, + () => undefined, + ); + return result; +}