From 37e0f50ef225bfddc3bd3e3de28fbc568385018f Mon Sep 17 00:00:00 2001
From: Tanuj Mittal <tamitta@microsoft.com>
Date: Fri, 9 Apr 2021 00:22:41 -0700
Subject: [PATCH] Fix telemetry from child windows of Data Explorer (#633)

* Fix telemetry from child windows of Data Explorer

* Address feedback
---
 src/Common/MessageHandler.ts     | 46 ++++++++++++++++----------------
 src/hooks/useKnockoutExplorer.ts |  9 ++++++-
 2 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/src/Common/MessageHandler.ts b/src/Common/MessageHandler.ts
index 73636a5fd..97201b8b2 100644
--- a/src/Common/MessageHandler.ts
+++ b/src/Common/MessageHandler.ts
@@ -48,32 +48,18 @@ export function sendCachedDataMessage<TResponseDataModel>(
 }
 
 export function sendMessage(data: any): void {
-  if (canSendMessage()) {
-    // We try to find data explorer window first, then fallback to current window
-    const portalChildWindow = getDataExplorerWindow(window) || window;
-    portalChildWindow.parent.postMessage(
-      {
-        signature: "pcIframe",
-        data: data,
-      },
-      portalChildWindow.document.referrer || "*"
-    );
-  }
+  _sendMessage({
+    signature: "pcIframe",
+    data: data,
+  });
 }
 
 export function sendReadyMessage(): void {
-  if (canSendMessage()) {
-    // We try to find data explorer window first, then fallback to current window
-    const portalChildWindow = getDataExplorerWindow(window) || window;
-    portalChildWindow.parent.postMessage(
-      {
-        signature: "pcIframe",
-        kind: "ready",
-        data: "ready",
-      },
-      portalChildWindow.document.referrer || "*"
-    );
-  }
+  _sendMessage({
+    signature: "pcIframe",
+    kind: "ready",
+    data: "ready",
+  });
 }
 
 export function canSendMessage(): boolean {
@@ -89,3 +75,17 @@ export function runGarbageCollector() {
     }
   });
 }
+
+const _sendMessage = (message: any): void => {
+  if (canSendMessage()) {
+    // Portal window can receive messages from only child windows
+    const portalChildWindow = getDataExplorerWindow(window) || window;
+    if (portalChildWindow === window) {
+      // Current window is a child of portal, send message to portal window
+      portalChildWindow.parent.postMessage(message, portalChildWindow.document.referrer || "*");
+    } else {
+      // Current window is not a child of portal, send message to the child window instead (which is data explorer)
+      portalChildWindow.postMessage(message, portalChildWindow.location.origin || "*");
+    }
+  }
+};
diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts
index f1fce5827..e7d5a6853 100644
--- a/src/hooks/useKnockoutExplorer.ts
+++ b/src/hooks/useKnockoutExplorer.ts
@@ -3,7 +3,7 @@ import { applyExplorerBindings } from "../applyExplorerBindings";
 import { AuthType } from "../AuthType";
 import { AccountKind, DefaultAccountExperience } from "../Common/Constants";
 import { normalizeArmEndpoint } from "../Common/EnvironmentUtility";
-import { sendReadyMessage } from "../Common/MessageHandler";
+import { sendMessage, sendReadyMessage } from "../Common/MessageHandler";
 import { configContext, Platform, updateConfigContext } from "../ConfigContext";
 import { ActionType, DataExplorerAction } from "../Contracts/ActionContracts";
 import { MessageTypes } from "../Contracts/ExplorerContracts";
@@ -266,6 +266,8 @@ async function configurePortal(explorerParams: ExplorerParams): Promise<Explorer
           if (openAction) {
             handleOpenAction(openAction, explorer.databases(), explorer);
           }
+        } else if (shouldForwardMessage(message, event.origin)) {
+          sendMessage(message);
         }
       },
       false
@@ -275,6 +277,11 @@ async function configurePortal(explorerParams: ExplorerParams): Promise<Explorer
   });
 }
 
+function shouldForwardMessage(message: PortalMessage, messageOrigin: string) {
+  // Only allow forwarding messages from the same origin
+  return messageOrigin === window.document.location.origin && message.type === MessageTypes.TelemetryInfo;
+}
+
 function shouldProcessMessage(event: MessageEvent): boolean {
   if (typeof event.data !== "object") {
     return false;