From 37ed2455eb04d8fee0239e7cc430f44b7b28b9c0 Mon Sep 17 00:00:00 2001
From: Andrea Marchesini <amarchesini@mozilla.com>
Date: Fri, 22 Oct 2021 17:25:25 +0200
Subject: [PATCH] Network key isolation when the VPN changes state

---
 src/js/background/assignManager.js        | 10 ++++++--
 src/js/background/mozillaVpnBackground.js | 30 +++++++++++++++++++----
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/js/background/assignManager.js b/src/js/background/assignManager.js
index b754548..a18317c 100644
--- a/src/js/background/assignManager.js
+++ b/src/js/background/assignManager.js
@@ -200,10 +200,16 @@ window.assignManager = {
 
     const tab = await browser.tabs.get(requestInfo.tabId);
     const result = await proxifiedContainers.retrieve(tab.cookieStoreId);
-    if (result) {
+    if (!result || !result.proxy) {
+      return Utils.DEFAULT_PROXY;
+    }
+
+    if (!result.proxy.mozProxyEnabled) {
       return result.proxy;
     }
-    return Utils.DEFAULT_PROXY;
+
+    // Let's add the isolation key.
+    return [{ ...result.proxy, connectionIsolationKey: "" + MozillaVPN_Background.isolationKey }];
   },
 
   // Before a request is handled by the browser we decide if we should
diff --git a/src/js/background/mozillaVpnBackground.js b/src/js/background/mozillaVpnBackground.js
index 5502a42..f9a147e 100644
--- a/src/js/background/mozillaVpnBackground.js
+++ b/src/js/background/mozillaVpnBackground.js
@@ -5,6 +5,8 @@ const MozillaVPN_Background = {
   MOZILLA_VPN_HIDE_MAIN_TOUT_KEY: "mozillaVpnHideMainTout",
   MOZILLA_VPN_SERVERS_KEY: "mozillaVpnServers",
 
+  _isolationKey: 0,
+
   async maybeInitPort() {
     if (this.port && this.port.error === null) {
       return;
@@ -17,10 +19,17 @@ const MozillaVPN_Background = {
       */
       this.port = await browser.runtime.connectNative("mozillavpn");
       await browser.storage.local.set({ [this.MOZILLA_VPN_INSTALLED_KEY]: true});
-      this.port.onMessage.addListener(this.handleResponse);
+      this.port.onMessage.addListener(response => this.handleResponse(response));
+
       this.postToApp("status");
       this.postToApp("servers");
 
+      // When the mozillavpn dies or the VPN disconnects, we need to increase
+      // the isolation key in order to create new proxy connections. Otherwise
+      // we could see random timeout when the browser tries to connect to an
+      // invalid proxy connection.
+      this.port.onDisconnect.addListener(() => this.increaseIsolationKey());
+
     } catch(e) {
       browser.storage.local.set({ [this.MOZILLA_VPN_INSTALLED_KEY]: false });
       browser.storage.local.set({ [this.MOZILLA_VPN_CONNECTED_KEY]: false });
@@ -54,9 +63,9 @@ const MozillaVPN_Background = {
 
   // Handle responses from MozillaVPN client
   async handleResponse(response) {
-
     if (response.error && response.error === "vpn-client-down") {
       browser.storage.local.set({ [MozillaVPN_Background.MOZILLA_VPN_CONNECTED_KEY]: false });
+      this.increaseIsolationKey();
       return;
     }
     if (response.servers) {
@@ -65,10 +74,10 @@ const MozillaVPN_Background = {
       return;
     }
 
-    if (response.status && response.status.vpn) {
+    if ((response.status && response.status.vpn) || response.t === "status") {
       browser.storage.local.set({ [MozillaVPN_Background.MOZILLA_VPN_INSTALLED_KEY]: true });
 
-      const status = response.status.vpn;
+      const status = response.status ? response.status.vpn : response.vpn;
 
       if (status === "StateOn") {
         browser.storage.local.set({ [MozillaVPN_Background.MOZILLA_VPN_CONNECTED_KEY]: true });
@@ -77,8 +86,19 @@ const MozillaVPN_Background = {
       if (status === "StateOff" || status === "StateDisconnecting") {
         browser.storage.local.set({ [MozillaVPN_Background.MOZILLA_VPN_CONNECTED_KEY]: false });
       }
+
+      // Let's increase the network key isolation at any vpn status change.
+      this.increaseIsolationKey();
     }
-  }
+  },
+
+  increaseIsolationKey() {
+    ++this._isolationKey;
+  },
+
+  get isolationKey() {
+    return this._isolationKey;
+  },
 };
 
 MozillaVPN_Background.init();