headless: Allow protocol handler customization

headless/lib/browser/headless_url_request_context_getter.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "headless/lib/browser/headless_url_request_context_getter.h"
 
 #include <memory>
 
 #include "base/memory/ptr_util.h"
 #include "base/single_thread_task_runner.h"
@@ skipping 54 matching lines @@
       file_task_runner_(std::move(file_task_runner)),
       net_log_(net_log),
       user_agent_(options->user_agent),
       host_resolver_rules_(options->host_resolver_rules),
       proxy_server_(options->proxy_server),
       request_interceptors_(std::move(request_interceptors)) {
   // Must first be created on the UI thread.
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
   std::swap(protocol_handlers_, *protocol_handlers);
+  for (auto& pair : options->protocol_handlers) {
+    protocol_handlers_[pair.first] =
+        linked_ptr<net::URLRequestJobFactory::ProtocolHandler>(
+            pair.second.release());
+  }
+  options->protocol_handlers.clear();
 
   // We must create the proxy config service on the UI loop on Linux because it
   // must synchronously run on the glib message loop. This will be passed to
   // the URLRequestContextStorage on the IO thread in GetURLRequestContext().
   if (proxy_server_.IsEmpty())
     proxy_config_service_ = GetProxyConfigService();
 }
 
 HeadlessURLRequestContextGetter::~HeadlessURLRequestContextGetter() {}
 
@@ skipping 13 matching lines @@
   if (!proxy_server_.IsEmpty())
     return net::ProxyService::CreateFixed(proxy_server_.ToString());
   return net::ProxyService::CreateUsingSystemProxyResolver(
       std::move(proxy_config_service_), 0, url_request_context_->net_log());
 }
 
 net::URLRequestContext*
 HeadlessURLRequestContextGetter::GetURLRequestContext() {
   DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
 
   if (!url_request_context_) {

[mmenke, 2016/06/01 15:32:19]

I'd recommend you switch to using URLRequestContextBuilder - it's updated as the
magic incantation to set up a URLRequestContext changes.  It also avoids
duplicating a lot of code.  If it doesn't support setting something you need to
set on the URLREquestContext, it's fine to add it.  Beyond the scope of this CL,
of course.

[Sami, 2016/06/01 15:44:51]

Thanks for the tip; filed
https://bugs.chromium.org/p/chromium/issues/detail?id=616475.

     url_request_context_.reset(new net::URLRequestContext());
     url_request_context_->set_net_log(net_log_);
     network_delegate_ = CreateNetworkDelegate();
     url_request_context_->set_network_delegate(network_delegate_.get());
     storage_.reset(
         new net::URLRequestContextStorage(url_request_context_.get()));
     storage_->set_cookie_store(
         content::CreateCookieStore(content::CookieStoreConfig()));
     storage_->set_channel_id_service(base::WrapUnique(
         new net::ChannelIDService(new net::DefaultChannelIDStore(nullptr),
@@ skipping 55 matching lines @@
 
     storage_->set_http_network_session(
         base::WrapUnique(new net::HttpNetworkSession(network_session_params)));
     storage_->set_http_transaction_factory(base::WrapUnique(new net::HttpCache(
         storage_->http_network_session(), std::move(main_backend),
         true /* set_up_quic_server_info */)));
 
     std::unique_ptr<net::URLRequestJobFactoryImpl> job_factory(
         new net::URLRequestJobFactoryImpl());
 
+    // Install handlers for default protocols which aren't handled by the
+    // network layer.
+    if (protocol_handlers_.find(url::kDataScheme) == protocol_handlers_.end()) {
+      protocol_handlers_[url::kDataScheme] =
+          linked_ptr<net::URLRequestJobFactory::ProtocolHandler>(
+              new net::DataProtocolHandler());
+    }
+    if (protocol_handlers_.find(url::kFileScheme) == protocol_handlers_.end()) {
+      protocol_handlers_[url::kFileScheme] =
+          linked_ptr<net::URLRequestJobFactory::ProtocolHandler>(
+              new net::FileProtocolHandler(
+                  content::BrowserThread::GetBlockingPool()
+                      ->GetTaskRunnerWithShutdownBehavior(
+                          base::SequencedWorkerPool::SKIP_ON_SHUTDOWN)));
+    }
     InstallProtocolHandlers(job_factory.get(), &protocol_handlers_);
-    bool set_protocol = job_factory->SetProtocolHandler(
-        url::kDataScheme, base::WrapUnique(new net::DataProtocolHandler));
-    DCHECK(set_protocol);
-    set_protocol = job_factory->SetProtocolHandler(
-        url::kFileScheme,
-        base::WrapUnique(new net::FileProtocolHandler(
-            content::BrowserThread::GetBlockingPool()
-                ->GetTaskRunnerWithShutdownBehavior(
-                    base::SequencedWorkerPool::SKIP_ON_SHUTDOWN))));
-    DCHECK(set_protocol);
 
     // Set up interceptors in the reverse order so that the last inceptor is at
     // the end of the linked list of job factories.
     std::unique_ptr<net::URLRequestJobFactory> top_job_factory =
         std::move(job_factory);
     for (auto i = request_interceptors_.rbegin();
          i != request_interceptors_.rend(); ++i) {
       top_job_factory.reset(new net::URLRequestInterceptingJobFactory(
           std::move(top_job_factory), base::WrapUnique(*i)));
     }
     request_interceptors_.weak_clear();
     // Save the head of the job factory list at storage_.
     storage_->set_job_factory(std::move(top_job_factory));
   }
 
   return url_request_context_.get();
 }
 
 scoped_refptr<base::SingleThreadTaskRunner>
 HeadlessURLRequestContextGetter::GetNetworkTaskRunner() const {
   return content::BrowserThread::GetMessageLoopProxyForThread(
       content::BrowserThread::IO);
 }
 
 net::HostResolver* HeadlessURLRequestContextGetter::host_resolver() const {
   return url_request_context_->host_resolver();
 }
 
 }  // namespace headless