headless: Allow per-context protocol handlers

headless/lib/browser/headless_browser_context_impl.cc

 // Copyright 2015 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_browser_context_impl.h"
 
 #include <memory>
 
 #include "base/memory/ptr_util.h"
 #include "base/path_service.h"
@@ skipping 49 matching lines @@
   CHECK(url_request_context_getter_);
   return url_request_context_getter_->GetURLRequestContext()->host_resolver();
 }
 
 net::URLRequestContext* HeadlessResourceContext::GetRequestContext() {
   CHECK(url_request_context_getter_);
   return url_request_context_getter_->GetURLRequestContext();
 }
 
 HeadlessBrowserContextImpl::HeadlessBrowserContextImpl(
+    ProtocolHandlerMap protocol_handlers,
     HeadlessBrowser::Options* options)
-    : resource_context_(new HeadlessResourceContext), options_(options) {
+    : protocol_handlers_(std::move(protocol_handlers)),
+      options_(options),
+      resource_context_(new HeadlessResourceContext) {
   InitWhileIOAllowed();
 }
 
 HeadlessBrowserContextImpl::~HeadlessBrowserContextImpl() {
   if (resource_context_) {
     content::BrowserThread::DeleteSoon(content::BrowserThread::IO, FROM_HERE,
                                        resource_context_.release());
   }
 }
 
@@ skipping 64 matching lines @@
 
 net::URLRequestContextGetter* HeadlessBrowserContextImpl::CreateRequestContext(
     content::ProtocolHandlerMap* protocol_handlers,
     content::URLRequestInterceptorScopedVector request_interceptors) {
   scoped_refptr<HeadlessURLRequestContextGetter> url_request_context_getter(
       new HeadlessURLRequestContextGetter(
           content::BrowserThread::GetMessageLoopProxyForThread(
               content::BrowserThread::IO),
           content::BrowserThread::GetMessageLoopProxyForThread(
               content::BrowserThread::FILE),
-          protocol_handlers, std::move(request_interceptors), options()));
+          protocol_handlers, std::move(protocol_handlers_),
+          std::move(request_interceptors), options()));
   resource_context_->set_url_request_context_getter(url_request_context_getter);
   return url_request_context_getter.get();
 }
 
 net::URLRequestContextGetter*
 HeadlessBrowserContextImpl::CreateRequestContextForStoragePartition(
     const base::FilePath& partition_path,
     bool in_memory,
     content::ProtocolHandlerMap* protocol_handlers,
     content::URLRequestInterceptorScopedVector request_interceptors) {
@@ skipping 17 matching lines @@
   options_ = options;
 }
 
 HeadlessBrowserContext::Builder::Builder(HeadlessBrowserImpl* browser)
     : browser_(browser) {}
 
 HeadlessBrowserContext::Builder::~Builder() = default;
 
 HeadlessBrowserContext::Builder::Builder(Builder&&) = default;
 
+HeadlessBrowserContext::Builder&
+HeadlessBrowserContext::Builder::SetProtocolHandlers(
+    ProtocolHandlerMap protocol_handlers) {
+  protocol_handlers_ = std::move(protocol_handlers);
+  return *this;
+}
+
 std::unique_ptr<HeadlessBrowserContext>
 HeadlessBrowserContext::Builder::Build() {
-  return base::WrapUnique(new HeadlessBrowserContextImpl(browser_->options()));
+  return base::WrapUnique(new HeadlessBrowserContextImpl(
+      std::move(protocol_handlers_), browser_->options()));
 }
 
 }  // namespace headless