headless: Allow protocol handler customization

headless/lib/headless_browser_browsertest.cc

Index: headless/lib/headless_browser_browsertest.cc
diff --git a/headless/lib/headless_browser_browsertest.cc b/headless/lib/headless_browser_browsertest.cc
index d6eb82b21e9b41905c6b7898eb95b0011e6ec9f7..b85b49e05bd448f0e14edf7cc237b752d918d02a 100644
--- a/headless/lib/headless_browser_browsertest.cc
+++ b/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