From 27b5409535f2cd90eb68b4d396b44bf860c5ec93 Mon Sep 17 00:00:00 2001
From: hensm <m+github@matt.tf>
Date: Tue, 17 Jul 2018 16:56:56 +0100
Subject: [PATCH] Add messageRouter

---
 ext/src/main.js          | 62 +++++++++++++++++++---------------------
 ext/src/messageRouter.js | 25 ++++++++++++++++
 ext/webpack.config.js    |  1 +
 3 files changed, 55 insertions(+), 33 deletions(-)
 create mode 100644 ext/src/messageRouter.js

diff --git a/ext/src/main.js b/ext/src/main.js
index eb2822b..f0c24bc 100755
--- a/ext/src/main.js
+++ b/ext/src/main.js
@@ -1,5 +1,7 @@
 "use strict";
 
+import messageRouter from "./messageRouter";
+
 const _ = browser.i18n.getMessage;
 
 
@@ -20,7 +22,6 @@ const SENDER_SCRIPT_FRAMEWORK_URL =
  */
 browser.webRequest.onBeforeRequest.addListener(
         async details => {
-            console.log(details);
             switch (details.url) {
                 case SENDER_SCRIPT_URL:
                     // Content/Page script bridge
@@ -141,11 +142,9 @@ function initBridge (tabId, frameId) {
 
     port.onMessage.addListener(message => {
         // Forward shim: messages
+        // TODO: Integrate into messageRouter
         if (message.subject.startsWith("shim:")) {
-            browser.tabs.sendMessage(tabId, message, {
-                frameId
-            });
-            return;
+            browser.tabs.sendMessage(tabId, message, { frameId });
         }
     });
 }
@@ -205,38 +204,12 @@ async function openPopup (tabId) {
 }
 
 
-// Extension messages
-browser.runtime.onMessage.addListener(async (message, sender, respond) => {
+messageRouter.register("main", async (message, sender) => {
     const tabId = sender.tab.id;
-    const { frameId } = sender.tab;
-
-    // Forward bridge: messages
-    if (message.subject.startsWith("bridge:")) {
-        bridgeMap.get(tabId).postMessage(message);
-        return;
-    }
-
-    // Forward shim: messages
-    if (message.subject.startsWith("shim:")) {
-        browser.tabs.sendMessage(popupOpenerTabId, message, { frameId });
-        return;
-    }
-
-    // Forward popup messages
-    if (message.subject.startsWith("popup:")) {
-        if (popupTabId) {
-            try {
-                browser.tabs.sendMessage(popupTabId, message);
-            } catch (err) {
-                // Popup is closed
-            }
-        }
-        return;
-    }
 
     switch (message.subject) {
         case "main:initialize":
-            initBridge(tabId);
+            initBridge(tabId, sender.tab.frameId);
             break;
 
         case "main:openPopup": {
@@ -245,3 +218,26 @@ browser.runtime.onMessage.addListener(async (message, sender, respond) => {
         }
     }
 });
+
+messageRouter.register("bridge", (message, sender) => {
+    bridgeMap.get(sender.tab.id).postMessage(message);
+});
+
+messageRouter.register("shim", (message, sender) => {
+    browser.tabs.sendMessage(popupOpenerTabId, message
+          , { frameId: sender.tab.frameId });
+});
+
+messageRouter.register("popup", (message, sender) => {
+    if (!popupTabId) return;
+
+    try {
+        browser.tabs.sendMessage(popupTabId, message);
+    } catch (err) {
+        // Popup is closed
+    }
+});
+
+browser.runtime.onMessage.addListener((message, sender) => {
+    messageRouter.handleMessage(message, sender);
+});
diff --git a/ext/src/messageRouter.js b/ext/src/messageRouter.js
new file mode 100644
index 0000000..e745f15
--- /dev/null
+++ b/ext/src/messageRouter.js
@@ -0,0 +1,25 @@
+"use strict";
+
+const routeMap = new Map();
+
+
+function register (routeName, senderCallback) {
+    routeMap.set(routeName, senderCallback);
+}
+
+function deregister (routeName) {
+    routeMap.delete(routeName);
+}
+
+function handleMessage (message, sender) {
+    const destination = message.subject.split(":")[0];
+    if (routeMap.has(destination)) {
+        routeMap.get(destination)(message, sender);
+    }
+}
+
+export default {
+    register
+  , deregister
+  , handleMessage
+}
diff --git a/ext/webpack.config.js b/ext/webpack.config.js
index 461efdf..c20f1bc 100755
--- a/ext/webpack.config.js
+++ b/ext/webpack.config.js
@@ -17,6 +17,7 @@ module.exports = {
       , "contentSetup"  : `${include_path}/contentSetup.js`
       , "mediaCast"     : `${include_path}/mediaCast.js`
       , "mirroringCast" : `${include_path}/mirroringCast.js`
+      , "messageRouter" : `${include_path}/messageRouter.js`
     }
   , output: {
         filename: "[name].js"