Add protocol interceptor API.

browser/api/atom_api_protocol.cc

@@ -4,6 +4,7 @@
 
 #include "browser/api/atom_api_protocol.h"
 
+#include "base/stl_util.h"
 #include "browser/atom_browser_context.h"
 #include "browser/net/adapter_request_job.h"
 #include "browser/net/atom_url_request_job_factory.h"
@@ -54,7 +55,6 @@ v8::Handle<v8::Object> ConvertURLRequestToV8Object(
 
 // Get the job factory.
 AtomURLRequestJobFactory* GetRequestJobFactory() {
-  DCHECK(content::BrowserThread::CurrentlyOn(content::BrowserThread::IO));
   return static_cast<AtomURLRequestJobFactory*>(
       const_cast<net::URLRequestJobFactory*>(
           static_cast<content::BrowserContext*>(AtomBrowserContext::Get())->
@@ -128,12 +128,14 @@ class CustomProtocolRequestJob : public AdapterRequestJob {
     }
 
     // Try the default protocol handler if we have.
-    if (default_protocol_handler())
+    if (default_protocol_handler()) {
       content::BrowserThread::PostTask(
           content::BrowserThread::IO,
           FROM_HERE,
           base::Bind(&AdapterRequestJob::CreateJobFromProtocolHandlerAndStart,
                      GetWeakPtr()));
+      return;
+    }
 
     // Fallback to the not implemented error.
     content::BrowserThread::PostTask(
@@ -165,6 +167,12 @@ class CustomProtocolHandler : public ProtocolHandler {
                                         network_delegate);
   }
 
+  ProtocolHandler* ReleaseDefaultProtocolHandler() {
+    return protocol_handler_.release();
+  }
+
+  ProtocolHandler* original_handler() { return protocol_handler_.get(); }
+
  private:
   scoped_ptr<ProtocolHandler> protocol_handler_;
 
@@ -219,18 +227,38 @@ v8::Handle<v8::Value> Protocol::IsHandledProtocol(const v8::Arguments& args) {
 // static
 v8::Handle<v8::Value> Protocol::InterceptProtocol(const v8::Arguments& args) {
   std::string scheme(*v8::String::Utf8Value(args[0]));
-  if (scheme == "https" || scheme == "http")
-    return node::ThrowError("Intercepting http protocol is not supported.");
+  if (!GetRequestJobFactory()->HasProtocolHandler(scheme))
+    return node::ThrowError("Cannot intercept procotol");
+
+  if (ContainsKey(g_handlers, scheme))
+    return node::ThrowError("Cannot intercept custom procotols");
+
+  // Store the handler in a map.
+  if (!args[1]->IsFunction())
+    return node::ThrowError("Handler must be a function");
+  g_handlers[scheme] = v8::Persistent<v8::Function>::New(
+      node::node_isolate, v8::Handle<v8::Function>::Cast(args[1]));
 
   content::BrowserThread::PostTask(content::BrowserThread::IO,
                                    FROM_HERE,
-                                   base::Bind(&UnregisterProtocolInIO, scheme));
-
+                                   base::Bind(&InterceptProtocolInIO, scheme));
   return v8::Undefined();
 }
 
 // static
 v8::Handle<v8::Value> Protocol::UninterceptProtocol(const v8::Arguments& args) {
+  std::string scheme(*v8::String::Utf8Value(args[0]));
+
+  // Erase the handler from map.
+  HandlersMap::iterator it(g_handlers.find(scheme));
+  if (it == g_handlers.end())
+    return node::ThrowError("The scheme has not been registered");
+  g_handlers.erase(it);
+
+  content::BrowserThread::PostTask(content::BrowserThread::IO,
+                                   FROM_HERE,
+                                   base::Bind(&UninterceptProtocolInIO,
+                                              scheme));
   return v8::Undefined();
 }
 
@@ -263,11 +291,39 @@ void Protocol::UnregisterProtocolInIO(const std::string& scheme) {
 // static
 void Protocol::InterceptProtocolInIO(const std::string& scheme) {
   DCHECK(content::BrowserThread::CurrentlyOn(content::BrowserThread::IO));
+  AtomURLRequestJobFactory* job_factory(GetRequestJobFactory());
+  ProtocolHandler* original_handler = job_factory->GetProtocolHandler(scheme);
+  job_factory->ReplaceProtocol(scheme,
+                               new CustomProtocolHandler(original_handler));
+
+  content::BrowserThread::PostTask(content::BrowserThread::UI,
+                                   FROM_HERE,
+                                   base::Bind(&EmitEventInUI,
+                                              "intercepted",
+                                              scheme));
 }
 
 // static
 void Protocol::UninterceptProtocolInIO(const std::string& scheme) {
   DCHECK(content::BrowserThread::CurrentlyOn(content::BrowserThread::IO));
+  AtomURLRequestJobFactory* job_factory(GetRequestJobFactory());
+
+  // Check if the protocol handler is intercepted.
+  CustomProtocolHandler* handler = static_cast<CustomProtocolHandler*>(
+      job_factory->GetProtocolHandler(scheme));
+  if (!handler->original_handler())
+    return;
+
+  // Reset the protocol handler to the orignal one and delete current
+  // protocol handler.
+  ProtocolHandler* original_handler = handler->ReleaseDefaultProtocolHandler();
+  delete job_factory->ReplaceProtocol(scheme, original_handler);
+
+  content::BrowserThread::PostTask(content::BrowserThread::UI,
+                                   FROM_HERE,
+                                   base::Bind(&EmitEventInUI,
+                                              "unintercepted",
+                                              scheme));
 }
 
 // static
@@ -279,6 +335,8 @@ void Protocol::Initialize(v8::Handle<v8::Object> target) {
   node::SetMethod(target, "registerProtocol", RegisterProtocol);
   node::SetMethod(target, "unregisterProtocol", UnregisterProtocol);
   node::SetMethod(target, "isHandledProtocol", IsHandledProtocol);
+  node::SetMethod(target, "interceptProtocol", InterceptProtocol);
+  node::SetMethod(target, "uninterceptProtocol", UninterceptProtocol);
 }
 
 }  // namespace api

browser/net/adapter_request_job.cc

@@ -96,7 +96,7 @@ void AdapterRequestJob::CreateJobFromProtocolHandlerAndStart() {
   DCHECK(protocol_handler_);
   real_job_ = protocol_handler_->MaybeCreateJob(request(),
                                                 network_delegate());
-  if (!real_job_)
+  if (!real_job_.get())
     CreateErrorJobAndStart(net::ERR_NOT_IMPLEMENTED);
   else
     real_job_->Start();

browser/net/atom_url_request_job_factory.cc

@@ -25,6 +25,8 @@ bool AtomURLRequestJobFactory::SetProtocolHandler(
     ProtocolHandler* protocol_handler) {
   DCHECK(CalledOnValidThread());
 
+  base::AutoLock locked(lock_);
+
   if (!protocol_handler) {
     ProtocolHandlerMap::iterator it = protocol_handler_map_.find(scheme);
     if (it == protocol_handler_map_.end())
@@ -41,12 +43,13 @@ bool AtomURLRequestJobFactory::SetProtocolHandler(
   return true;
 }
 
-ProtocolHandler* AtomURLRequestJobFactory::InterceptProtocol(
+ProtocolHandler* AtomURLRequestJobFactory::ReplaceProtocol(
     const std::string& scheme,
     ProtocolHandler* protocol_handler) {
   DCHECK(CalledOnValidThread());
   DCHECK(protocol_handler);
 
+  base::AutoLock locked(lock_);
   if (!ContainsKey(protocol_handler_map_, scheme))
     return NULL;
   ProtocolHandler* original_protocol_handler = protocol_handler_map_[scheme];
@@ -54,11 +57,30 @@ ProtocolHandler* AtomURLRequestJobFactory::InterceptProtocol(
   return original_protocol_handler;
 }
 
+ProtocolHandler* AtomURLRequestJobFactory::GetProtocolHandler(
+    const std::string& scheme) const {
+  DCHECK(CalledOnValidThread());
+
+  base::AutoLock locked(lock_);
+  ProtocolHandlerMap::const_iterator it = protocol_handler_map_.find(scheme);
+  if (it == protocol_handler_map_.end())
+    return NULL;
+  return it->second;
+}
+
+bool AtomURLRequestJobFactory::HasProtocolHandler(
+    const std::string& scheme) const {
+  base::AutoLock locked(lock_);
+  return ContainsKey(protocol_handler_map_, scheme);
+}
+
 net::URLRequestJob* AtomURLRequestJobFactory::MaybeCreateJobWithProtocolHandler(
     const std::string& scheme,
     net::URLRequest* request,
     net::NetworkDelegate* network_delegate) const {
   DCHECK(CalledOnValidThread());
+
+  base::AutoLock locked(lock_);
   ProtocolHandlerMap::const_iterator it = protocol_handler_map_.find(scheme);
   if (it == protocol_handler_map_.end())
     return NULL;
@@ -68,7 +90,7 @@ net::URLRequestJob* AtomURLRequestJobFactory::MaybeCreateJobWithProtocolHandler(
 bool AtomURLRequestJobFactory::IsHandledProtocol(
     const std::string& scheme) const {
   DCHECK(CalledOnValidThread());
-  return ContainsKey(protocol_handler_map_, scheme) ||
+  return HasProtocolHandler(scheme) ||
       net::URLRequest::IsHandledProtocol(scheme);
 }
 

browser/net/atom_url_request_job_factory.h

@@ -11,6 +11,7 @@
 
 #include "base/basictypes.h"
 #include "base/compiler_specific.h"
+#include "base/synchronization/lock.h"
 #include "net/url_request/url_request_job_factory.h"
 
 namespace atom {
@@ -28,9 +29,15 @@ class AtomURLRequestJobFactory : public net::URLRequestJobFactory {
 
   // Intercepts the ProtocolHandler for a scheme. Returns the original protocol
   // handler on success, otherwise returns NULL.
-  ProtocolHandler* InterceptProtocol(const std::string& scheme,
+  ProtocolHandler* ReplaceProtocol(const std::string& scheme,
                                      ProtocolHandler* protocol_handler);
 
+  // Returns the protocol handler registered with scheme.
+  ProtocolHandler* GetProtocolHandler(const std::string& scheme) const;
+
+  // Whether the protocol handler is registered by the job factory.
+  bool HasProtocolHandler(const std::string& scheme) const;
+
   // URLRequestJobFactory implementation
   virtual net::URLRequestJob* MaybeCreateJobWithProtocolHandler(
       const std::string& scheme,
@@ -44,6 +51,8 @@ class AtomURLRequestJobFactory : public net::URLRequestJobFactory {
 
   ProtocolHandlerMap protocol_handler_map_;
 
+  mutable base::Lock lock_;
+
   DISALLOW_COPY_AND_ASSIGN(AtomURLRequestJobFactory);
 };
 

spec/api/protocol.coffee

@@ -80,3 +80,53 @@ describe 'protocol API', ->
       assert.equal protocol.isHandledProtocol('http'), true
       assert.equal protocol.isHandledProtocol('https'), true
       assert.equal protocol.isHandledProtocol('atom'), false
+
+  describe 'protocol.interceptProtocol', ->
+    it 'throws error when scheme is not a registered one', ->
+      register = -> protocol.interceptProtocol('test-intercept', ->)
+      assert.throws register, /Cannot intercept procotol/
+
+    it 'throws error when scheme is a custom protocol', (done) ->
+      protocol.once 'unregistered', (scheme) ->
+        assert.equal scheme, 'atom'
+        done()
+      protocol.once 'registered', (scheme) ->
+        assert.equal scheme, 'atom'
+        register = -> protocol.interceptProtocol('test-intercept', ->)
+        assert.throws register, /Cannot intercept procotol/
+        protocol.unregisterProtocol scheme
+      protocol.registerProtocol('atom', ->)
+
+    it 'returns original job when callback returns nothing', (done) ->
+      targetScheme = 'file'
+      protocol.once 'intercepted', (scheme) ->
+        assert.equal scheme, targetScheme
+        free = -> protocol.uninterceptProtocol targetScheme
+        $.ajax
+          url: "#{targetScheme}://#{__filename}",
+          success: ->
+            protocol.once 'unintercepted', (scheme) ->
+              assert.equal scheme, targetScheme
+              done()
+            free()
+          error: (xhr, errorType, error) ->
+            free()
+            assert false, 'Got error: ' + errorType + ' ' + error
+      protocol.interceptProtocol targetScheme, (request) ->
+        assert.equal request.url, "#{targetScheme}://#{__filename}"
+
+    it 'can override original protocol handler', (done) ->
+      handler = remote.createFunctionWithReturnValue 'valar morghulis'
+      protocol.once 'intercepted', ->
+        free = -> protocol.uninterceptProtocol 'file'
+        $.ajax
+          url: 'file://fake-host'
+          success: (data) ->
+            protocol.once 'unintercepted', ->
+              assert.equal data, handler()
+              done()
+            free()
+          error: (xhr, errorType, error) ->
+            assert false, 'Got error: ' + errorType + ' ' + error
+            free()
+      protocol.interceptProtocol 'file', handler