headless: Allow protocol handler customization

Allow headless clients to provide custom protocol handlers. This lets
the clients override how URL requests are fulfilled for different
schemes types.

Design doc: https://docs.google.com/document/d/1m4CCZGJWOJdHY5MVwQqY3sT-5kBoz7NqmGvIb6-l3fM/edit#bookmark=id.5bqp8w95liri

BUG=546953

Review-Url: https://codereview.chromium.org/2024973002
Cr-Commit-Position: refs/heads/master@{#397184}

headless/lib/browser/headless_url_request_context_getter.cc

@@ -72,6 +72,12 @@ HeadlessURLRequestContextGetter::HeadlessURLRequestContextGetter(
   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
@@ -181,17 +187,22 @@ HeadlessURLRequestContextGetter::GetURLRequestContext() {
     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.

headless/lib/headless_browser_browsertest.cc

@@ -4,16 +4,94 @@
 
 #include <memory>
 
+#include "base/memory/ptr_util.h"
 #include "base/strings/stringprintf.h"
 #include "content/public/test/browser_test.h"
+#include "headless/public/domains/types.h"
 #include "headless/public/headless_browser.h"
 #include "headless/public/headless_web_contents.h"
 #include "headless/test/headless_browser_test.h"
+#include "net/base/io_buffer.h"
+#include "net/http/http_response_headers.h"
 #include "net/test/spawned_test_server/spawned_test_server.h"
+#include "net/url_request/url_request_job.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "ui/gfx/geometry/size.h"
 
 namespace headless {
+namespace {
+
+class TestURLRequestJob : public net::URLRequestJob {
+ public:
+  TestURLRequestJob(net::URLRequest* request,
+                    net::NetworkDelegate* network_delegate,
+                    const std::string& body);
+  ~TestURLRequestJob() override {}
+
+  // net::URLRequestJob implementation:
+  void Start() override;
+  void GetResponseInfo(net::HttpResponseInfo* info) override;
+  int ReadRawData(net::IOBuffer* buf, int buf_size) override;
+
+ private:
+  scoped_refptr<net::StringIOBuffer> body_;
+  scoped_refptr<net::DrainableIOBuffer> src_buf_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestURLRequestJob);
+};
+
+TestURLRequestJob::TestURLRequestJob(net::URLRequest* request,
+                                     net::NetworkDelegate* network_delegate,
+                                     const std::string& body)
+    : net::URLRequestJob(request, network_delegate),
+      body_(new net::StringIOBuffer(body)),
+      src_buf_(new net::DrainableIOBuffer(body_.get(), body_->size())) {}
+
+void TestURLRequestJob::Start() {
+  NotifyHeadersComplete();
+}
+
+void TestURLRequestJob::GetResponseInfo(net::HttpResponseInfo* info) {
+  info->headers =
+      new net::HttpResponseHeaders("Content-Type: text/html\r\n\r\n");
+}
+
+int TestURLRequestJob::ReadRawData(net::IOBuffer* buf, int buf_size) {
+  scoped_refptr<net::DrainableIOBuffer> dest_buf(
+      new net::DrainableIOBuffer(buf, buf_size));
+  while (src_buf_->BytesRemaining() > 0 && dest_buf->BytesRemaining() > 0) {
+    *dest_buf->data() = *src_buf_->data();
+    src_buf_->DidConsume(1);
+    dest_buf->DidConsume(1);
+  }
+  return dest_buf->BytesConsumed();
+}
+
+class TestProtocolHandler : public net::URLRequestJobFactory::ProtocolHandler {
+ public:
+  TestProtocolHandler(const std::string& body);
+  ~TestProtocolHandler() override {}
+
+  net::URLRequestJob* MaybeCreateJob(
+      net::URLRequest* request,
+      net::NetworkDelegate* network_delegate) const override;
+
+ private:
+  std::string body_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestProtocolHandler);
+};
+
+TestProtocolHandler::TestProtocolHandler(const std::string& body)
+    : body_(body) {}
+
+net::URLRequestJob* TestProtocolHandler::MaybeCreateJob(
+    net::URLRequest* request,
+    net::NetworkDelegate* network_delegate) const {
+  return new TestURLRequestJob(request, network_delegate, body_);
+}
+
+}  // namespace
 
 IN_PROC_BROWSER_TEST_F(HeadlessBrowserTest, CreateAndDestroyWebContents) {
   HeadlessWebContents* web_contents =
@@ -94,4 +172,52 @@ IN_PROC_BROWSER_TEST_F(HeadlessBrowserTest, SetHostResolverRules) {
   EXPECT_TRUE(WaitForLoad(web_contents));
 }
 
+IN_PROC_BROWSER_TEST_F(HeadlessBrowserTest, HttpProtocolHandler) {
+  const std::string kResponseBody = "<p>HTTP response body</p>";
+  ProtocolHandlerMap protocol_handlers;
+  protocol_handlers[url::kHttpScheme] =
+      base::WrapUnique(new TestProtocolHandler(kResponseBody));
+
+  HeadlessBrowser::Options::Builder builder;
+  builder.SetProtocolHandlers(std::move(protocol_handlers));
+  SetBrowserOptions(builder.Build());
+
+  // Load a page which doesn't actually exist, but which is fetched by our
+  // custom protocol handler.
+  HeadlessWebContents* web_contents = browser()->CreateWebContents(
+      GURL("http://not-an-actual-domain.tld/hello.html"), gfx::Size(800, 600));
+  EXPECT_TRUE(WaitForLoad(web_contents));
+
+  std::string inner_html;
+  EXPECT_TRUE(EvaluateScript(web_contents, "document.body.innerHTML")
+                  ->GetResult()
+                  ->GetValue()
+                  ->GetAsString(&inner_html));
+  EXPECT_EQ(kResponseBody, inner_html);
+}
+
+IN_PROC_BROWSER_TEST_F(HeadlessBrowserTest, HttpsProtocolHandler) {
+  const std::string kResponseBody = "<p>HTTPS response body</p>";
+  ProtocolHandlerMap protocol_handlers;
+  protocol_handlers[url::kHttpsScheme] =
+      base::WrapUnique(new TestProtocolHandler(kResponseBody));
+
+  HeadlessBrowser::Options::Builder builder;
+  builder.SetProtocolHandlers(std::move(protocol_handlers));
+  SetBrowserOptions(builder.Build());
+
+  // Load a page which doesn't actually exist, but which is fetched by our
+  // custom protocol handler.
+  HeadlessWebContents* web_contents = browser()->CreateWebContents(
+      GURL("https://not-an-actual-domain.tld/hello.html"), gfx::Size(800, 600));
+  EXPECT_TRUE(WaitForLoad(web_contents));
+
+  std::string inner_html;
+  EXPECT_TRUE(EvaluateScript(web_contents, "document.body.innerHTML")
+                  ->GetResult()
+                  ->GetValue()
+                  ->GetAsString(&inner_html));
+  EXPECT_EQ(kResponseBody, inner_html);
+}
+
 }  // namespace headless

headless/public/headless_browser.cc

@@ -65,6 +65,11 @@ Builder& Builder::SetSingleProcessMode(bool single_process_mode) {
   return *this;
 }
 
+Builder& Builder::SetProtocolHandlers(ProtocolHandlerMap protocol_handlers) {
+  options_.protocol_handlers = std::move(protocol_handlers);
+  return *this;
+}
+
 Options Builder::Build() {
   return std::move(options_);
 }

headless/public/headless_browser.h

@@ -7,6 +7,7 @@
 
 #include <memory>
 #include <string>
+#include <unordered_map>
 
 #include "base/callback.h"
 #include "base/macros.h"
@@ -14,6 +15,7 @@
 #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;
@@ -61,6 +63,10 @@ class HEADLESS_EXPORT HeadlessBrowser {
   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;
@@ -97,6 +103,10 @@ struct HeadlessBrowser::Options {
   // 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);
 
@@ -115,6 +125,7 @@ class HeadlessBrowser::Options::Builder {
   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();
 

headless/test/headless_browser_test.cc

@@ -15,6 +15,7 @@
 #include "headless/lib/headless_content_main_delegate.h"
 #include "headless/public/domains/network.h"
 #include "headless/public/domains/page.h"
+#include "headless/public/domains/runtime.h"
 #include "headless/public/headless_devtools_client.h"
 #include "headless/public/headless_devtools_target.h"
 #include "headless/public/headless_web_contents.h"
@@ -67,6 +68,43 @@ class WaitForLoadObserver : public page::Observer, public network::Observer {
   DISALLOW_COPY_AND_ASSIGN(WaitForLoadObserver);
 };
 
+class EvaluateHelper {
+ public:
+  EvaluateHelper(HeadlessBrowserTest* browser_test,
+                 HeadlessWebContents* web_contents,
+                 const std::string& script_to_eval)
+      : browser_test_(browser_test),
+        web_contents_(web_contents),
+        devtools_client_(HeadlessDevToolsClient::Create()) {
+    web_contents_->GetDevToolsTarget()->AttachClient(devtools_client_.get());
+    devtools_client_->GetRuntime()->Evaluate(
+        script_to_eval,
+        base::Bind(&EvaluateHelper::OnEvaluateResult, base::Unretained(this)));
+  }
+
+  ~EvaluateHelper() {
+    web_contents_->GetDevToolsTarget()->DetachClient(devtools_client_.get());
+  }
+
+  void OnEvaluateResult(std::unique_ptr<runtime::EvaluateResult> result) {
+    result_ = std::move(result);
+    browser_test_->FinishAsynchronousTest();
+  }
+
+  std::unique_ptr<runtime::EvaluateResult> TakeResult() {
+    return std::move(result_);
+  }
+
+ private:
+  HeadlessBrowserTest* browser_test_;  // Not owned.
+  HeadlessWebContents* web_contents_;  // Not owned.
+  std::unique_ptr<HeadlessDevToolsClient> devtools_client_;
+
+  std::unique_ptr<runtime::EvaluateResult> result_;
+
+  DISALLOW_COPY_AND_ASSIGN(EvaluateHelper);
+};
+
 }  // namespace
 
 HeadlessBrowserTest::HeadlessBrowserTest() {
@@ -114,6 +152,14 @@ bool HeadlessBrowserTest::WaitForLoad(HeadlessWebContents* web_contents) {
   return observer.navigation_succeeded();
 }
 
+std::unique_ptr<runtime::EvaluateResult> HeadlessBrowserTest::EvaluateScript(
+    HeadlessWebContents* web_contents,
+    const std::string& script) {
+  EvaluateHelper helper(this, web_contents, script);
+  RunAsynchronousTest();
+  return helper.TakeResult();
+}
+
 void HeadlessBrowserTest::RunAsynchronousTest() {
   base::MessageLoop::ScopedNestableTaskAllower nestable_allower(
       base::MessageLoop::current());

headless/test/headless_browser_test.h

@@ -14,6 +14,9 @@ class RunLoop;
 }
 
 namespace headless {
+namespace runtime {
+class EvaluateResult;
+}
 class HeadlessWebContents;
 
 // Base class for tests which require a full instance of the headless browser.
@@ -44,6 +47,11 @@ class HeadlessBrowserTest : public content::BrowserTestBase {
   // Synchronously waits for a tab to finish loading.
   bool WaitForLoad(HeadlessWebContents* web_contents);
 
+  // Synchronously evaluates a script and returns the result.
+  std::unique_ptr<runtime::EvaluateResult> EvaluateScript(
+      HeadlessWebContents* web_contents,
+      const std::string& script);
+
  protected:
   // Returns the browser for the test.
   HeadlessBrowser* browser() const;