From b019fbae6f110a27a12fc2e371973aa00381cfcd Mon Sep 17 00:00:00 2001
From: Pallav Agarwal <pallavagarwal07@gmail.com>
Date: Mon, 28 Mar 2022 13:12:12 +0530
Subject: [PATCH] fix: iframe srcDoc was not able to postMessage that can be
 received by onMessageReceived event. Match iframe by object not origin
 (#11809)

---
 .../widgets/IframeWidget/component/index.tsx  | 19 ++++++++++++++-----
 .../src/widgets/IframeWidget/widget/index.tsx |  5 -----
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/app/client/src/widgets/IframeWidget/component/index.tsx b/app/client/src/widgets/IframeWidget/component/index.tsx
index 30eaf7a89f..375ae25db2 100644
--- a/app/client/src/widgets/IframeWidget/component/index.tsx
+++ b/app/client/src/widgets/IframeWidget/component/index.tsx
@@ -75,16 +75,25 @@ function IframeComponent(props: IframeComponentProps) {
     widgetId,
   } = props;
 
+  const frameRef = useRef<HTMLIFrameElement>(null);
+
   const isFirstRender = useRef(true);
 
   const [message, setMessage] = useState("");
 
   useEffect(() => {
+    const handler = (event: MessageEvent) => {
+      const iframeWindow =
+        frameRef.current?.contentWindow ||
+        frameRef.current?.contentDocument?.defaultView;
+      // Accept messages only from the current iframe
+      if (event.source !== iframeWindow) return;
+      onMessageReceived(event);
+    };
     // add a listener
-    window.addEventListener("message", onMessageReceived, false);
+    window.addEventListener("message", handler, false);
     // clean up
-    return () =>
-      window.removeEventListener("message", onMessageReceived, false);
+    return () => window.removeEventListener("message", handler, false);
   }, []);
 
   useEffect(() => {
@@ -129,9 +138,9 @@ function IframeComponent(props: IframeComponentProps) {
       {message ? (
         message
       ) : srcDoc ? (
-        <iframe src={source} srcDoc={srcDoc} title={title} />
+        <iframe ref={frameRef} src={source} srcDoc={srcDoc} title={title} />
       ) : (
-        <iframe src={source} title={title} />
+        <iframe ref={frameRef} src={source} title={title} />
       )}
     </IframeContainer>
   );
diff --git a/app/client/src/widgets/IframeWidget/widget/index.tsx b/app/client/src/widgets/IframeWidget/widget/index.tsx
index 96f3aeba22..8c98c41702 100644
--- a/app/client/src/widgets/IframeWidget/widget/index.tsx
+++ b/app/client/src/widgets/IframeWidget/widget/index.tsx
@@ -162,11 +162,6 @@ class IframeWidget extends BaseWidget<IframeWidgetProps, WidgetState> {
   };
 
   handleMessageReceive = (event: MessageEvent) => {
-    // Accept messages only from the current iframe
-    if (!this.props.source?.includes(event.origin)) {
-      return;
-    }
-
     this.props.updateWidgetMetaProperty("message", event.data, {
       triggerPropertyName: "onMessageReceived",
       dynamicString: this.props.onMessageReceived,