move protocol to session properties for working with partitions

atom/browser/api/atom_api_protocol.cc

@@ -173,17 +173,10 @@ void RegisterStandardSchemes(
                                   base::JoinString(schemes, ","));
 }
 
-mate::Handle<atom::api::Protocol> CreateProtocol(v8::Isolate* isolate) {
-  auto browser_context = static_cast<atom::AtomBrowserContext*>(
-      atom::AtomBrowserMainParts::Get()->browser_context());
-  return atom::api::Protocol::Create(isolate, browser_context);
-}
-
 void Initialize(v8::Local<v8::Object> exports, v8::Local<v8::Value> unused,
                 v8::Local<v8::Context> context, void* priv) {
   v8::Isolate* isolate = context->GetIsolate();
   mate::Dictionary dict(isolate, exports);
-  dict.SetMethod("createProtocolObject", base::Bind(&CreateProtocol, isolate));
   dict.SetMethod("registerStandardSchemes", &RegisterStandardSchemes);
 }
 

atom/browser/api/atom_api_protocol.h

@@ -9,6 +9,7 @@
 #include <map>
 #include <vector>
 
+#include "atom/browser/api/trackable_object.h"
 #include "atom/browser/net/atom_url_request_job_factory.h"
 #include "base/callback.h"
 #include "base/containers/scoped_ptr_hash_map.h"
@@ -16,7 +17,6 @@
 #include "native_mate/arguments.h"
 #include "native_mate/dictionary.h"
 #include "native_mate/handle.h"
-#include "native_mate/wrappable.h"
 
 namespace base {
 class DictionaryValue;
@@ -34,7 +34,7 @@ class AtomURLRequestJobFactory;
 
 namespace api {
 
-class Protocol : public mate::Wrappable<Protocol> {
+class Protocol : public mate::TrackableObject<Protocol> {
  public:
   using Handler =
       base::Callback<void(const base::DictionaryValue&, v8::Local<v8::Value>)>;

atom/browser/api/atom_api_session.cc

@@ -9,6 +9,7 @@
 
 #include "atom/browser/api/atom_api_cookies.h"
 #include "atom/browser/api/atom_api_download_item.h"
+#include "atom/browser/api/atom_api_protocol.h"
 #include "atom/browser/api/atom_api_web_request.h"
 #include "atom/browser/atom_browser_context.h"
 #include "atom/browser/atom_browser_main_parts.h"
@@ -462,6 +463,14 @@ v8::Local<v8::Value> Session::Cookies(v8::Isolate* isolate) {
   return v8::Local<v8::Value>::New(isolate, cookies_);
 }
 
+v8::Local<v8::Value> Session::Protocol(v8::Isolate* isolate) {
+  if (protocol_.IsEmpty()) {
+    auto handle = atom::api::Protocol::Create(isolate, browser_context());
+    protocol_.Reset(isolate, handle.ToV8());
+  }
+  return v8::Local<v8::Value>::New(isolate, protocol_);
+}
+
 v8::Local<v8::Value> Session::WebRequest(v8::Isolate* isolate) {
   if (web_request_.IsEmpty()) {
     auto handle = atom::api::WebRequest::Create(isolate, browser_context());
@@ -512,6 +521,7 @@ void Session::BuildPrototype(v8::Isolate* isolate,
       .SetMethod("allowNTLMCredentialsForDomains",
                  &Session::AllowNTLMCredentialsForDomains)
       .SetProperty("cookies", &Session::Cookies)
+      .SetProperty("protocol", &Session::Protocol)
       .SetProperty("webRequest", &Session::WebRequest);
 }
 

atom/browser/api/atom_api_session.h

@@ -81,10 +81,12 @@ class Session: public mate::TrackableObject<Session>,
   void ClearHostResolverCache(mate::Arguments* args);
   void AllowNTLMCredentialsForDomains(const std::string& domains);
   v8::Local<v8::Value> Cookies(v8::Isolate* isolate);
+  v8::Local<v8::Value> Protocol(v8::Isolate* isolate);
   v8::Local<v8::Value> WebRequest(v8::Isolate* isolate);
 
   // Cached object.
   v8::Global<v8::Value> cookies_;
+  v8::Global<v8::Value> protocol_;
   v8::Global<v8::Value> web_request_;
 
   // The X-DevTools-Emulate-Network-Conditions-Client-Id.

lib/browser/api/protocol.js

@@ -1,5 +1,5 @@
-const {app} = require('electron')
-const {createProtocolObject, registerStandardSchemes} = process.atomBinding('protocol')
+const {app, session} = require('electron')
+const {registerStandardSchemes} = process.atomBinding('protocol')
 
 exports.registerStandardSchemes = function (schemes) {
   if (app.isReady()) {
@@ -10,7 +10,7 @@ exports.registerStandardSchemes = function (schemes) {
 }
 
 app.once('ready', function () {
-  let protocol = createProtocolObject()
+  let protocol = session.defaultSession.protocol
   for (let method in protocol) {
     exports[method] = protocol[method].bind(protocol)
   }

spec/api-session-spec.js

@@ -16,8 +16,15 @@ describe('session module', function () {
   var fixtures = path.resolve(__dirname, 'fixtures')
   var w = null
   var url = 'http://127.0.0.1'
+  var partitionName = 'temp'
+  var protocolName = 'sp'
+  const tempProtocol = session.fromPartition(partitionName).protocol
+  const protocol = session.defaultSession.protocol
 
   beforeEach(function () {
+    if (w != null) {
+      w.destroy()
+    }
     w = new BrowserWindow({
       show: false,
       width: 400,
@@ -26,7 +33,10 @@ describe('session module', function () {
   })
 
   afterEach(function () {
-    w.destroy()
+    if (w != null) {
+      w.destroy()
+    }
+    w = null
   })
 
   describe('session.cookies', function () {
@@ -262,4 +272,50 @@ describe('session module', function () {
       })
     })
   })
+
+  describe('session.protocol', function () {
+    beforeEach(function () {
+      if (w != null) {
+        w.destroy()
+      }
+      w = new BrowserWindow({
+        show: false,
+        width: 400,
+        height: 400,
+        webPreferences: {
+          partition: partitionName
+        }
+      })
+    })
+
+    afterEach(function () {
+      if (w != null) {
+        w.destroy()
+      }
+      w = null
+    })
+
+    it('handles requests from a partition', function (done) {
+      var handler = function (error, callback) {
+        callback({
+          data: 'test'
+        })
+      }
+      tempProtocol.registerStringProtocol(protocolName, handler, function (error) {
+        if (error) {
+          return done(error)
+        }
+        protocol.isProtocolHandled(protocolName, function (result) {
+          assert.equal(result, false)
+          tempProtocol.isProtocolHandled(protocolName, function (result) {
+            assert.equal(result, true)
+            w.webContents.on('did-finish-load', function () {
+              done()
+            })
+            w.loadURL(protocolName + "://fake-host")
+          })
+        })
+      })
+    })
+  })
 })