feat: stricter iframe validation

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 = {
@@ -47,7 +47,7 @@ const Markdown = React.memo(({ content, message, showCursor }: TContentProps) =>
   const isInitializing = content === '<span className="result-streaming">█</span>';
   const isLatestMessage = message?.messageId === latestMessage?.messageId;
   const currentContent = content?.replace('z-index: 1;', '') ?? '';
-  const isIFrame = currentContent.includes('<iframe');
+  const isValidIFrame = validateIframe(currentContent);
 
   useEffect(() => {
     let timer1: NodeJS.Timeout, timer2: NodeJS.Timeout;
@@ -88,7 +88,7 @@ const Markdown = React.memo(({ content, message, showCursor }: TContentProps) =>
     [rehypeRaw],
   ];
 
-  if ((!isInitializing || !isLatestMessage) && !isIFrame) {
+  if ((!isInitializing || !isLatestMessage) && !isValidIFrame) {
     rehypePlugins.pop();
   }
 

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');
+}