From 2b8b76a5276edef7dd3ce22b3824fbec8e416cbe Mon Sep 17 00:00:00 2001
From: deepak1556 <hop2deep@gmail.com>
Date: Mon, 28 Nov 2016 00:37:50 +0530
Subject: [PATCH] protocol: support custom session with interceptHttpProtocol

---
 atom/browser/net/url_request_fetch_job.cc | 26 +++++++++++++-------
 spec/api-protocol-spec.js                 | 29 ++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/atom/browser/net/url_request_fetch_job.cc b/atom/browser/net/url_request_fetch_job.cc
index 5f11c67476..6cd2f447ea 100644
--- a/atom/browser/net/url_request_fetch_job.cc
+++ b/atom/browser/net/url_request_fetch_job.cc
@@ -7,9 +7,12 @@
 #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"
+#include "native_mate/handle.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_errors.h"
 #include "net/http/http_response_headers.h"
@@ -89,15 +92,20 @@ 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());
+  mate::Handle<api::Session> session;
+  if (options.Get("session", &session)) {
+    if (session.IsEmpty()) {
+      // 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 {
+      AtomBrowserContext* browser_context = session->browser_context();
+      url_request_context_getter_ =
+          browser_context->url_request_context_getter();
+    }
   }
 }
 
diff --git a/spec/api-protocol-spec.js b/spec/api-protocol-spec.js
index 51003ebf42..089c0b39d7 100644
--- a/spec/api-protocol-spec.js
+++ b/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 () {