chore(ui): html/markdown security (#1006)

defenseunicorns · Sep 16, 2024 · 6a0038a · 6a0038a
1 parent 6d17187
commit 6a0038a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/leapfrogai_ui/src/lib/components/Message.svelte b/src/leapfrogai_ui/src/lib/components/Message.svelte
@@ -41,6 +41,7 @@
 
   // used for code formatting and handling
   const md = markdownit({
+    html: true,
     highlight: function (str: string, language: string) {
       let code: string;
       if (language && hljs.getLanguage(language)) {
@@ -178,7 +179,13 @@
               <MessagePendingSkeleton size="sm" class="mt-4" darkColor="bg-gray-500" />
             {:else}
               <!--eslint-disable-next-line svelte/no-at-html-tags -- We use DomPurity to sanitize the code snippet-->
-              {@html md.render(DOMPurify.sanitize(messageText))}
+              {@html DOMPurify.sanitize(md.render(messageText), {
+                CUSTOM_ELEMENT_HANDLING: {
+                  tagNameCheck: /^code-block$/,
+                  attributeNameCheck: /^(code|language)$/,
+                  allowCustomizedBuiltInElements: false
+                }
+              })}
               <div class="flex flex-col items-start">
                 {#each getCitations(message, $page.data.files) as { component: Component, props }}
                   <svelte:component this={Component} {...props} />

diff --git a/src/leapfrogai_ui/src/lib/web-components/CodeBlock.js b/src/leapfrogai_ui/src/lib/web-components/CodeBlock.js
@@ -79,7 +79,7 @@ export class CodeBlock extends LitElement {
 
   removeHtml = () => {
     const tmp = document.createElement('DIV');
-    tmp.innerHTML = this.code;
+    tmp.innerHTML = DOMPurify.sanitize(this.code);
     return tmp.textContent || tmp.innerText || '';
   };