From bc76f356911d4d51a707cd9ead14fc090dcd07d0 Mon Sep 17 00:00:00 2001
From: Thiago de Arruda <tpadilha84@gmail.com>
Date: Thu, 1 Feb 2018 08:26:05 -0300
Subject: [PATCH] Fix protocol filtering of net.request

net::URLRequest inherits from base::SupportsUserData, which allows
associating arbitrary data with the request. Use this mechanism as a
condition for filtering requests from custom protocols.

Close #11657
---
 atom/browser/api/atom_api_protocol.h             |  4 ----
 atom/browser/net/atom_url_request.cc             |  4 ++++
 atom/browser/net/atom_url_request_job_factory.cc | 12 ++++++++++++
 atom/browser/net/atom_url_request_job_factory.h  |  2 ++
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/atom/browser/api/atom_api_protocol.h b/atom/browser/api/atom_api_protocol.h
index 40dc30600f936..dfc32be6bc57c 100644
--- a/atom/browser/api/atom_api_protocol.h
+++ b/atom/browser/api/atom_api_protocol.h
@@ -78,10 +78,6 @@ class Protocol : public mate::TrackableObject<Protocol> {
     net::URLRequestJob* MaybeCreateJob(
         net::URLRequest* request,
         net::NetworkDelegate* network_delegate) const override {
-      if (!request->initiator().has_value()) {
-        // Don't intercept this request as it was created by `net.request`.
-        return nullptr;
-      }
       RequestJob* request_job = new RequestJob(request, network_delegate);
       request_job->SetHandlerInfo(isolate_, request_context_.get(), handler_);
       return request_job;
diff --git a/atom/browser/net/atom_url_request.cc b/atom/browser/net/atom_url_request.cc
index e75aba9137322..1f0bdae4c08aa 100644
--- a/atom/browser/net/atom_url_request.cc
+++ b/atom/browser/net/atom_url_request.cc
@@ -7,6 +7,7 @@
 #include <string>
 #include "atom/browser/api/atom_api_url_request.h"
 #include "atom/browser/atom_browser_context.h"
+#include "atom/browser/net/atom_url_request_job_factory.h"
 #include "base/callback.h"
 #include "content/public/browser/browser_thread.h"
 #include "net/base/elements_upload_data_stream.h"
@@ -121,6 +122,9 @@ void AtomURLRequest::DoInitialize(
   request_->set_method(method);
   // Do not send cookies from the cookie store.
   DoSetLoadFlags(net::LOAD_DO_NOT_SEND_COOKIES);
+  // Set a flag to stop custom protocol from intercepting this request.
+  request_->SetUserData(DisableProtocolInterceptFlagKey(),
+                        base::WrapUnique(new base::SupportsUserData::Data()));
 }
 
 void AtomURLRequest::DoTerminate() {
diff --git a/atom/browser/net/atom_url_request_job_factory.cc b/atom/browser/net/atom_url_request_job_factory.cc
index 79dd80b586600..20680adf7da82 100644
--- a/atom/browser/net/atom_url_request_job_factory.cc
+++ b/atom/browser/net/atom_url_request_job_factory.cc
@@ -15,8 +15,18 @@ using content::BrowserThread;
 
 namespace atom {
 
+namespace {
+
+int disable_protocol_intercept_flag_key = 0;
+
+}  // namespace
+
 typedef net::URLRequestJobFactory::ProtocolHandler ProtocolHandler;
 
+const void* DisableProtocolInterceptFlagKey() {
+  return &disable_protocol_intercept_flag_key;
+}
+
 AtomURLRequestJobFactory::AtomURLRequestJobFactory() {}
 
 AtomURLRequestJobFactory::~AtomURLRequestJobFactory() {
@@ -93,6 +103,8 @@ net::URLRequestJob* AtomURLRequestJobFactory::MaybeCreateJobWithProtocolHandler(
   auto it = protocol_handler_map_.find(scheme);
   if (it == protocol_handler_map_.end())
     return nullptr;
+  if (request->GetUserData(DisableProtocolInterceptFlagKey()))
+    return nullptr;
   return it->second->MaybeCreateJob(request, network_delegate);
 }
 
diff --git a/atom/browser/net/atom_url_request_job_factory.h b/atom/browser/net/atom_url_request_job_factory.h
index 56f979c2ff60b..a560839b68e0c 100644
--- a/atom/browser/net/atom_url_request_job_factory.h
+++ b/atom/browser/net/atom_url_request_job_factory.h
@@ -16,6 +16,8 @@
 
 namespace atom {
 
+const void* DisableProtocolInterceptFlagKey();
+
 class AtomURLRequestJobFactory : public net::URLRequestJobFactory {
  public:
   AtomURLRequestJobFactory();