From e8b79d62605be17c55da827991bff47a69f60dc6 Mon Sep 17 00:00:00 2001
From: Tanuj Mittal <tamitta@microsoft.com>
Date: Tue, 27 Apr 2021 12:52:52 -0400
Subject: [PATCH] Use postRobot to listen for GitHub OAuth messages (#729)

---
 src/GitHub/GitHubConnector.ts         | 25 ++++++++++++++-----------
 src/GitHub/GitHubOAuthService.test.ts | 27 ++++-----------------------
 src/GitHub/GitHubOAuthService.ts      | 21 +++++++++++----------
 3 files changed, 29 insertions(+), 44 deletions(-)

diff --git a/src/GitHub/GitHubConnector.ts b/src/GitHub/GitHubConnector.ts
index 866f86f51..5c9cbff37 100644
--- a/src/GitHub/GitHubConnector.ts
+++ b/src/GitHub/GitHubConnector.ts
@@ -1,3 +1,5 @@
+import postRobot from "post-robot";
+
 export interface IGitHubConnectorParams {
   state: string;
   code: string;
@@ -6,25 +8,26 @@ export interface IGitHubConnectorParams {
 export const GitHubConnectorMsgType = "GitHubConnectorMsgType";
 
 export class GitHubConnector {
-  public start(params: URLSearchParams, window: Window & typeof globalThis) {
-    window.postMessage(
+  public async start(params: URLSearchParams, window: Window & typeof globalThis): Promise<void> {
+    await postRobot.send(
+      window,
+      GitHubConnectorMsgType,
       {
-        type: GitHubConnectorMsgType,
-        data: {
-          state: params.get("state"),
-          code: params.get("code"),
-        } as IGitHubConnectorParams,
-      },
-      window.location.origin
+        state: params.get("state"),
+        code: params.get("code"),
+      } as IGitHubConnectorParams,
+      {
+        domain: window.location.origin,
+      }
     );
   }
 }
 
 var connector = new GitHubConnector();
-window.addEventListener("load", () => {
+window.addEventListener("load", async () => {
   const openerWindow = window.opener;
   if (openerWindow) {
-    connector.start(new URLSearchParams(document.location.search), openerWindow);
+    await connector.start(new URLSearchParams(document.location.search), openerWindow);
     window.close();
   }
 });
diff --git a/src/GitHub/GitHubOAuthService.test.ts b/src/GitHub/GitHubOAuthService.test.ts
index 9a2b092af..43f510c16 100644
--- a/src/GitHub/GitHubOAuthService.test.ts
+++ b/src/GitHub/GitHubOAuthService.test.ts
@@ -1,12 +1,12 @@
 import ko from "knockout";
 import { HttpStatusCodes } from "../Common/Constants";
 import * as DataModels from "../Contracts/DataModels";
-import { JunoClient } from "../Juno/JunoClient";
-import { GitHubConnector, IGitHubConnectorParams } from "./GitHubConnector";
-import { GitHubOAuthService } from "./GitHubOAuthService";
+import Explorer from "../Explorer/Explorer";
 import { ConsoleDataType } from "../Explorer/Menus/NotificationConsole/NotificationConsoleComponent";
 import NotebookManager from "../Explorer/Notebook/NotebookManager";
-import Explorer from "../Explorer/Explorer";
+import { JunoClient } from "../Juno/JunoClient";
+import { IGitHubConnectorParams } from "./GitHubConnector";
+import { GitHubOAuthService } from "./GitHubOAuthService";
 
 const sampleDatabaseAccount: DataModels.DatabaseAccount = {
   id: "id",
@@ -85,25 +85,6 @@ describe("GitHubOAuthService", () => {
     expect(newParams.get("state")).not.toEqual(initialParams.get("state"));
   });
 
-  it("finishOAuth is called whenever GitHubConnector is started", async () => {
-    const finishOAuthCallback = jest.fn().mockImplementation();
-    gitHubOAuthService.finishOAuth = finishOAuthCallback;
-
-    const params: IGitHubConnectorParams = {
-      state: "state",
-      code: "code",
-    };
-    const searchParams = new URLSearchParams({ ...params });
-
-    const gitHubConnector = new GitHubConnector();
-    gitHubConnector.start(searchParams, window);
-
-    // GitHubConnector uses Window.postMessage and there's no good way to know when the message has received
-    await new Promise((resolve) => setTimeout(resolve, 100));
-
-    expect(finishOAuthCallback).toBeCalledWith(params);
-  });
-
   it("finishOAuth updates token", async () => {
     const data = { key: "value" };
     const getGitHubTokenCallback = jest.fn().mockReturnValue({ status: HttpStatusCodes.OK, data });
diff --git a/src/GitHub/GitHubOAuthService.ts b/src/GitHub/GitHubOAuthService.ts
index 779d70bd7..6f800e66d 100644
--- a/src/GitHub/GitHubOAuthService.ts
+++ b/src/GitHub/GitHubOAuthService.ts
@@ -1,24 +1,25 @@
 import ko from "knockout";
+import postRobot from "post-robot";
 import { HttpStatusCodes } from "../Common/Constants";
 import { handleError } from "../Common/ErrorHandlingUtils";
 import { configContext } from "../ConfigContext";
 import { AuthorizeAccessComponent } from "../Explorer/Controls/GitHub/AuthorizeAccessComponent";
 import { JunoClient } from "../Juno/JunoClient";
-import { isInvalidParentFrameOrigin } from "../Utils/MessageValidation";
 import { logConsoleInfo } from "../Utils/NotificationConsoleUtils";
 import { GitHubConnectorMsgType, IGitHubConnectorParams } from "./GitHubConnector";
 
-window.addEventListener("message", (event: MessageEvent) => {
-  if (isInvalidParentFrameOrigin(event)) {
-    return;
-  }
-
-  const msg = event.data;
-  if (msg.type === GitHubConnectorMsgType) {
-    const params = msg.data as IGitHubConnectorParams;
+postRobot.on(
+  GitHubConnectorMsgType,
+  {
+    domain: window.location.origin,
+  },
+  (event) => {
+    // Typescript definition for event is wrong. So read params by casting to <any>
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const params = (event as any).data as IGitHubConnectorParams;
     window.dataExplorer.notebookManager?.gitHubOAuthService.finishOAuth(params);
   }
-});
+);
 
 export interface IGitHubOAuthToken {
   // API properties