From 3e542871c63be899aa272346d822a9eed15056e4 Mon Sep 17 00:00:00 2001
From: Tamius Han <tamius.han@gmail.com>
Date: Thu, 2 Mar 2023 01:08:16 +0100
Subject: [PATCH] Attach source frame info when communicating between frames on
 the same page

---
 src/ext/lib/EventBus.ts          |  1 +
 src/ext/lib/comms/CommsClient.ts | 22 ++++++++++++++++++++--
 src/ext/lib/comms/CommsServer.ts |  8 +++++++-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/ext/lib/EventBus.ts b/src/ext/lib/EventBus.ts
index c09d71c..283bd1d 100644
--- a/src/ext/lib/EventBus.ts
+++ b/src/ext/lib/EventBus.ts
@@ -15,6 +15,7 @@ export interface EventBusContext {
     sender?: any,
     port?: any,
     frame?: any,
+    sourceFrame?: IframeData
     forwardTo?: 'all' | 'active' | 'contentScript' | 'server' | 'sameOrigin' | 'popup' | 'all-frames',
   }
 }
diff --git a/src/ext/lib/comms/CommsClient.ts b/src/ext/lib/comms/CommsClient.ts
index 18b3f59..7619601 100644
--- a/src/ext/lib/comms/CommsClient.ts
+++ b/src/ext/lib/comms/CommsClient.ts
@@ -133,8 +133,26 @@ class CommsClient {
     return browser.runtime.sendMessage(null, message, null);
   }
 
-  processReceivedMessage(message){
-    this.eventBus.send(message.command, message.config, {origin: CommsOrigin.Server});
+  processReceivedMessage(receivedMessage){
+    // when sending between frames, message will be enriched with two new properties
+    const {_sourceFrame, _sourcePort, ...message} = receivedMessage;
+
+    let comms;
+    if (_sourceFrame || _sourcePort) {
+      comms = {
+        port: _sourcePort,
+        sourceFrame: _sourceFrame
+      }
+    }
+
+    this.eventBus.send(
+      message.command,
+      message.config,
+      {
+        comms,
+        origin: CommsOrigin.Server
+      }
+    );
   }
 }
 
diff --git a/src/ext/lib/comms/CommsServer.ts b/src/ext/lib/comms/CommsServer.ts
index 6648b87..cbe1ee3 100644
--- a/src/ext/lib/comms/CommsServer.ts
+++ b/src/ext/lib/comms/CommsServer.ts
@@ -185,9 +185,15 @@ class CommsServer {
   private async sendToOtherFrames(message, context) {
     const sender = context.comms.sourceFrame;
 
+    const enrichedMessage = {
+      message,
+      _sourceFrame: context.comms.sourceFrame,
+      _sourcePort: context.comms.port
+    }
+
     for (const frame in this.ports[sender.tabId]) {
       if (frame !== sender.frameId) {
-        this.sendToFrameContentScripts(message, sender.tabId, sender.frameId);
+        this.sendToFrameContentScripts(enrichedMessage, sender.tabId, sender.frameId);
       }
     }
   }