From 4dff5d7152e20daeeab5909ffff1ea78ed1f735b Mon Sep 17 00:00:00 2001
From: Nathan Sarrazin <sarrazin.nathan@gmail.com>
Date: Fri, 12 Jul 2024 09:13:15 +0000
Subject: [PATCH] Improve sanitization of model output

---
 package-lock.json                          | 26 ++++++++++++++++++++++
 package.json                               |  2 ++
 src/lib/components/CodeBlock.svelte        |  6 +++--
 src/lib/components/chat/ChatMessage.svelte | 13 +++++------
 4 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 72deea78fcc..d0e5e6bede5 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -59,6 +59,7 @@
 				"@sveltejs/adapter-node": "^1.3.1",
 				"@sveltejs/kit": "^1.30.4",
 				"@tailwindcss/typography": "^0.5.9",
+				"@types/dompurify": "^3.0.5",
 				"@types/express": "^4.17.21",
 				"@types/js-yaml": "^4.0.9",
 				"@types/jsdom": "^21.1.1",
@@ -68,6 +69,7 @@
 				"@types/uuid": "^9.0.8",
 				"@typescript-eslint/eslint-plugin": "^6.x",
 				"@typescript-eslint/parser": "^6.x",
+				"dompurify": "^3.1.6",
 				"eslint": "^8.28.0",
 				"eslint-config-prettier": "^8.5.0",
 				"eslint-plugin-svelte": "^2.30.0",
@@ -3460,6 +3462,16 @@
 			"integrity": "sha512-COUnqfB2+ckwXXSFInsFdOAWQzCCx+a5hq2ruyj+Vjund94RJQd4LG2u9hnvJrTgunKAaax7ancBYlDrNYxA0g==",
 			"dev": true
 		},
+		"node_modules/@types/dompurify": {
+			"version": "3.0.5",
+			"resolved": "https://registry.npmjs.org/@types/dompurify/-/dompurify-3.0.5.tgz",
+			"integrity": "sha512-1Wg0g3BtQF7sSb27fJQAKck1HECM6zV1EB66j8JH9i3LCjYabJa0FSdiSgsD5K/RbrsR0SiraKacLB+T8ZVYAg==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@types/trusted-types": "*"
+			}
+		},
 		"node_modules/@types/estree": {
 			"version": "1.0.5",
 			"resolved": "https://registry.npmjs.org/@types/estree/-/estree-1.0.5.tgz",
@@ -3678,6 +3690,13 @@
 			"integrity": "sha512-Q5vtl1W5ue16D+nIaW8JWebSSraJVlK+EthKn7e7UcD4KWsaSJ8BqGPXNaPghgtcn/fhvrN17Tv8ksUsQpiplw==",
 			"dev": true
 		},
+		"node_modules/@types/trusted-types": {
+			"version": "2.0.7",
+			"resolved": "https://registry.npmjs.org/@types/trusted-types/-/trusted-types-2.0.7.tgz",
+			"integrity": "sha512-ScaPdn1dQczgbl0QFTeTOmVHFULt394XJgOQNoyVhZ6r2vLnMLJfBPd53SB52T/3G36VI1/g2MZaX0cwDuXsfw==",
+			"dev": true,
+			"license": "MIT"
+		},
 		"node_modules/@types/uuid": {
 			"version": "9.0.8",
 			"resolved": "https://registry.npmjs.org/@types/uuid/-/uuid-9.0.8.tgz",
@@ -5350,6 +5369,13 @@
 				"url": "https://github.com/fb55/domhandler?sponsor=1"
 			}
 		},
+		"node_modules/dompurify": {
+			"version": "3.1.6",
+			"resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.1.6.tgz",
+			"integrity": "sha512-cTOAhc36AalkjtBpfG6O8JimdTMWNXjiePT2xQH/ppBGi/4uIpmj8eKyIkMJErXWARyINV/sB38yf8JCLF5pbQ==",
+			"dev": true,
+			"license": "(MPL-2.0 OR Apache-2.0)"
+		},
 		"node_modules/domutils": {
 			"version": "3.1.0",
 			"resolved": "https://registry.npmjs.org/domutils/-/domutils-3.1.0.tgz",
diff --git a/package.json b/package.json
index 628e5bdfdfa..8ac4b495266 100644
--- a/package.json
+++ b/package.json
@@ -22,6 +22,7 @@
 		"@sveltejs/adapter-node": "^1.3.1",
 		"@sveltejs/kit": "^1.30.4",
 		"@tailwindcss/typography": "^0.5.9",
+		"@types/dompurify": "^3.0.5",
 		"@types/express": "^4.17.21",
 		"@types/js-yaml": "^4.0.9",
 		"@types/jsdom": "^21.1.1",
@@ -31,6 +32,7 @@
 		"@types/uuid": "^9.0.8",
 		"@typescript-eslint/eslint-plugin": "^6.x",
 		"@typescript-eslint/parser": "^6.x",
+		"dompurify": "^3.1.6",
 		"eslint": "^8.28.0",
 		"eslint-config-prettier": "^8.5.0",
 		"eslint-plugin-svelte": "^2.30.0",
diff --git a/src/lib/components/CodeBlock.svelte b/src/lib/components/CodeBlock.svelte
index dc7cbc500d9..2556bfdc2d5 100644
--- a/src/lib/components/CodeBlock.svelte
+++ b/src/lib/components/CodeBlock.svelte
@@ -1,6 +1,7 @@
 <script lang="ts">
 	import { afterUpdate } from "svelte";
 	import CopyToClipBoardBtn from "./CopyToClipBoardBtn.svelte";
+	import DOMPurify from "dompurify";
 
 	export let code = "";
 	export let lang = "";
@@ -19,8 +20,9 @@
 	<!-- eslint-disable svelte/no-at-html-tags -->
 	<pre
 		class="scrollbar-custom overflow-auto px-5 scrollbar-thumb-gray-500 hover:scrollbar-thumb-gray-400 dark:scrollbar-thumb-white/10 dark:hover:scrollbar-thumb-white/20"><code
-			class="language-{lang}">{@html highlightedCode || code.replaceAll("<", "&lt;")}</code
-		></pre>
+			class="language-{lang}">
+			{@html DOMPurify.sanitize(highlightedCode || code)}
+		</code></pre>
 	<CopyToClipBoardBtn
 		classNames="absolute top-2 right-2 invisible opacity-0 group-hover:visible group-hover:opacity-100"
 		value={code}
diff --git a/src/lib/components/chat/ChatMessage.svelte b/src/lib/components/chat/ChatMessage.svelte
index 920e3bd9442..2754f48547b 100644
--- a/src/lib/components/chat/ChatMessage.svelte
+++ b/src/lib/components/chat/ChatMessage.svelte
@@ -33,6 +33,7 @@
 	import Modal from "../Modal.svelte";
 	import ToolUpdate from "./ToolUpdate.svelte";
 	import { useSettingsStore } from "$lib/stores/settings";
+	import DOMPurify from "dompurify";
 
 	function sanitizeMd(md: string) {
 		let ret = md
@@ -53,9 +54,6 @@
 
 		return ret;
 	}
-	function unsanitizeMd(md: string) {
-		return md.replaceAll("&lt;", "<");
-	}
 
 	export let model: Model;
 	export let id: Message["id"];
@@ -106,7 +104,6 @@
 	marked.use(
 		markedKatex({
 			throwOnError: false,
-			// output: "html",
 		})
 	);
 
@@ -301,10 +298,12 @@
 				{/if}
 				{#each tokens as token}
 					{#if token.type === "code"}
-						<CodeBlock lang={token.lang} code={unsanitizeMd(token.text)} />
+						<CodeBlock lang={token.lang} code={token.text} />
 					{:else}
-						<!-- eslint-disable-next-line svelte/no-at-html-tags -->
-						{@html marked.parse(token.raw, options)}
+						{#await marked.parse(token.raw, options) then parsed}
+							<!-- eslint-disable-next-line svelte/no-at-html-tags -->
+							{@html DOMPurify.sanitize(parsed)}
+						{/await}
 					{/if}
 				{/each}
 			</div>