From ef823abe547aaece2c3f0f59064816d3685c0506 Mon Sep 17 00:00:00 2001
From: SuZhou-Joe <suzhou@amazon.com>
Date: Thu, 30 Nov 2023 13:46:38 +0800
Subject: [PATCH] feat: sanitize content of basic_output (#24)

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>
---
 .../parsers/basic_input_output_parser.test.ts | 26 +++++++++++++++++++
 server/parsers/basic_input_output_parser.ts   | 10 ++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/server/parsers/basic_input_output_parser.test.ts b/server/parsers/basic_input_output_parser.test.ts
index 9f72f941..6570cc47 100644
--- a/server/parsers/basic_input_output_parser.test.ts
+++ b/server/parsers/basic_input_output_parser.test.ts
@@ -29,4 +29,30 @@ describe('BasicInputOutputParser', () => {
       },
     ]);
   });
+
+  it('sanitizes markdown outputs', async () => {
+    const outputs = await BasicInputOutputParser.parserProvider({
+      input: 'test question',
+      response:
+        'normal text<b onmouseover=alert("XSS testing!")></b> <img src="image.jpg" alt="image" width="500" height="600"> !!!!!!![](http://evil.com/) ![image](http://evil.com/) [good link](https://link)',
+      conversation_id: 'test-session',
+      interaction_id: 'interaction_id',
+      create_time: '',
+    });
+
+    expect(outputs).toEqual([
+      {
+        type: 'input',
+        contentType: 'text',
+        content: 'test question',
+      },
+      {
+        content:
+          'normal text<b></b>  [](http://evil.com/) [image](http://evil.com/) [good link](https://link)',
+        contentType: 'markdown',
+        traceId: 'interaction_id',
+        type: 'output',
+      },
+    ]);
+  });
 });
diff --git a/server/parsers/basic_input_output_parser.ts b/server/parsers/basic_input_output_parser.ts
index 7febe7b7..c1769059 100644
--- a/server/parsers/basic_input_output_parser.ts
+++ b/server/parsers/basic_input_output_parser.ts
@@ -3,8 +3,16 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import createDOMPurify from 'dompurify';
+import { JSDOM } from 'jsdom';
 import { IInput, IOutput, Interaction } from '../../common/types/chat_saved_object_attributes';
 
+const sanitize = (content: string) => {
+  const window = new JSDOM('').window;
+  const DOMPurify = createDOMPurify((window as unknown) as Window);
+  return DOMPurify.sanitize(content, { FORBID_TAGS: ['img'] }).replace(/!+\[/g, '[');
+};
+
 export const BasicInputOutputParser = {
   order: 0,
   id: 'output_message',
@@ -18,7 +26,7 @@ export const BasicInputOutputParser = {
       {
         type: 'output',
         contentType: 'markdown',
-        content: interaction.response,
+        content: sanitize(interaction.response),
         traceId: interaction.interaction_id,
       },
     ];