From 506479c02b14e9e289a5347ffd1d49445d654ce4 Mon Sep 17 00:00:00 2001
From: Branden Visser <mrvisser@gmail.com>
Date: Mon, 1 Jun 2020 21:51:58 -0400
Subject: [PATCH] Add 'connected' function to ping across to the other end of
 the comms

---
 src/connector.ts       | 54 +++++++++++++++++++++++++++++++++++++++++-
 test/connector.spec.ts | 40 +++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/src/connector.ts b/src/connector.ts
index f9913db..7ed1cd3 100644
--- a/src/connector.ts
+++ b/src/connector.ts
@@ -10,6 +10,13 @@ import {
 } from './types';
 import { hasValue } from './util';
 
+type SystemProtocol = {
+  ping: {
+    body: {};
+    response: {};
+  };
+};
+
 function mkPayloadType<
   P extends Protoframe,
   T extends ProtoframeMessageType<P>
@@ -376,6 +383,26 @@ export class ProtoframePublisher<P extends Protoframe>
 
 export class ProtoframePubsub<P extends Protoframe>
   implements AbstractProtoframePubsub<P> {
+  public static async connect<P extends Protoframe>(
+    pubsub: ProtoframePubsub<P>,
+    tries = 25,
+    timeout = 500,
+  ): Promise<ProtoframePubsub<P>> {
+    for (let i = 0; i < tries; i++) {
+      try {
+        await pubsub.ping({ timeout });
+        return pubsub;
+      } catch (_) {
+        continue;
+      }
+    }
+    throw new Error(
+      `Could not connect on protocol ${pubsub.protocol.type} after ${
+        tries * timeout
+      }ms`,
+    );
+  }
+
   /**
    * We are a "parent" page that is embedding an iframe, and we wish to connect
    * to that iframe for communication.
@@ -432,6 +459,9 @@ export class ProtoframePubsub<P extends Protoframe>
     );
   }
 
+  private systemProtocol: ProtoframeDescriptor<SystemProtocol> = {
+    type: `system|${this.protocol.type}`,
+  };
   private listeners: [Window, (ev: MessageEvent) => void][] = [];
 
   constructor(
@@ -439,7 +469,29 @@ export class ProtoframePubsub<P extends Protoframe>
     private readonly targetWindow: Window,
     private readonly thisWindow: Window = window,
     private readonly targetOrigin: string = '*',
-  ) {}
+  ) {
+    // Answer to ping requests
+    handleAsk0(
+      thisWindow,
+      targetWindow,
+      this.systemProtocol,
+      'ping',
+      targetOrigin,
+      () => Promise.resolve({}),
+    );
+  }
+
+  public async ping({ timeout = 10000 }: { timeout?: number }): Promise<void> {
+    await ask0(
+      this.thisWindow,
+      this.targetWindow,
+      this.systemProtocol,
+      'ping',
+      {},
+      this.targetOrigin,
+      timeout,
+    );
+  }
 
   public handleTell<
     T extends ProtoframeMessageType<P>,
diff --git a/test/connector.spec.ts b/test/connector.spec.ts
index 56decdc..25c2929 100644
--- a/test/connector.spec.ts
+++ b/test/connector.spec.ts
@@ -120,6 +120,46 @@ describe('ProtoframeSubscriber', () => {
 });
 
 describe('ProtoframePubsub', () => {
+  describe('connect', () => {
+    it('should fail if we cannot connect within the allocated time', async () => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const iframe: any = {
+        contentWindow: {
+          // This iframe window we will post messages on is not linked to the
+          // one we are listening to. So any `ask` messages will be dropped
+          postMessage: (): void => undefined,
+        },
+      };
+      spyOn(iframe.contentWindow, 'postMessage');
+
+      const pubsub = ProtoframePubsub.parent(cacheProtocol, iframe);
+      try {
+        await expectAsync(
+          ProtoframePubsub.connect(pubsub, 5, 10),
+        ).toBeRejectedWithError(
+          'Could not connect on protocol cache after 50ms',
+        );
+
+        // Ensure it attempted to connect 5 times
+        expect(iframe.contentWindow.postMessage).toHaveBeenCalledTimes(5);
+      } finally {
+        pubsub.destroy();
+      }
+    });
+    it('should connect if there is a connector on both ends', async () => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const iframe: any = {
+        contentWindow: window,
+      };
+      const pubsub = ProtoframePubsub.parent(cacheProtocol, iframe);
+      try {
+        const connectedPubsub = await ProtoframePubsub.connect(pubsub, 5, 10);
+        expect(pubsub).toBe(connectedPubsub);
+      } finally {
+        pubsub.destroy();
+      }
+    });
+  });
   describe('parent', () => {
     it('should fail if the child iframe has no contentWindow', () => {
       // eslint-disable-next-line @typescript-eslint/no-explicit-any