From 86f72591a93cda1e8a243ee0f6c78763ba6e0d89 Mon Sep 17 00:00:00 2001
From: Ben Oehlkers <boehlkers@athenahealth.com>
Date: Wed, 9 Mar 2022 11:03:06 -0500
Subject: [PATCH] feat: find apollo client instance on embedded iframes

---
 development/client/public/iframe.html    | 12 +++++
 src/extension/tab/__tests__/hook.test.ts | 57 ++++++++++++++++++++++++
 src/extension/tab/hook.ts                | 43 +++++++++++++++---
 3 files changed, 106 insertions(+), 6 deletions(-)
 create mode 100644 development/client/public/iframe.html
 create mode 100644 src/extension/tab/__tests__/hook.test.ts
diff --git a/development/client/public/iframe.html b/development/client/public/iframe.html
new file mode 100644
index 000000000..72fd298e0
--- /dev/null
+++ b/development/client/public/iframe.html
@@ -0,0 +1,12 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+  <title>Color App in an IFrame</title>
+</head>
+
+<body>
+  <iframe src="../index.html" style="width: 100%; height: 600px;">
+</body>
+
+</html>
\ No newline at end of file
diff --git a/src/extension/tab/__tests__/hook.test.ts b/src/extension/tab/__tests__/hook.test.ts
new file mode 100644
index 000000000..8bcbde5dd
--- /dev/null
+++ b/src/extension/tab/__tests__/hook.test.ts
@@ -0,0 +1,57 @@
+import { ApolloClient, InMemoryCache } from '@apollo/client';
+import { findClient } from '../hook';
+
+const cache = new InMemoryCache();
+
+describe('findClient', () => {
+  beforeEach(() => {
+    window.__APOLLO_CLIENT__ = undefined;
+    window.document.body.innerHTML = ``;
+  });
+
+  it('should find the client on the provided window object', () => {
+    const expectedClient = new ApolloClient({ cache });
+
+    const client = findClient();
+    expect(client).toEqual(expectedClient);
+  });
+
+  it('should find the client on an embedded iframe', () => {
+    const expectedClient = new ApolloClient({ cache, connectToDevTools: false });
+
+    const iframe = document.createElement('iframe');
+
+    window.document.body.appendChild(iframe);
+
+    iframe.contentWindow!.__APOLLO_CLIENT__ = expectedClient;
+
+    const client = findClient();
+    expect(client).toEqual(expectedClient);
+  });
+
+  it('should find the client in a nested iframe', () => {
+    const expectedClient = new ApolloClient({ cache, connectToDevTools: false });
+
+    const iframe = document.createElement('iframe');
+    const nestedIframe = document.createElement('iframe');
+
+    window.document.body.appendChild(iframe);
+
+    iframe.contentDocument!.body.appendChild(nestedIframe);
+    nestedIframe.contentWindow!.__APOLLO_CLIENT__ = expectedClient;
+
+    const client = findClient();
+    expect(client).toEqual(expectedClient);
+  });
+
+  it('should return undefined if no client instance is found', () => {
+    const iframe = document.createElement('iframe');
+    const nestedIframe = document.createElement('iframe');
+
+    window.document.body.appendChild(iframe);
+    iframe.contentDocument!.body.appendChild(nestedIframe);
+
+    const client = findClient();
+    expect(client).toBeUndefined();
+  });
+});
\ No newline at end of file
diff --git a/src/extension/tab/hook.ts b/src/extension/tab/hook.ts
index 72dc3d956..13df401bf 100755
--- a/src/extension/tab/hook.ts
+++ b/src/extension/tab/hook.ts
@@ -34,7 +34,7 @@ declare global {
   type TCache = any;
 
   interface Window {
-    __APOLLO_CLIENT__: ApolloClient<TCache>;
+    __APOLLO_CLIENT__?: ApolloClient<TCache>;
   }
 }
 
@@ -46,6 +46,36 @@ type Hook = {
   getCache: () => void;
 };
 
+/**
+ * Look through the window and any available iframes to find a client instance
+ * in an __APOLLO_CLIENT__ property.
+ */
+export function findClient(): ApolloClient<any> | undefined {
+  let currentWindow: Window | undefined = window;
+
+  // Account for nested iframes by doing a breadth-first-search through
+  // all iframes, and all iframes contained within those iframes, and so on.
+  const windowsQueue: Window[] = [];
+  do {
+    if (currentWindow.__APOLLO_CLIENT__) {
+      return currentWindow.__APOLLO_CLIENT__;
+    }
+
+    const iframesInDocument = currentWindow.document.getElementsByTagName('iframe');
+    for (const iframe of iframesInDocument) {
+      // Check for contentDocument to ensure we have access to it; i.e. same-origin policy
+      if (iframe.contentWindow && iframe.contentDocument) {
+        windowsQueue.unshift(iframe.contentWindow);
+      }
+    }
+
+    currentWindow = windowsQueue.shift();
+  }
+  while (currentWindow);
+
+  return undefined;
+}
+
 function initializeHook() {
   const hook: Hook = {
     ApolloClient: undefined,
@@ -213,14 +243,15 @@ function initializeHook() {
     }
   });
 
-  function findClient() {
+  function findAndRegisterClient() {
     let interval;
     let count = 0;
 
     function initializeDevtoolsHook() {
       if (count++ > 10) clearInterval(interval);
-      if (window.__APOLLO_CLIENT__) {
-        hook.ApolloClient = window.__APOLLO_CLIENT__;
+      const foundClient = findClient();
+      if (foundClient) {
+        hook.ApolloClient = foundClient;
         hook.ApolloClient.__actionHookForDevTools(handleActionHookForDevtools);
         hook.getQueries = () =>
           getQueries((hook.ApolloClient as any).queryManager.queries);
@@ -231,7 +262,7 @@ function initializeHook() {
                 (
                   hook.ApolloClient as any
                 ).queryManager.mutationStore?.getStore()
-              : // Apollo Client 3.3
+                : // Apollo Client 3.3
                 (hook.ApolloClient as any).queryManager.mutationStore
           );
         hook.getCache = () => hook.ApolloClient?.cache.extract(true);
@@ -245,7 +276,7 @@ function initializeHook() {
   }
 
   // Attempt to find the client on a 1-second interval for 10 seconds max
-  findClient();
+  findAndRegisterClient();
 }
 
 initializeHook();
