added ability to resurrect connections after service worker deactivation

andreadev-it · andreadev-it · commit 37215314f0ab · 2024-06-26T18:22:58.000+02:00
diff --git a/build/bgscript.js b/build/bgscript.js
@@ -19,12 +19,10 @@ var _CustomEventTarget = _interopRequireDefault(require("./CustomEventTarget.js"
 var _Connection = require("./Connection.js");
 var _Errors = require("./Errors.js");
 function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }
+const STORED_CONNECTIONS_KEY = "bgscript.connections";
+
 /** 
  * Class that will handle all the content scripts that will connect to the background script.
- * 
- * @property {Map<string, Connection>} scriptConnections A Map that will relate every script ID to its Connection object.
- * @property {object} exposedData The properties and methods exposed to the connecting scripts.
- * @property {function} errorCallback A callback that gets fired whenever there is an error in the script. It will get passed some details about the error.
  */
 class BackgroundHandler extends _CustomEventTarget.default {
   /**
@@ -33,24 +31,53 @@ class BackgroundHandler extends _CustomEventTarget.default {
    * @param {object} exposedData An object containing all properties and methods to be exposed to the content scripts
    * @param {object} options An object that will customize how this class works.
    */
-  constructor(exposedData = {}, options = {
-    runtime: chrome.runtime
-  }) {
-    var _options$errorCallbac, _options$runtime;
+  constructor(exposedData = {}, options = {}) {
+    var _options$errorCallbac, _options$runtime, _options$storage, _options$chromeTabs;
     super();
+
+    /** @type {Map<string, Connection>} scriptConnections A Map that will relate every script ID to its Connection object. */
     this.scriptConnections = new Map(); // script-id --> connection
+    /** @property {object} exposedData The properties and methods exposed to the connecting scripts. */
     this.exposedData = exposedData;
+    /** @property {function} errorCallback A callback that gets fired whenever there is an error in the script. It will get passed some details about the error. */
     this.errorCallback = (_options$errorCallbac = options.errorCallback) !== null && _options$errorCallbac !== void 0 ? _options$errorCallbac : null;
+
+    /** @property {chrome.runtime} runtime The runtime that will be used to create the connections. */
     this.runtime = (_options$runtime = options.runtime) !== null && _options$runtime !== void 0 ? _options$runtime : chrome.runtime;
+    this.storage = (_options$storage = options.storage) !== null && _options$storage !== void 0 ? _options$storage : chrome.storage;
+    this.chromeTabs = (_options$chromeTabs = options.chromeTabs) !== null && _options$chromeTabs !== void 0 ? _options$chromeTabs : chrome.tabs;
     this.runtime.onConnect.addListener(port => this.handleNewConnection(port));
+    this.restoreConnections();
+  }
+  async restoreConnections() {
+    let data = await this.storage.local.get(STORED_CONNECTIONS_KEY);
+    if (data[STORED_CONNECTIONS_KEY]) {
+      let connections = data[STORED_CONNECTIONS_KEY];
+      for (let [scriptId, tab] of connections) {
+        this.chromeTabs.sendMessage(tab, {
+          type: _Connection.CONNECTION_PREFIX + "ping",
+          scriptId
+        });
+      }
+    }
+  }
+  async saveConnections() {
+    let conns = [];
+    for (let [scriptId, _] of this.scriptConnections) {
+      let [name, tabId] = scriptId.split("-");
+      conns.push([name, parseInt(tabId)]);
+    }
+    await this.storage.local.set({
+      [STORED_CONNECTIONS_KEY]: conns
+    });
   }
 
   /**
    * Handle a new incoming connection
    * 
    * @param {chrome.runtime.Port} port The newly created connection to a content script
    */
-  handleNewConnection(port) {
+  async handleNewConnection(port) {
     var _port$sender$tab$id, _port$sender;
     if (!this.isInternalConnection(port)) return;
     let [name, scriptId] = this.parsePortName(port);
@@ -73,6 +100,7 @@ class BackgroundHandler extends _CustomEventTarget.default {
     // see BackgroundScript.js:68
     connection.addListener("disconnect", () => this.disconnectScript(name, tabId));
     this.scriptConnections.set(scriptId, connection);
+    await this.saveConnections();
 
     // Fire the connection event
     this.fireEvent("connectionreceived", {
@@ -176,6 +204,25 @@ class BackgroundHandler extends _CustomEventTarget.default {
     return proxy;
   }
 
+  /**
+   * Get all tab ids where a specific scriptId is present.
+   *
+   * @param {string} scriptId
+   * @returns {number[]} The tab IDs
+   */
+  getScriptTabs(scriptId) {
+    let tabs = [];
+    for (let [id, connection] of this.scriptConnections) {
+      if (id.startsWith(scriptId)) {
+        var _connection$port$send;
+        if ((_connection$port$send = connection.port.sender) !== null && _connection$port$send !== void 0 && (_connection$port$send = _connection$port$send.tab) !== null && _connection$port$send !== void 0 && _connection$port$send.id) {
+          tabs.push(connection.port.sender.tab.id);
+        }
+      }
+    }
+    return tabs;
+  }
+
   /**
    * Check if a script with a specific id associated to a specific tab has made a connection to the background page.
    * 
@@ -251,6 +298,7 @@ class BackgroundScript extends _CustomEventTarget.default {
     this.context = (_options$context = options.context) !== null && _options$context !== void 0 ? _options$context : "content";
     this.runtime = (_options$runtime = options.runtime) !== null && _options$runtime !== void 0 ? _options$runtime : chrome.runtime;
     this.connectBackgroundScript();
+    this.checkForReconnection();
   }
 
   /**
@@ -311,6 +359,20 @@ class BackgroundScript extends _CustomEventTarget.default {
     return proxy;
   }
 
+  /**
+   * Check if the background script is pinging us
+   */
+  checkForReconnection() {
+    this.runtime.onMessage.addListener(async (req, sender, sendResponse) => {
+      if (this.connection != null) return;
+      if (req.type == _Connection.CONNECTION_PREFIX + "ping") {
+        if (req.scriptId == this.scriptId) {
+          this.connectBackgroundScript();
+        }
+      }
+    });
+  }
+
   /**
    * Function that returns a uuid version 4 formatted string.
    * 
@@ -474,7 +536,7 @@ class Connection extends _CustomEventTarget.default {
   }
 
   /**
-   * Decides how to answer based on the incomin message type.
+   * Decides how to answer based on the incoming message type.
    * 
    * @param {Object} message The incoming message
    */
diff --git a/src/BackgroundHandler.js b/src/BackgroundHandler.js
@@ -2,6 +2,8 @@ import CustomEventTarget from './CustomEventTarget.js';
 import { Connection, CONNECTION_PREFIX, CONNECTION_PREFIX_NOTAB } from './Connection.js';
 import { BgHandlerErrors as ERRORS, Error } from './Errors.js';
 
+const STORED_CONNECTIONS_KEY = "bgscript.connections";
+
 /** 
  * Class that will handle all the content scripts that will connect to the background script.
  */
@@ -13,7 +15,7 @@ class BackgroundHandler extends CustomEventTarget {
      * @param {object} exposedData An object containing all properties and methods to be exposed to the content scripts
      * @param {object} options An object that will customize how this class works.
      */
-    constructor(exposedData = {}, options = { runtime: chrome.runtime }) {
+    constructor(exposedData = {}, options = {}) {
         super();
         
         /** @type {Map<string, Connection>} scriptConnections A Map that will relate every script ID to its Connection object. */
@@ -26,15 +28,45 @@ class BackgroundHandler extends CustomEventTarget {
         /** @property {chrome.runtime} runtime The runtime that will be used to create the connections. */
         this.runtime = options.runtime ?? chrome.runtime;
 
+        this.storage = options.storage ?? chrome.storage;
+        this.chromeTabs = options.chromeTabs ?? chrome.tabs;
+
         this.runtime.onConnect.addListener( (port) => this.handleNewConnection(port) );
+
+        this.restoreConnections();
+    }
+
+    async restoreConnections() {
+        let data = await this.storage.local.get(STORED_CONNECTIONS_KEY);
+        
+        if (data[STORED_CONNECTIONS_KEY]) {
+            let connections = data[STORED_CONNECTIONS_KEY];
+
+            for (let [scriptId, tab] of connections) {
+                this.chromeTabs.sendMessage(tab, { 
+                    type: CONNECTION_PREFIX + "ping",
+                    scriptId
+                });
+            }
+        }
+    }
+
+    async saveConnections() {
+        let conns = [];
+        for (let [scriptId, _] of this.scriptConnections) {
+            let [name, tabId] = scriptId.split("-");
+            conns.push([name, parseInt(tabId)]);
+        }
+
+        await this.storage.local.set({ [STORED_CONNECTIONS_KEY]: conns });
     }
 
     /**
      * Handle a new incoming connection
      * 
      * @param {chrome.runtime.Port} port The newly created connection to a content script
      */
-    handleNewConnection(port) {
+    async handleNewConnection(port) {
 
         if (!this.isInternalConnection(port)) return;
 
@@ -58,6 +90,7 @@ class BackgroundHandler extends CustomEventTarget {
         connection.addListener("disconnect", () => this.disconnectScript(name, tabId) );
 
         this.scriptConnections.set(scriptId, connection);
+        await this.saveConnections();
 
         // Fire the connection event
         this.fireEvent("connectionreceived", {
diff --git a/src/BackgroundScript.js b/src/BackgroundScript.js
@@ -32,6 +32,8 @@ class BackgroundScript extends CustomEventTarget {
         this.runtime = options.runtime ?? chrome.runtime;
 
         this.connectBackgroundScript();
+
+        this.checkForReconnection();
     }
 
     /**
@@ -103,6 +105,21 @@ class BackgroundScript extends CustomEventTarget {
         let proxy = await this.connection.getProxy();
         return proxy;
     }
+    
+    /**
+     * Check if the background script is pinging us
+     */
+    checkForReconnection() {
+        this.runtime.onMessage.addListener(async (req, sender, sendResponse) => {
+            if (this.connection != null) return;
+
+            if (req.type == CONNECTION_PREFIX + "ping") {
+                if (req.scriptId == this.scriptId) {
+                    this.connectBackgroundScript();
+                }
+            }
+        });
+    }
 
     /**
      * Function that returns a uuid version 4 formatted string.
diff --git a/src/Connection.js b/src/Connection.js
@@ -146,7 +146,7 @@ export class Connection extends CustomEventTarget {
     }
 
     /**
-     * Decides how to answer based on the incomin message type.
+     * Decides how to answer based on the incoming message type.
      * 
      * @param {Object} message The incoming message
      */
diff --git a/tests/Integrations.test.js b/tests/Integrations.test.js
@@ -5,19 +5,22 @@ import { MockedWindow } from "./mocks/mockedWindow.js";
 import BackgroundHandler from "../src/BackgroundHandler.js";
 import BackgroundScript from "../src/BackgroundScript.js";
 import { waitFor, setupScripts } from "./utilities.js";
+import { MockedChromeStorage } from "./mocks/mockedChromeStorage.js";
 
 global.window = new MockedWindow();
 
 describe("The Background Handler", () => {
     test("should be constructed correctly", () => {
         let runtime = new MockedChromeRuntime();
+        let storage = new MockedChromeStorage();
+        let chromeTabs = {};
 
         const foo = () => {};
         const exposed = {
             foo,
             prop: 1
         };
-        let bgHandler = new BackgroundHandler(exposed, { runtime });
+        let bgHandler = new BackgroundHandler(exposed, { runtime, storage, chromeTabs });
 
         assert.deepEqual(bgHandler.exposedData, { foo, prop: 1});
         assert.equal(bgHandler.scriptConnections.size, 0);
diff --git a/tests/mocks/mockedChromeRuntime.js b/tests/mocks/mockedChromeRuntime.js
@@ -4,6 +4,7 @@ import { MockedEvent } from "./mockedEvent.js";
 export class MockedChromeRuntime {
     constructor(fixedTabId = null) {
         this.onConnect = new MockedEvent();
+        this.onMessage = new MockedEvent();
         this.fixedTabId = fixedTabId;
     }
 
diff --git a/tests/mocks/mockedChromeStorage.js b/tests/mocks/mockedChromeStorage.js
@@ -0,0 +1,21 @@
+class Storage {
+    data = new Map();
+
+    async get(keys) {
+        let result = new Map();
+        for (let key of keys) {
+            result.set(key, this.data.get(key));
+        }
+        return result;
+    }
+
+    async set(entries) {
+        for (let [key, value] of Object.entries(entries)) {
+            this.data.set(key, value);
+        }
+    }
+}
+
+export class MockedChromeStorage {
+    local = new Storage();
+}
diff --git a/tests/utilities.js b/tests/utilities.js
@@ -1,6 +1,7 @@
 import BackgroundHandler from "../src/BackgroundHandler.js";
 import BackgroundScript from "../src/BackgroundScript.js";
 import { MockedChromeRuntime } from "./mocks/mockedChromeRuntime.js";
+import { MockedChromeStorage } from "./mocks/mockedChromeStorage.js";
 
 export function waitFor(ms) {
     return new Promise((resolve, _) => {
@@ -10,7 +11,8 @@ export function waitFor(ms) {
 
 export function setupScripts(handlerData, scriptsData, fixedTabId = null) {
     let runtime = new MockedChromeRuntime(fixedTabId);
-    let bgHandler = new BackgroundHandler(handlerData, { runtime });
+    let storage = new MockedChromeStorage();
+    let bgHandler = new BackgroundHandler(handlerData, { runtime, storage, chromeTabs: {} });
     
     let scripts = [];
     for (let data of scriptsData) {

Original file line number	Diff line number	Diff line change
`@@ -146,7 +146,7 @@ export class Connection extends CustomEventTarget {`
`146`	`146`	`}`
`147`	`147`
`148`	`148`	`/**`
`149`		`- * Decides how to answer based on the incomin message type.`
	`149`	`+ * Decides how to answer based on the incoming message type.`
`150`	`150`	`*`
`151`	`151`	`* @param {Object} message The incoming message`
`152`	`152`	`*/`
Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@ import { MockedEvent } from "./mockedEvent.js";`
`4`	`4`	`export class MockedChromeRuntime {`
`5`	`5`	`constructor(fixedTabId = null) {`
`6`	`6`	`this.onConnect = new MockedEvent();`
	`7`	`+ this.onMessage = new MockedEvent();`
`7`	`8`	`this.fixedTabId = fixedTabId;`
`8`	`9`	`}`
`9`	`10`