protocol: support custom session with interceptHttpProtocol

atom/browser/net/url_request_fetch_job.cc

@@ -7,6 +7,8 @@
 #include <algorithm>
 #include <string>
 
+#include "atom/browser/api/atom_api_session.h"
+#include "atom/browser/atom_browser_context.h"
 #include "base/memory/ptr_util.h"
 #include "base/strings/string_util.h"
 #include "native_mate/dictionary.h"
@@ -89,15 +91,23 @@ void URLRequestFetchJob::BeforeStartInUI(
     return;
 
   // When |session| is set to |null| we use a new request context for fetch job.
-  // TODO(zcbenz): Handle the case when it is not null.
-  v8::Local<v8::Value> session;
-  if (options.Get("session", &session) && session->IsNull()) {
-    // We have to create the URLRequestContextGetter on UI thread.
-    url_request_context_getter_ = new brightray::URLRequestContextGetter(
-        this, nullptr, nullptr, base::FilePath(), true,
-        BrowserThread::UnsafeGetMessageLoopForThread(BrowserThread::IO),
-        BrowserThread::UnsafeGetMessageLoopForThread(BrowserThread::FILE),
-        nullptr, content::URLRequestInterceptorScopedVector());
+  v8::Local<v8::Value> val;
+  if (options.Get("session", &val)) {
+    if (val->IsNull()) {
+      // We have to create the URLRequestContextGetter on UI thread.
+      url_request_context_getter_ = new brightray::URLRequestContextGetter(
+          this, nullptr, nullptr, base::FilePath(), true,
+          BrowserThread::UnsafeGetMessageLoopForThread(BrowserThread::IO),
+          BrowserThread::UnsafeGetMessageLoopForThread(BrowserThread::FILE),
+          nullptr, content::URLRequestInterceptorScopedVector());
+    } else {
+      mate::Handle<api::Session> session;
+      if (mate::ConvertFromV8(isolate, val, &session) && !session.IsEmpty()) {
+        AtomBrowserContext* browser_context = session->browser_context();
+        url_request_context_getter_ =
+            browser_context->url_request_context_getter();
+      }
+    }
   }
 }
 
@@ -237,21 +247,21 @@ int URLRequestFetchJob::GetResponseCode() const {
 }
 
 void URLRequestFetchJob::OnURLFetchComplete(const net::URLFetcher* source) {
-  if (!response_info_) {
-    // Since we notify header completion only after first write there will be
-    // no response object constructed for http respones with no content 204.
-    // We notify header completion here.
-    HeadersCompleted();
-    return;
-  }
-
   ClearPendingBuffer();
   ClearWriteBuffer();
 
-  if (fetcher_->GetStatus().is_success())
+  if (fetcher_->GetStatus().is_success()) {
+    if (!response_info_) {
+      // Since we notify header completion only after first write there will be
+      // no response object constructed for http respones with no content 204.
+      // We notify header completion here.
+      HeadersCompleted();
+      return;
+    }
     ReadRawDataComplete(0);
-  else
+  } else {
     NotifyStartError(fetcher_->GetStatus());
+  }
 }
 
 int URLRequestFetchJob::BufferCopy(net::IOBuffer* source, int num_bytes,

spec/api-protocol-spec.js

@@ -4,7 +4,7 @@ const path = require('path')
 const qs = require('querystring')
 const {closeWindow} = require('./window-helpers')
 const remote = require('electron').remote
-const {BrowserWindow, ipcMain, protocol, webContents} = remote
+const {BrowserWindow, ipcMain, protocol, session, webContents} = remote
 
 describe('protocol module', function () {
   var protocolName = 'sp'
@@ -870,6 +870,33 @@ describe('protocol module', function () {
         })
       })
     })
+
+    it('can use custom session', function (done) {
+      const customSession = session.fromPartition('custom-ses', {
+        cache: false
+      })
+      customSession.webRequest.onBeforeRequest(function (details, callback) {
+        assert.equal(details.url, 'http://fake-host/')
+        callback({cancel: true})
+      })
+      const handler = function (request, callback) {
+        callback({
+          url: request.url,
+          session: customSession
+        })
+      }
+      protocol.interceptHttpProtocol('http', handler, function (error) {
+        if (error) {
+          return done(error)
+        }
+        fetch('http://fake-host').then(function () {
+          done('request succeeded but it should not')
+        }).catch(function () {
+          customSession.webRequest.onBeforeRequest(null)
+          done()
+        })
+      })
+    })
   })
 
   describe('protocol.uninterceptProtocol', function () {