feat: Stricter AI Message Rendering to Prevent XSS

api/app/clients/BaseClient.js

@@ -473,6 +473,7 @@ class BaseClient {
       conversationId,
       parentMessageId: userMessage.messageId,
       isCreatedByUser: false,
+      isEdited,
       model: this.modelOptions.model,
       sender: this.sender,
       text: addSpaceIfNeeded(generation) + (await this.sendCompletion(payload, opts)),

api/app/clients/PluginsClient.js

@@ -242,6 +242,7 @@ class PluginsClient extends OpenAIClient {
     }
     const {
       user,
+      isEdited,
       conversationId,
       responseMessageId,
       saveOptions,
@@ -293,6 +294,7 @@ class PluginsClient extends OpenAIClient {
       conversationId,
       parentMessageId: userMessage.messageId,
       isCreatedByUser: false,
+      isEdited,
       model: this.modelOptions.model,
       sender: this.sender,
       promptTokens,

api/models/Message.js

@@ -14,6 +14,7 @@ module.exports = {
     error,
     unfinished,
     cancelled,
+    isEdited = false,
     finish_reason = null,
     tokenCount = null,
     plugin = null,
@@ -34,6 +35,7 @@ module.exports = {
           sender,
           text,
           isCreatedByUser,
+          isEdited,
           finish_reason,
           error,
           unfinished,
@@ -63,6 +65,7 @@ module.exports = {
   async updateMessage(message) {
     try {
       const { messageId, ...update } = message;
+      update.isEdited = true;
       const updatedMessage = await Message.findOneAndUpdate({ messageId }, update, { new: true });
 
       if (!updatedMessage) {
@@ -77,6 +80,7 @@ module.exports = {
         text: updatedMessage.text,
         isCreatedByUser: updatedMessage.isCreatedByUser,
         tokenCount: updatedMessage.tokenCount,
+        isEdited: true,
       };
     } catch (err) {
       console.error(`Error updating message: ${err}`);

api/models/schema/messageSchema.js

@@ -55,6 +55,10 @@ const messageSchema = mongoose.Schema(
       required: true,
       default: false,
     },
+    isEdited: {
+      type: Boolean,
+      default: false,
+    },
     unfinished: {
       type: Boolean,
       default: false,

api/server/routes/edit/anthropic.js

@@ -61,6 +61,7 @@ router.post(
             text: partialText,
             unfinished: true,
             cancelled: false,
+            isEdited: true,
             error: false,
           });
         }

api/server/routes/edit/gptPlugins.js

@@ -80,6 +80,7 @@ router.post(
             model: endpointOption.modelOptions.model,
             unfinished: true,
             cancelled: false,
+            isEdited: true,
             error: false,
           });
         }

api/server/routes/edit/openAI.js

@@ -63,6 +63,7 @@ router.post(
             model: endpointOption.modelOptions.model,
             unfinished: true,
             cancelled: false,
+            isEdited: true,
             error: false,
           });
         }

api/server/routes/endpoints/schemas.js

@@ -12,47 +12,6 @@ const EModelEndpoint = {
 
 const eModelEndpointSchema = z.nativeEnum(EModelEndpoint);
 
-/*
-const tMessageSchema = z.object({
-  messageId: z.string(),
-  clientId: z.string().nullable().optional(),
-  conversationId: z.string().nullable(),
-  parentMessageId: z.string().nullable(),
-  sender: z.string(),
-  text: z.string(),
-  isCreatedByUser: z.boolean(),
-  error: z.boolean(),
-  createdAt: z
-    .string()
-    .optional()
-    .default(() => new Date().toISOString()),
-  updatedAt: z
-    .string()
-    .optional()
-    .default(() => new Date().toISOString()),
-  current: z.boolean().optional(),
-  unfinished: z.boolean().optional(),
-  submitting: z.boolean().optional(),
-  searchResult: z.boolean().optional(),
-  finish_reason: z.string().optional(),
-});
-
-const tPresetSchema = tConversationSchema
-  .omit({
-    conversationId: true,
-    createdAt: true,
-    updatedAt: true,
-    title: true,
-  })
-  .merge(
-    z.object({
-      conversationId: z.string().optional(),
-      presetId: z.string().nullable().optional(),
-      title: z.string().nullable().optional(),
-    }),
-  );
-*/
-
 const tPluginAuthConfigSchema = z.object({
   authField: z.string(),
   label: z.string(),

client/src/components/Messages/Content/EditMessage.tsx

@@ -69,6 +69,7 @@ const EditMessage = ({
           ? {
             ...msg,
             text,
+            isEdited: true,
           }
           : msg,
       ),

client/src/components/Messages/Content/Markdown.tsx

@@ -10,7 +10,7 @@ import supersub from 'remark-supersub';
 import remarkGfm from 'remark-gfm';
 import rehypeRaw from 'rehype-raw';
 import CodeBlock from './CodeBlock';
-import { langSubset } from '~/utils';
+import { langSubset, validateIframe } from '~/utils';
 import store from '~/store';
 
 type TCodeProps = {
@@ -45,9 +45,10 @@ const Markdown = React.memo(({ content, message, showCursor }: TContentProps) =>
   const isSubmitting = useRecoilValue(store.isSubmitting);
   const latestMessage = useRecoilValue(store.latestMessage);
   const isInitializing = content === '<span className="result-streaming">█</span>';
-  const isLatestMessage = message?.messageId === latestMessage?.messageId;
+
+  const { isEdited, messageId } = message ?? {};
+  const isLatestMessage = messageId === latestMessage?.messageId;
   const currentContent = content?.replace('z-index: 1;', '') ?? '';
-  const isIFrame = currentContent.includes('<iframe');
 
   useEffect(() => {
     let timer1: NodeJS.Timeout, timer2: NodeJS.Timeout;
@@ -88,7 +89,12 @@ const Markdown = React.memo(({ content, message, showCursor }: TContentProps) =>
     [rehypeRaw],
   ];
 
-  if ((!isInitializing || !isLatestMessage) && !isIFrame) {
+  let isValidIframe: string | boolean | null = false;
+  if (!isEdited) {
+    isValidIframe = validateIframe(currentContent);
+  }
+
+  if (isEdited || ((!isInitializing || !isLatestMessage) && !isValidIframe)) {
     rehypePlugins.pop();
   }
 

client/src/hooks/useMessageHandler.ts

@@ -101,6 +101,7 @@ const useMessageHandler = () => {
       unfinished: false,
       submitting: true,
       isCreatedByUser: false,
+      isEdited: isEditOrContinue,
       error: false,
     };
 

client/src/utils/index.ts

@@ -5,6 +5,7 @@ export * from './languages';
 export { default as getError } from './getError';
 export { default as buildTree } from './buildTree';
 export { default as cleanupPreset } from './cleanupPreset';
+export { default as validateIframe } from './validateIframe';
 export { default as getLocalStorageItems } from './getLocalStorageItems';
 export { default as getDefaultConversation } from './getDefaultConversation';
 

client/src/utils/validateIframe.ts

@@ -0,0 +1,42 @@
+export default function validateIframe(content: string): string | boolean | null {
+  const hasValidIframe =
+    content.includes('<iframe role="presentation" style="') &&
+    content.includes('src="https://www.bing.com/images/create');
+
+  if (!hasValidIframe) {
+    return false;
+  }
+
+  const iframeRegex = /<iframe\s[^>]*?>/g;
+  const iframeMatches = content.match(iframeRegex);
+
+  if (!iframeMatches || iframeMatches.length > 1) {
+    return false;
+  }
+
+  const parser = new DOMParser();
+  const parsedHtml = parser.parseFromString(content, 'text/html');
+
+  const potentiallyHarmfulTags = ['script', 'img', 'style', 'div', 'a', 'input', 'button', 'form'];
+  for (const tag of potentiallyHarmfulTags) {
+    const elements = parsedHtml.getElementsByTagName(tag);
+
+    if (elements.length > 0) {
+      return false;
+    }
+  }
+
+  const iframes = parsedHtml.getElementsByTagName('iframe');
+
+  if (iframes.length !== 1) {
+    return false;
+  }
+
+  const iframe = iframes[0];
+
+  // Verify role and src attributes
+  const role = iframe.getAttribute('role');
+  const src = iframe.getAttribute('src');
+
+  return role === 'presentation' && src && src.startsWith('https://www.bing.com/images/create');
+}

packages/data-provider/src/schemas.ts

@@ -77,6 +77,7 @@ export const tMessageSchema = z.object({
   sender: z.string(),
   text: z.string(),
   generation: z.string().nullable().optional(),
+  isEdited: z.boolean().optional(),
   isCreatedByUser: z.boolean(),
   error: z.boolean(),
   createdAt: z