From 00cd5bfbd1b6578845a319a0172193296e35b41f Mon Sep 17 00:00:00 2001
From: Cheok <30773478+jiafuei@users.noreply.github.com>
Date: Wed, 13 Jul 2022 15:11:25 +0800
Subject: [PATCH] feat: logout sync for embed

feat: logout for embed
---
 src/controllers/TorusController.ts | 53 +++++++++++++++++-------------
 src/modules/controllers.ts         |  7 +++-
 src/utils/helpers.ts               |  5 ++-
 3 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/src/controllers/TorusController.ts b/src/controllers/TorusController.ts
index c5fd5c4c..d2928a47 100644
--- a/src/controllers/TorusController.ts
+++ b/src/controllers/TorusController.ts
@@ -219,6 +219,8 @@ export default class TorusController extends BaseController<TorusControllerConfi
 
   private instanceId = "";
 
+  private logoutBcAttached!: boolean;
+
   constructor({ _config, _state }: { _config: Partial<TorusControllerConfig>; _state: Partial<TorusControllerState> }) {
     super({ config: _config, state: _state });
   }
@@ -777,7 +779,7 @@ export default class TorusController extends BaseController<TorusControllerConfi
   logout(req: JRPCRequest<[]>, res: JRPCResponse<boolean>, _: JRPCEngineNextCallback, end: JRPCEngineEndCallback): void {
     this.handleLogout();
     res.result = true;
-    end();
+    setTimeout(() => end(), 100); // Make sure all async ops are executed.
   }
 
   public handleLogout(): void {
@@ -1241,9 +1243,6 @@ export default class TorusController extends BaseController<TorusControllerConfi
       });
       if (waitSaving) await saveToOpenLogin;
 
-      if (isMain) {
-        this.attachLogoutBC();
-      }
       this.emit("LOGIN_RESPONSE", null, address);
       return result;
     } catch (error) {
@@ -1389,11 +1388,7 @@ export default class TorusController extends BaseController<TorusControllerConfi
           // This call sync and refresh blockchain state
           this.setSelectedAccount(selectedAddress, true);
 
-          // Listen to logout events across tabs
-          if (isMain) {
-            this.attachLogoutBC();
-          }
-
+          this.attachLogoutBC();
           return true;
         } catch (e) {
           log.error(e, "Error restoring state after successful decrypt!");
@@ -1411,6 +1406,32 @@ export default class TorusController extends BaseController<TorusControllerConfi
     return this.preferencesController.getDappList();
   }
 
+  attachLogoutBC() {
+    if (this.logoutBcAttached) {
+      log.warn("Logout BC already attached");
+      return;
+    }
+
+    const channelName = getLogoutBcChannelName(this.origin, this.userInfo);
+    const bc = new BroadcastChannel<LogoutMessage>(channelName);
+    this.logoutBcAttached = true;
+
+    const thisInstance = this.instanceId.slice(0, 8);
+    const eventListener = (msg: LogoutMessage) => {
+      if (thisInstance === msg.instanceId) return;
+      bc.removeEventListener("message", eventListener);
+      bc.close()
+        .then(() => {
+          this.logoutBcAttached = false;
+          this.emit("logout", true);
+          if (!isMain) this.notifyEmbedLogout();
+          return null;
+        })
+        .catch((err) => log.error("broadcastchannel close error", err));
+    };
+    bc.addEventListener("message", eventListener);
+  }
+
   private async providerRequestAccounts(req: JRPCRequest<unknown>) {
     const accounts = await this.requestAccounts(req);
 
@@ -1700,18 +1721,4 @@ export default class TorusController extends BaseController<TorusControllerConfi
     };
     this.embedController.initializeProvider(commProviderHandlers);
   }
-
-  private attachLogoutBC() {
-    const channelName = getLogoutBcChannelName(this.origin, this.userInfo);
-    const bc = new BroadcastChannel<LogoutMessage>(channelName);
-    const thisInstance = this.instanceId.slice(0, 8);
-    const eventListener = (msg: LogoutMessage) => {
-      if (thisInstance === msg.instanceId) return;
-      bc.removeEventListener("message", eventListener);
-      bc.close()
-        .then(() => this.emit("logout", true))
-        .catch((err) => log.error("broadcastchannel close error", err));
-    };
-    bc.addEventListener("message", eventListener);
-  }
 }
diff --git a/src/modules/controllers.ts b/src/modules/controllers.ts
index 3be321ca..730471aa 100644
--- a/src/modules/controllers.ts
+++ b/src/modules/controllers.ts
@@ -418,6 +418,11 @@ class ControllerModule extends VuexModule {
     this.torus.on("logout", (fromBC?: boolean) => {
       this.logout(fromBC);
     });
+    this.torus.on("LOGIN_RESPONSE", (message?: string, address?: string) => {
+      if (message === null && address) {
+        this.torus.attachLogoutBC();
+      }
+    });
     this.setInstanceId(instanceId);
 
     if (!isMain) {
@@ -490,8 +495,8 @@ class ControllerModule extends VuexModule {
   async logout(fromBC?: boolean): Promise<void> {
     if (isMain && this.selectedAddress) {
       this.openloginLogout();
-      if (!fromBC) logoutWithBC(this.torus.origin, this.instanceId, this.torus.userInfo);
     }
+    if (!fromBC) await logoutWithBC(this.torus.origin, this.instanceId, this.torus.userInfo);
     const initialState = { ...cloneDeep(DEFAULT_STATE) };
     // this.updateTorusState(initialState);
 
diff --git a/src/utils/helpers.ts b/src/utils/helpers.ts
index 61d3c8a5..ad5b9c2b 100644
--- a/src/utils/helpers.ts
+++ b/src/utils/helpers.ts
@@ -210,9 +210,8 @@ export const logoutWithBC = async (origin: string, _instanceId: string, userInfo
   const bc = new BroadcastChannel<LogoutMessage>(`${channelName}`, { server: { timeout: 5 } });
   const timestamp = new Date().getTime();
   const instanceId = _instanceId.slice(0, 8);
-  bc.postMessage({ instanceId, timestamp })
-    .then(() => bc.close())
-    .catch((err) => log.error(err));
+  await bc.postMessage({ instanceId, timestamp });
+  await bc.close();
 };
 
 export function getBrowserKey() {