headless: Allow per-context protocol handlers

headless/public/headless_browser_context.h

Index: headless/public/headless_browser_context.h
diff --git a/headless/public/headless_browser_context.h b/headless/public/headless_browser_context.h
index b72abb3c3ec669722981d974d6eeb34a47e17f4f..f8c581c2fd5fb0a17cfce226cca63e9f380cdaff 100644
--- a/headless/public/headless_browser_context.h
+++ b/headless/public/headless_browser_context.h
@@ -6,10 +6,15 @@
 #define HEADLESS_PUBLIC_HEADLESS_BROWSER_CONTEXT_H_
 
 #include "headless/public/headless_export.h"
+#include "net/url_request/url_request_job_factory.h"
 
 namespace headless {
 class HeadlessBrowserImpl;
 
+using ProtocolHandlerMap = std::unordered_map<
+    std::string,
+    std::unique_ptr<net::URLRequestJobFactory::ProtocolHandler>>;
+
 // Represents an isolated session with a unique cache, cookies, and other
 // profile/session related data.
 class HEADLESS_EXPORT HeadlessBrowserContext {
@@ -32,7 +37,9 @@ class HEADLESS_EXPORT HeadlessBrowserContext::Builder {
   Builder(Builder&&);
   ~Builder();
 
-  // TODO(skyostil): Allow overriding protocol handlers.
+  // Set custom network protocol handlers. These can be used to override URL
+  // fetching for different network schemes.
+  Builder& SetProtocolHandlers(ProtocolHandlerMap protocol_handlers);
 
   std::unique_ptr<HeadlessBrowserContext> Build();
 
@@ -42,6 +49,7 @@ class HEADLESS_EXPORT HeadlessBrowserContext::Builder {
   explicit Builder(HeadlessBrowserImpl* browser);
 
   HeadlessBrowserImpl* browser_;
+  ProtocolHandlerMap protocol_handlers_;
 
   DISALLOW_COPY_AND_ASSIGN(Builder);
 };