fix: eliminate XSS vulnerability in CodeBlock component (#5157)

Co-authored-by: Eric Wheeler <[email protected]>
Co-authored-by: Daniel Riccio <[email protected]>

Author: 3 people

pnpm-lock.yaml

@@ -42,6 +42,7 @@
 		"debounce": "^2.1.1",
 		"fast-deep-equal": "^3.1.3",
 		"fzf": "^0.5.2",
+		"hast-util-to-jsx-runtime": "^2.3.6",
 		"i18next": "^25.0.0",
 		"i18next-http-backend": "^3.0.2",
 		"katex": "^0.16.11",

webview-ui/package.json

@@ -4,6 +4,8 @@ import { useCopyToClipboard } from "@src/utils/clipboard"
 import { getHighlighter, isLanguageLoaded, normalizeLanguage, ExtendedLanguage } from "@src/utils/highlighter"
 import { bundledLanguages } from "shiki"
 import type { ShikiTransformer } from "shiki"
+import { toJsxRuntime } from "hast-util-to-jsx-runtime"
+import { Fragment, jsx, jsxs } from "react/jsx-runtime"
 import { ChevronDown, ChevronUp, WrapText, AlignJustify, Copy, Check } from "lucide-react"
 import { useAppTranslation } from "@src/i18n/TranslationContext"
 import { StandardTooltip } from "@/components/ui"
@@ -226,7 +228,7 @@ const CodeBlock = memo(
 		const [windowShade, setWindowShade] = useState(initialWindowShade)
 		const [currentLanguage, setCurrentLanguage] = useState<ExtendedLanguage>(() => normalizeLanguage(language))
 		const userChangedLanguageRef = useRef(false)
-		const [highlightedCode, setHighlightedCode] = useState<string>("")
+		const [highlightedCode, setHighlightedCode] = useState<React.ReactNode>(null)
 		const [showCollapseButton, setShowCollapseButton] = useState(true)
 		const codeBlockRef = useRef<HTMLDivElement>(null)
 		const preRef = useRef<HTMLDivElement>(null)
@@ -253,7 +255,12 @@ const CodeBlock = memo(
 			// Set mounted state at the beginning of this effect
 			isMountedRef.current = true
 
-			const fallback = `<pre style="padding: 0; margin: 0;"><code class="hljs language-${currentLanguage || "txt"}">${source || ""}</code></pre>`
+			// Create a safe fallback using React elements instead of HTML string
+			const fallback = (
+				<pre style={{ padding: 0, margin: 0 }}>
+					<code className={`hljs language-${currentLanguage || "txt"}`}>{source || ""}</code>
+				</pre>
+			)
 
 			const highlight = async () => {
 				// Show plain text if language needs to be loaded.
@@ -266,7 +273,7 @@ const CodeBlock = memo(
 				const highlighter = await getHighlighter(currentLanguage)
 				if (!isMountedRef.current) return
 
-				const html = await highlighter.codeToHtml(source || "", {
+				const hast = await highlighter.codeToHast(source || "", {
 					lang: currentLanguage || "txt",
 					theme: document.body.className.toLowerCase().includes("light") ? "github-light" : "github-dark",
 					transformers: [
@@ -290,8 +297,25 @@ const CodeBlock = memo(
 				})
 				if (!isMountedRef.current) return
 
-				if (isMountedRef.current) {
-					setHighlightedCode(html)
+				// Convert HAST to React elements using hast-util-to-jsx-runtime
+				// This approach eliminates XSS vulnerabilities by avoiding dangerouslySetInnerHTML
+				// while maintaining the exact same visual output and syntax highlighting
+				try {
+					const reactElement = toJsxRuntime(hast, {
+						Fragment,
+						jsx,
+						jsxs,
+						// Don't override components - let them render as-is to maintain exact output
+					})
+
+					if (isMountedRef.current) {
+						setHighlightedCode(reactElement)
+					}
+				} catch (error) {
+					console.error("[CodeBlock] Error converting HAST to JSX:", error)
+					if (isMountedRef.current) {
+						setHighlightedCode(fallback)
+					}
 				}
 			}
 
@@ -783,7 +807,7 @@ const CodeBlock = memo(
 )
 
 // Memoized content component to prevent unnecessary re-renders of highlighted code
-const MemoizedCodeContent = memo(({ html }: { html: string }) => <div dangerouslySetInnerHTML={{ __html: html }} />)
+const MemoizedCodeContent = memo(({ children }: { children: React.ReactNode }) => <>{children}</>)
 
 // Memoized StyledPre component
 const MemoizedStyledPre = memo(
@@ -801,7 +825,7 @@ const MemoizedStyledPre = memo(
 		wordWrap: boolean
 		windowShade: boolean
 		collapsedHeight?: number
-		highlightedCode: string
+		highlightedCode: React.ReactNode
 		updateCodeBlockButtonPosition: (forceHide?: boolean) => void
 	}) => (
 		<StyledPre
@@ -812,7 +836,7 @@ const MemoizedStyledPre = memo(
 			collapsedHeight={collapsedHeight}
 			onMouseDown={() => updateCodeBlockButtonPosition(true)}
 			onMouseUp={() => updateCodeBlockButtonPosition(false)}>
-			<MemoizedCodeContent html={highlightedCode} />
+			<MemoizedCodeContent>{highlightedCode}</MemoizedCodeContent>
 		</StyledPre>
 	),
 )

webview-ui/src/components/common/CodeBlock.tsx

@@ -37,6 +37,43 @@ vi.mock("../../../utils/highlighter", () => {
 			const theme = options.theme === "github-light" ? "light" : "dark"
 			return `<pre><code class="hljs language-${options.lang}">${code} [${theme}-theme]</code></pre>`
 		}),
+		codeToHast: vi.fn().mockImplementation((code, options) => {
+			const theme = options.theme === "github-light" ? "light" : "dark"
+			// Return a comprehensive HAST node structure that matches Shiki's output
+			// Apply transformers if provided
+			const preNode = {
+				type: "element",
+				tagName: "pre",
+				properties: {},
+				children: [
+					{
+						type: "element",
+						tagName: "code",
+						properties: { className: [`hljs`, `language-${options.lang}`] },
+						children: [
+							{
+								type: "text",
+								value: `${code} [${theme}-theme]`,
+							},
+						],
+					},
+				],
+			}
+
+			// Apply transformers if they exist
+			if (options.transformers) {
+				for (const transformer of options.transformers) {
+					if (transformer.pre) {
+						transformer.pre(preNode)
+					}
+					if (transformer.code && preNode.children[0]) {
+						transformer.code(preNode.children[0])
+					}
+				}
+			}
+
+			return preNode
+		}),
 	}
 
 	return {