browser(firefox): support Browser.setProxy method in juggler (#2464)

This lets us support network proxies per browser context.
microsoft · Jun 4, 2020 · 3c9699d · 3c9699d
1 parent d5c992e
commit 3c9699d
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 1 deletion.
diff --git a/browser_patches/firefox/BUILD_NUMBER b/browser_patches/firefox/BUILD_NUMBER
@@ -1 +1 @@
-1101
+1102
diff --git a/browser_patches/firefox/juggler/NetworkObserver.js b/browser_patches/firefox/juggler/NetworkObserver.js
@@ -19,6 +19,8 @@ const Cm = Components.manager;
 const CC = Components.Constructor;
 const helper = new Helper();
 
+const UINT32_MAX = Math.pow(2, 32)-1;
+
 const BinaryInputStream = CC('@mozilla.org/binaryinputstream;1', 'nsIBinaryInputStream', 'setInputStream');
 const BinaryOutputStream = CC('@mozilla.org/binaryoutputstream;1', 'nsIBinaryOutputStream', 'setOutputStream');
 const StorageStream = CC('@mozilla.org/storagestream;1', 'nsIStorageStream', 'init');
@@ -153,6 +155,31 @@ class NetworkObserver {
     this._postAuthChannelIdToRequestId = new Map();  // pre-auth id => post-auth id
     this._bodyListeners = new Map();  // channel id => ResponseBodyListener.
 
+    const protocolProxyService = Cc['@mozilla.org/network/protocol-proxy-service;1'].getService();
+    this._channelProxyFilter = {
+      QueryInterface: ChromeUtils.generateQI([Ci.nsIProtocolProxyChannelFilter]),
+      applyFilter: (channel, defaultProxyInfo, proxyFilter) => {
+        const originAttributes = channel.loadInfo && channel.loadInfo.originAttributes;
+        const browserContext = originAttributes ? this._targetRegistry.browserContextForUserContextId(originAttributes.userContextId) : null;
+        const proxy = browserContext ? browserContext.proxy : null;
+        if (!proxy) {
+          proxyFilter.onProxyFilterResult(defaultProxyInfo);
+          return;
+        }
+        proxyFilter.onProxyFilterResult(protocolProxyService.newProxyInfo(
+            proxy.type,
+            proxy.host,
+            proxy.port,
+            '', /* aProxyAuthorizationHeader */
+            '', /* aConnectionIsolationKey */
+            0, /* aFlags */
+            UINT32_MAX, /* aFailoverTimeout */
+            null, /* failover proxy */
+        ));
+      },
+    };
+    protocolProxyService.registerChannelFilter(this._channelProxyFilter, 0 /* position */);
+
     this._channelSink = {
       QueryInterface: ChromeUtils.generateQI([Ci.nsIChannelEventSink]),
       asyncOnChannelRedirect: (oldChannel, newChannel, flags, callback) => {

diff --git a/browser_patches/firefox/juggler/TargetRegistry.js b/browser_patches/firefox/juggler/TargetRegistry.js
@@ -252,6 +252,10 @@ class TargetRegistry {
     return this._browserContextIdToBrowserContext.get(browserContextId);
   }
 
+  browserContextForUserContextId(userContextId) {
+    return this._userContextIdToBrowserContext.get(userContextId);
+  }
+
   async newPage({browserContextId}) {
     let window;
     let created = false;
@@ -458,6 +462,7 @@ class BrowserContext {
     this._registry._browserContextIdToBrowserContext.set(this.browserContextId, this);
     this._registry._userContextIdToBrowserContext.set(this.userContextId, this);
     this.removeOnDetach = removeOnDetach;
+    this.proxy = undefined;
     this.extraHTTPHeaders = undefined;
     this.httpCredentials = undefined;
     this.requestInterceptionEnabled = undefined;

diff --git a/browser_patches/firefox/juggler/protocol/BrowserHandler.js b/browser_patches/firefox/juggler/protocol/BrowserHandler.js
@@ -151,6 +151,10 @@ class BrowserHandler {
     this._targetRegistry.browserContextForId(browserContextId).httpCredentials = nullToUndefined(credentials);
   }
 
+  async setProxy({browserContextId, type, host, port}) {
+    this._targetRegistry.browserContextForId(browserContextId).proxy = { type, host, port };
+  }
+
   setRequestInterception({browserContextId, enabled}) {
     this._targetRegistry.browserContextForId(browserContextId).requestInterceptionEnabled = enabled;
   }

diff --git a/browser_patches/firefox/juggler/protocol/Protocol.js b/browser_patches/firefox/juggler/protocol/Protocol.js
@@ -259,6 +259,14 @@ const Browser = {
         headers: t.Array(networkTypes.HTTPHeader),
       },
     },
+    'setProxy': {
+      params: {
+        browserContextId: t.Optional(t.String),
+        type: t.Enum(['http', 'https', 'socks', 'socks4']),
+        host: t.String,
+        port: t.Number,
+      },
+    },
     'setHTTPCredentials': {
       params: {
         browserContextId: t.Optional(t.String),