From ee6060ce352f9c2128f327b09aa315ec59fe6752 Mon Sep 17 00:00:00 2001
From: Flavien David <flavien.david74@gmail.com>
Date: Thu, 25 Jul 2024 20:21:19 +0200
Subject: [PATCH] Flav/iframe origin (#6526)

* Enforce viz iframe origin

* Add CSP on viz
---
 .../actions/VisualizationActionIframe.tsx     | 21 +++++++++++++------
 viz/next.config.mjs                           |  6 ++++++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
index 21a6b708964a..e850a1d921f5 100644
--- a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
+++ b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
@@ -12,7 +12,7 @@ import {
   visualizationExtractCode,
 } from "@dust-tt/types";
 import type { SetStateAction } from "react";
-import { useCallback, useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 
 import { RenderMessageMarkdown } from "@app/components/assistant/RenderMessageMarkdown";
 import { classNames } from "@app/lib/utils";
@@ -38,13 +38,15 @@ const sendResponseToIframe = <T extends VisualizationRPCCommand>(
 function useVisualizationDataHandler(
   action: VisualizationActionType,
   {
-    workspaceId,
     onRetry,
     setContentHeight,
+    vizIframeRef,
+    workspaceId,
   }: {
-    workspaceId: string;
     onRetry: () => void;
     setContentHeight: (v: SetStateAction<number>) => void;
+    vizIframeRef: React.MutableRefObject<HTMLIFrameElement | null>;
+    workspaceId: string;
   }
 ) {
   const extractedCode = useMemo(
@@ -74,10 +76,12 @@ function useVisualizationDataHandler(
     const listener = async (event: MessageEvent) => {
       const { data } = event;
 
-      // TODO(2024-07-24 flav) Check origin.
+      const isOriginatingFromViz =
+        event.source && event.source === vizIframeRef.current?.contentWindow;
+
       if (
         !isVisualizationRPCRequest(data) ||
-        !event.source ||
+        !isOriginatingFromViz ||
         data.actionId !== action.id
       ) {
         return;
@@ -138,6 +142,7 @@ export function VisualizationActionIframe({
 }) {
   const [showIframe, setShowIframe] = useState<boolean | null>(null);
   const [contentHeight, setContentHeight] = useState(0);
+  const vizIframeRef = useRef(null);
 
   const workspaceId = owner.sId;
 
@@ -145,6 +150,7 @@ export function VisualizationActionIframe({
     workspaceId,
     onRetry,
     setContentHeight,
+    vizIframeRef,
   });
 
   useEffect(() => {
@@ -181,10 +187,13 @@ export function VisualizationActionIframe({
             style={{ height: `${contentHeight}px` }}
             className={classNames(
               "absolute left-0 top-0 max-h-[60vh] w-full",
-              !showIframe && contentHeight > 0 ? "opacity-0" : "opacity-100"
+              !showIframe && contentHeight > 0
+                ? "pointer-events-none opacity-0"
+                : "pointer-events-auto opacity-100"
             )}
           >
             <iframe
+              ref={vizIframeRef}
               className="h-full w-full"
               src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?aId=${action.id}`}
               sandbox="allow-scripts"
diff --git a/viz/next.config.mjs b/viz/next.config.mjs
index 00e5ab523908..6366fe8f50a0 100644
--- a/viz/next.config.mjs
+++ b/viz/next.config.mjs
@@ -11,6 +11,12 @@ const nextConfig = {
             key: "Access-Control-Allow-Origin",
             value: isDev ? "http://localhost:3000" : "https://dust.tt",
           },
+          {
+            key: "Content-Security-Policy",
+            value: isDev
+              ? "frame-ancestors 'self' http://localhost:3000;"
+              : "frame-ancestors 'self' https://dust.tt;",
+          },
         ],
       },
       // Allow CORS for static files.