From bbc4247893bb2bff3ecbc980f71f4d513484f94f Mon Sep 17 00:00:00 2001
From: Tamius Han <tamius.han@gmail.com>
Date: Mon, 9 Mar 2020 19:39:24 +0100
Subject: [PATCH] Don't overwrite existing ports on connect

---
 src/ext/lib/comms/CommsServer.js | 52 ++++++++++++++++++++------------
 1 file changed, 33 insertions(+), 19 deletions(-)

diff --git a/src/ext/lib/comms/CommsServer.js b/src/ext/lib/comms/CommsServer.js
index 3d15c00..9432340 100644
--- a/src/ext/lib/comms/CommsServer.js
+++ b/src/ext/lib/comms/CommsServer.js
@@ -12,7 +12,7 @@ class CommsServer {
     var ths = this;
 
     if (BrowserDetect.firefox) {
-      browser.runtime.onConnect.addListener(p => ths.onConnect(p));
+      browser.runtime.onConnect.addListener(p => this.onConnect(p));
       browser.runtime.onMessage.addListener((m, sender) => ths.processReceivedMessage_nonpersistent(m, sender));
     } else {
       chrome.runtime.onConnect.addListener(p => ths.onConnect(p));
@@ -166,9 +166,11 @@ class CommsServer {
   }
 
   sendToAll(message){
-    for(var p of this.ports){
-      for(var frame in p){
-        p[frame].postMessage(message);
+    for(const tab of this.ports){
+      for(const frame in tab){
+        for (const port in tab[frame]) {
+          tab[frame][port].postMessage(message);
+        }
       }
     }
   }
@@ -186,6 +188,12 @@ class CommsServer {
     }
   }
 
+  async sendToContentScripts(message, tab, frame) {
+    for (const port in this.ports[tab][frame]) {
+      this.ports[tab][frame][port].postMessage(message);
+    }
+  }
+
   async sendToFrame(message, tab, frame) {
     this.logger.log('info', 'comms', `%c[CommsServer::sendToFrame] attempting to send message to tab ${tab}, frame ${frame}`, "background: #dda; color: #11D", message);
 
@@ -204,7 +212,7 @@ class CommsServer {
     this.logger.log('info', 'comms', `%c[CommsServer::sendToFrame] attempting to send message to tab ${tab}, frame ${frame}`, "background: #dda; color: #11D", message);
 
     try {
-      this.ports[tab][frame].postMessage(message);
+      this.sendToContentScripts(message, tab, frame);
     } catch (e) {
       this.logger.log('error', 'comms', `%c[CommsServer::sendToFrame] Sending message failed. Reason:`, "background: #dda; color: #11D", e);
     }
@@ -213,21 +221,20 @@ class CommsServer {
   async sendToActive(message) {
     this.logger.log('info', 'comms', "%c[CommsServer::sendToActive] trying to send a message to active tab. Message:", "background: #dda; color: #11D", message);
 
-    var tabs = await this._getActiveTab();
+    const tabs = await this._getActiveTab();
 
     this.logger.log('info', 'comms', "[CommsServer::_sendToActive] currently active tab(s)?", tabs);
-    for (var key in this.ports[tabs[0].id]) {
-      this.logger.log('info', 'comms', "key?", key, this.ports[tabs[0].id]);
+    for (const frame in this.ports[tabs[0].id]) {
+      this.logger.log('info', 'comms', "key?", frame, this.ports[tabs[0].id]);
     }
 
-    for (var key in this.ports[tabs[0].id]) {
-      this.ports[tabs[0].id][key].postMessage(message);
+    for (const frame in this.ports[tabs[0].id]) {
+      this.sendToContentScripts(message, tabs[0].id, frame);
+      this.ports[tabs[0].id][frame].postMessage(message);
     }
   }
 
   onConnect(port){
-    var ths = this;
-    
     // poseben primer | special case
     if (port.name === 'popup-port') {
       this.popupPort = port;
@@ -237,15 +244,22 @@ class CommsServer {
 
     var tabId = port.sender.tab.id;
     var frameId = port.sender.frameId;
-    if(! this.ports[tabId]){
+    if (! this.ports[tabId]){
       this.ports[tabId] = {}; 
     }
-    this.ports[tabId][frameId] = port;
-    this.ports[tabId][frameId].onMessage.addListener( (m,p) => ths.processReceivedMessage(m, p));
-    this.ports[tabId][frameId].onDisconnect.addListener( (p) => { 
-      delete ths.ports[p.sender.tab.id][p.sender.frameId]; 
-      if(Object.keys(ths.ports[p.sender.tab.id]).length === 0){
-        ths.ports[tabId] = undefined;
+    if (! this.ports[tabId][frameId]) {
+      this.ports[tabId][frameId] = {};
+    }
+    this.ports[tabId][frameId][port.name] = port;
+    this.ports[tabId][frameId][port.name].onMessage.addListener( (m,p) => ths.processReceivedMessage(m, p));
+
+    this.ports[tabId][frameId][port.name].onDisconnect.addListener( (p) => { 
+      delete this.ports[p.sender.tab.id][p.sender.frameId][port.name]; 
+      if (Object.keys(this.ports[tabId][frameId].length === 0)) {
+        delete this.ports[tabId][frameId];
+        if(Object.keys(this.ports[p.sender.tab.id]).length === 0) {
+          delete this.ports[tabId];
+        }
       }
     });
   }