From a4d2d2f81e8ea08970695103ceed0900d7b52a19 Mon Sep 17 00:00:00 2001
From: Hackall <36754621+hackall360@users.noreply.github.com>
Date: Sat, 13 Sep 2025 16:37:10 -0700
Subject: [PATCH] Sanitize fenced blocks before markdown parsing

---
 .gitignore                      |  1 +
 index.html                      |  1 +
 js/chat/chat-init.js            |  9 ++++++---
 js/chat/chat-storage.js         |  9 ++++++---
 js/chat/markdown-sanitizer.js   | 16 ++++++++++++++++
 package-lock.json               | 25 +++++++++++++++++++++++++
 package.json                    |  5 ++++-
 tests/markdown-sanitization.mjs | 27 +++++++++++++++++++++++++++
 8 files changed, 86 insertions(+), 7 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 js/chat/markdown-sanitizer.js
 create mode 100644 package-lock.json
 create mode 100644 tests/markdown-sanitization.mjs

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..c2658d7
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+node_modules/
diff --git a/index.html b/index.html
index 32d4c7b..0f8e302 100644
--- a/index.html
+++ b/index.html
@@ -485,6 +485,7 @@ <h3 class="modal-title">Voice Chat</h3>
     <script defer src="js/ui/ui.js"></script>
 
     <script defer src="js/chat/polli-utils.js"></script>
+    <script type="module" defer src="js/chat/markdown-sanitizer.js"></script>
     <script defer src="js/chat/chat-storage.js"></script>
     <script defer src="js/chat/chat-init.js"></script>
     <script defer src="js/ui/simple.js"></script>
diff --git a/js/chat/chat-init.js b/js/chat/chat-init.js
index 9ba933e..57b89a5 100644
--- a/js/chat/chat-init.js
+++ b/js/chat/chat-init.js
@@ -30,10 +30,13 @@ document.addEventListener("DOMContentLoaded", () => {
             marginLeft: role !== "user" ? "10px" : null,
         });
         container.classList.add(role === "user" ? "user-message" : "ai-message");
-        const bubbleContent = document.createElement("div");
-        bubbleContent.classList.add("message-text");
+        const bubbleContent = document.createElement("div");
+        bubbleContent.classList.add("message-text");
         if (role === "ai") {
-            bubbleContent.innerHTML = marked.parse(content);
+            const sanitized = window.sanitizeMarkdown
+                ? window.sanitizeMarkdown(content, window.blockedFenceTypes)
+                : content;
+            bubbleContent.innerHTML = marked.parse(sanitized);
             if (imageUrls.length > 0) {
                 imageUrls.forEach(url => {
                     const imageContainer = createImageElement(url, index);
diff --git a/js/chat/chat-storage.js b/js/chat/chat-storage.js
index 266d5f7..1229aa6 100644
--- a/js/chat/chat-storage.js
+++ b/js/chat/chat-storage.js
@@ -52,10 +52,13 @@ document.addEventListener("DOMContentLoaded", () => {
             container.style.maxWidth = "60%";
             container.style.marginLeft = "10px";
         }
-        const bubbleContent = document.createElement("div");
-        bubbleContent.classList.add("message-text");
+        const bubbleContent = document.createElement("div");
+        bubbleContent.classList.add("message-text");
         if (role === "ai") {
-            bubbleContent.innerHTML = marked.parse(content);
+            const sanitized = window.sanitizeMarkdown
+                ? window.sanitizeMarkdown(content, window.blockedFenceTypes)
+                : content;
+            bubbleContent.innerHTML = marked.parse(sanitized);
             if (imageUrls.length > 0) {
                 imageUrls.forEach(url => {
                     const imageContainer = createImageElement(url);
diff --git a/js/chat/markdown-sanitizer.js b/js/chat/markdown-sanitizer.js
new file mode 100644
index 0000000..b323e5d
--- /dev/null
+++ b/js/chat/markdown-sanitizer.js
@@ -0,0 +1,16 @@
+export const defaultBlockedFenceTypes = ['image', 'audio', 'ui'];
+
+export function sanitizeMarkdown(content, blocked = defaultBlockedFenceTypes) {
+  if (!content) return '';
+  const pattern = /```(\w+)\n[\s\S]*?```/g;
+  return content.replace(pattern, (match, type) => {
+    return blocked.includes(type.toLowerCase()) ? '' : match;
+  });
+}
+
+if (typeof window !== 'undefined') {
+  window.sanitizeMarkdown = sanitizeMarkdown;
+  if (!window.blockedFenceTypes) {
+    window.blockedFenceTypes = [...defaultBlockedFenceTypes];
+  }
+}
diff --git a/package-lock.json b/package-lock.json
new file mode 100644
index 0000000..1bf7421
--- /dev/null
+++ b/package-lock.json
@@ -0,0 +1,25 @@
+{
+  "name": "unity-chat",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {
+    "": {
+      "name": "unity-chat",
+      "devDependencies": {
+        "marked": "^11.2.0"
+      }
+    },
+    "node_modules/marked": {
+      "version": "11.2.0",
+      "resolved": "https://registry.npmjs.org/marked/-/marked-11.2.0.tgz",
+      "integrity": "sha512-HR0m3bvu0jAPYiIvLUUQtdg1g6D247//lvcekpHO1WMvbwDlwSkZAX9Lw4F4YHE1T0HaaNve0tuAWuV1UJ6vtw==",
+      "dev": true,
+      "bin": {
+        "marked": "bin/marked.js"
+      },
+      "engines": {
+        "node": ">= 18"
+      }
+    }
+  }
+}
diff --git a/package.json b/package.json
index 4a6ab06..d37a193 100644
--- a/package.json
+++ b/package.json
@@ -3,7 +3,10 @@
   "private": true,
   "type": "module",
   "scripts": {
-    "test": "node tests/pollilib-smoke.mjs"
+    "test": "node tests/pollilib-smoke.mjs && node tests/markdown-sanitization.mjs"
+  },
+  "devDependencies": {
+    "marked": "^11.2.0"
   }
 }
 
diff --git a/tests/markdown-sanitization.mjs b/tests/markdown-sanitization.mjs
new file mode 100644
index 0000000..e9c2dac
--- /dev/null
+++ b/tests/markdown-sanitization.mjs
@@ -0,0 +1,27 @@
+import { strict as assert } from 'node:assert';
+import { marked } from 'marked';
+import { sanitizeMarkdown } from '../js/chat/markdown-sanitizer.js';
+
+const input = [
+  'Hello',
+  '',
+  '```image',
+  'https://example.com/cat.png',
+  '```',
+  '',
+  '```javascript',
+  "console.log('hi');",
+  '```',
+  '',
+  '---',
+  ''
+].join('\n');
+
+const sanitized = sanitizeMarkdown(input);
+assert(!sanitized.includes('cat.png'), 'Blocked fence content should be removed');
+
+const html = marked.parse(sanitized);
+assert(html.includes('<code'), 'Code block should remain after sanitization');
+assert(html.includes('<hr'), 'Horizontal rule should be rendered');
+
+console.log('markdown-sanitization test passed');