headless: Allow protocol handler customization

headless/public/headless_browser.h

 // 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.
 
 #ifndef HEADLESS_PUBLIC_HEADLESS_BROWSER_H_
 #define HEADLESS_PUBLIC_HEADLESS_BROWSER_H_
 
 #include <memory>
 #include <string>
+#include <unordered_map>
 
 #include "base/callback.h"
 #include "base/macros.h"
 #include "base/memory/ref_counted.h"
 #include "headless/public/headless_export.h"
 #include "net/base/host_port_pair.h"
 #include "net/base/ip_endpoint.h"
+#include "net/url_request/url_request_job_factory.h"
 
 namespace base {
 class MessagePump;
 class SingleThreadTaskRunner;
 }
 
 namespace gfx {
 class Size;
 }
 
@@ skipping 27 matching lines @@
   virtual void Shutdown() = 0;
 
  protected:
   HeadlessBrowser() {}
   virtual ~HeadlessBrowser() {}
 
  private:
   DISALLOW_COPY_AND_ASSIGN(HeadlessBrowser);
 };
 
+using ProtocolHandlerMap = std::unordered_map<
+    std::string,
+    std::unique_ptr<net::URLRequestJobFactory::ProtocolHandler>>;
+
 // Embedding API overrides for the headless browser.
 struct HeadlessBrowser::Options {
   class Builder;
 
   Options(Options&& options);
   ~Options();
 
   Options& operator=(Options&& options);
 
   // Command line options to be passed to browser.
@@ skipping 16 matching lines @@
 
   // Comma-separated list of rules that control how hostnames are mapped. See
   // chrome::switches::kHostRules for a description for the format.
   std::string host_resolver_rules;
 
   // Run the browser in single process mode instead of using separate renderer
   // processes as per default. Note that this also disables any sandboxing of
   // web content, which can be a security risk.
   bool single_process_mode;
 
+  // Custom network protocol handlers. These can be used to override URL
+  // fetching for different network schemes.
+  ProtocolHandlerMap protocol_handlers;
+
  private:
   Options(int argc, const char** argv);
 
   DISALLOW_COPY_AND_ASSIGN(Options);
 };
 
 class HeadlessBrowser::Options::Builder {
  public:
   Builder(int argc, const char** argv);
   Builder();
   ~Builder();
 
   Builder& SetUserAgent(const std::string& user_agent);
   Builder& EnableDevToolsServer(const net::IPEndPoint& endpoint);
   Builder& SetMessagePump(base::MessagePump* message_pump);
   Builder& SetProxyServer(const net::HostPortPair& proxy_server);
   Builder& SetHostResolverRules(const std::string& host_resolver_rules);
   Builder& SetSingleProcessMode(bool single_process_mode);
+  Builder& SetProtocolHandlers(ProtocolHandlerMap protocol_handlers);
 
   Options Build();
 
  private:
   Options options_;
 
   DISALLOW_COPY_AND_ASSIGN(Builder);
 };
 
 // Main entry point for running the headless browser. This function constructs
 // the headless browser instance, passing it to the given
 // |on_browser_start_callback| callback. Note that since this function executes
 // the main loop, it will only return after HeadlessBrowser::Shutdown() is
 // called, returning the exit code for the process.
 int HeadlessBrowserMain(
     HeadlessBrowser::Options options,
     const base::Callback<void(HeadlessBrowser*)>& on_browser_start_callback);
 
 }  // namespace headless
 
 #endif  // HEADLESS_PUBLIC_HEADLESS_BROWSER_H_