headless: Implement screenshot capturing

(based on patch from skyostil@, also sets default window size to 800x600 to enable basic snapshot support)

With the --screenshot option, headless shell will save a PNG screenshot
of the loaded page.

patch from issue 2000723002 at patchset 1 (http://crrev.com/2000723002#ps1)

BUG=546953

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

headless/app/headless_shell.cc

@@ -4,10 +4,13 @@
 
 #include <iostream>
 #include <memory>
+#include <string>
 
+#include "base/base64.h"
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/command_line.h"
+#include "base/files/file_path.h"
 #include "base/json/json_writer.h"
 #include "base/location.h"
 #include "base/memory/ref_counted.h"
@@ -21,7 +24,10 @@
 #include "headless/public/headless_devtools_client.h"
 #include "headless/public/headless_devtools_target.h"
 #include "headless/public/headless_web_contents.h"
+#include "net/base/file_stream.h"
+#include "net/base/io_buffer.h"
 #include "net/base/ip_address.h"
+#include "net/base/net_errors.h"
 #include "ui/gfx/geometry/size.h"
 
 using headless::HeadlessBrowser;
@@ -33,6 +39,8 @@ namespace runtime = headless::runtime;
 namespace {
 // Address where to listen to incoming DevTools connections.
 const char kDevToolsHttpServerAddress[] = "127.0.0.1";
+// Default file name for screenshot. Can be overriden by "--screenshot" switch.
+const char kDefaultScreenshotFileName[] = "screenshot.png";
 }
 
 // A sample application which demonstrates the use of the headless API.
@@ -41,6 +49,7 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
   HeadlessShell()
       : browser_(nullptr),
         devtools_client_(HeadlessDevToolsClient::Create()),
+        web_contents_(nullptr),
         processed_page_ready_(false) {}
   ~HeadlessShell() override {}
 
@@ -135,6 +144,9 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
           << "Type a Javascript expression to evaluate or \"quit\" to exit."
           << std::endl;
       InputExpression();
+    } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+                   headless::switches::kScreenshot)) {
+      CaptureScreenshot();
     } else {
       Shutdown();
     }
@@ -181,6 +193,83 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
     InputExpression();
   }
 
+  void CaptureScreenshot() {
+    devtools_client_->GetPage()->GetExperimental()->CaptureScreenshot(
+        page::CaptureScreenshotParams::Builder().Build(),
+        base::Bind(&HeadlessShell::OnScreenshotCaptured,
+                   base::Unretained(this)));
+  }
+
+  void OnScreenshotCaptured(
+      std::unique_ptr<page::CaptureScreenshotResult> result) {
+    base::FilePath file_name =
+        base::CommandLine::ForCurrentProcess()->GetSwitchValuePath(
+            headless::switches::kScreenshot);
+    if (file_name.empty()) {
+      file_name = base::FilePath().AppendASCII(kDefaultScreenshotFileName);
+    }
+
+    screenshot_file_stream_.reset(
+        new net::FileStream(browser_->BrowserFileThread()));
+    const int open_result = screenshot_file_stream_->Open(
+        file_name, base::File::FLAG_CREATE_ALWAYS | base::File::FLAG_WRITE |
+                       base::File::FLAG_ASYNC,
+        base::Bind(&HeadlessShell::OnScreenshotFileOpened,
+                   base::Unretained(this), base::Passed(std::move(result)),
+                   file_name));
+    if (open_result != net::ERR_IO_PENDING) {
+      // Operation could not be started.
+      OnScreenshotFileOpened(nullptr, file_name, open_result);
+    }
+  }
+
+  void OnScreenshotFileOpened(
+      std::unique_ptr<page::CaptureScreenshotResult> result,
+      const base::FilePath file_name,
+      const int open_result) {
+    if (open_result != net::OK) {
+      LOG(ERROR) << "Writing screenshot to file " << file_name.value()
+                 << " was unsuccessful, could not open file: "
+                 << net::ErrorToString(open_result);
+      return;
+    }
+
+    std::string decoded_png;
+    base::Base64Decode(result->GetData(), &decoded_png);
+    scoped_refptr<net::IOBufferWithSize> buf =
+        new net::IOBufferWithSize(decoded_png.size());
+    memcpy(buf->data(), decoded_png.data(), decoded_png.size());
+    const int write_result = screenshot_file_stream_->Write(
+        buf.get(), buf->size(),
+        base::Bind(&HeadlessShell::OnScreenshotFileWritten,
+                   base::Unretained(this), file_name, buf->size()));
+    if (write_result != net::ERR_IO_PENDING) {
+      // Operation may have completed successfully or failed.
+      OnScreenshotFileWritten(file_name, buf->size(), write_result);
+    }
+  }
+
+  void OnScreenshotFileWritten(const base::FilePath file_name,
+                               const int length,
+                               const int write_result) {
+    if (write_result < length) {
+      // TODO(eseckler): Support recovering from partial writes.
+      LOG(ERROR) << "Writing screenshot to file " << file_name.value()
+                 << " was unsuccessful: " << net::ErrorToString(write_result);
+    } else {
+      std::cout << "Screenshot written to file " << file_name.value() << "."
+                << std::endl;
+    }
+    int close_result = screenshot_file_stream_->Close(base::Bind(
+        &HeadlessShell::OnScreenshotFileClosed, base::Unretained(this)));
+    if (close_result != net::ERR_IO_PENDING) {
+      // Operation could not be started.
+      OnScreenshotFileClosed(close_result);
+    }
+  }
+
+  void OnScreenshotFileClosed(const int close_result) { Shutdown(); }
+
   bool RemoteDebuggingEnabled() const {
     const base::CommandLine& command_line =
         *base::CommandLine::ForCurrentProcess();
@@ -193,6 +282,7 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
   std::unique_ptr<HeadlessDevToolsClient> devtools_client_;
   HeadlessWebContents* web_contents_;
   bool processed_page_ready_;
+  std::unique_ptr<net::FileStream> screenshot_file_stream_;
 
   DISALLOW_COPY_AND_ASSIGN(HeadlessShell);
 };

headless/app/headless_shell_switches.cc

@@ -24,5 +24,8 @@ const char kRemoteDebuggingAddress[] = "remote-debugging-address";
 // expressions.
 const char kRepl[] = "repl";
 
+// Save a screenshot of the loaded page.
+const char kScreenshot[] = "screenshot";
+
 }  // namespace switches
 }  // namespace headless

headless/app/headless_shell_switches.h

@@ -11,6 +11,7 @@ extern const char kDumpDom[];
 extern const char kProxyServer[];
 extern const char kRemoteDebuggingAddress[];
 extern const char kRepl[];
+extern const char kScreenshot[];
 }  // namespace switches
 }  // namespace headless
 

headless/lib/browser/headless_browser_impl.cc

@@ -4,6 +4,8 @@
 
 #include "headless/lib/browser/headless_browser_impl.h"
 
+#include <vector>
+
 #include "base/memory/ptr_util.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "content/public/app/content_main.h"
@@ -51,6 +53,12 @@ HeadlessBrowserImpl::BrowserMainThread() const {
       content::BrowserThread::UI);
 }
 
+scoped_refptr<base::SingleThreadTaskRunner>
+HeadlessBrowserImpl::BrowserFileThread() const {
+  return content::BrowserThread::GetMessageLoopProxyForThread(
+      content::BrowserThread::FILE);
+}
+
 void HeadlessBrowserImpl::Shutdown() {
   DCHECK(BrowserMainThread()->BelongsToCurrentThread());
   BrowserMainThread()->PostTask(FROM_HERE,
@@ -87,7 +95,10 @@ void HeadlessBrowserImpl::set_browser_main_parts(
 
 void HeadlessBrowserImpl::RunOnStartCallback() {
   DCHECK(aura::Env::GetInstance());
-  window_tree_host_.reset(aura::WindowTreeHost::Create(gfx::Rect()));
+  // TODO(eseckler): allow configuration of window (viewport) size by embedder.
+  const gfx::Size kDefaultSize(800, 600);
+  window_tree_host_.reset(
+      aura::WindowTreeHost::Create(gfx::Rect(kDefaultSize)));
   window_tree_host_->InitHost();
 
   window_tree_client_.reset(

headless/lib/browser/headless_browser_impl.h

@@ -9,6 +9,7 @@
 
 #include <memory>
 #include <unordered_map>
+#include <vector>
 
 #include "base/synchronization/lock.h"
 #include "headless/lib/browser/headless_web_contents_impl.h"
@@ -38,6 +39,8 @@ class HeadlessBrowserImpl : public HeadlessBrowser {
                                          const gfx::Size& size) override;
   scoped_refptr<base::SingleThreadTaskRunner> BrowserMainThread()
       const override;
+  scoped_refptr<base::SingleThreadTaskRunner> BrowserFileThread()
+      const override;
 
   void Shutdown() override;
 
@@ -73,6 +76,7 @@ class HeadlessBrowserImpl : public HeadlessBrowser {
   std::unordered_map<HeadlessWebContents*, std::unique_ptr<HeadlessWebContents>>
       web_contents_;
 
+ private:
   DISALLOW_COPY_AND_ASSIGN(HeadlessBrowserImpl);
 };
 

headless/lib/headless_devtools_client_browsertest.cc

@@ -24,6 +24,8 @@ namespace headless {
   class AsyncHeadlessBrowserTestNeedsSemicolon##TEST_FIXTURE_NAME {}
 
 // A test fixture which attaches a devtools client before starting the test.
+// TODO(eseckler): Use the reusable HeadlessAsyncDevTooledTest and macro from
+// headless_browser_test.h instead of this class.
 class HeadlessDevToolsClientTest : public HeadlessBrowserTest,
                                    public HeadlessWebContents::Observer {
  public:

headless/lib/headless_web_contents_browsertest.cc

@@ -3,9 +3,13 @@
 // found in the LICENSE file.
 
 #include <memory>
+#include <string>
+#include <vector>
 
 #include "content/public/test/browser_test.h"
+#include "headless/public/domains/page.h"
 #include "headless/public/headless_browser.h"
+#include "headless/public/headless_devtools_client.h"
 #include "headless/public/headless_web_contents.h"
 #include "headless/test/headless_browser_test.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -18,6 +22,7 @@ class HeadlessWebContentsTest : public HeadlessBrowserTest {};
 
 IN_PROC_BROWSER_TEST_F(HeadlessWebContentsTest, Navigation) {
   EXPECT_TRUE(embedded_test_server()->Start());
+
   HeadlessWebContents* web_contents = browser()->CreateWebContents(
       embedded_test_server()->GetURL("/hello.html"), gfx::Size(800, 600));
   EXPECT_TRUE(WaitForLoad(web_contents));
@@ -42,4 +47,23 @@ IN_PROC_BROWSER_TEST_F(HeadlessWebContentsTest, WindowOpen) {
   EXPECT_EQ(static_cast<size_t>(2), all_web_contents.size());
 }
 
+class HeadlessWebContentsScreenshotTest
+    : public HeadlessAsyncDevTooledBrowserTest {
+ public:
+  void RunDevTooledTest() override {
+    devtools_client_->GetPage()->GetExperimental()->CaptureScreenshot(
+        page::CaptureScreenshotParams::Builder().Build(),
+        base::Bind(&HeadlessWebContentsScreenshotTest::OnScreenshotCaptured,
+                   base::Unretained(this)));
+  }
+
+  void OnScreenshotCaptured(
+      std::unique_ptr<page::CaptureScreenshotResult> result) {
+    EXPECT_LT(0U, result->GetData().length());
+    FinishAsynchronousTest();
+  }
+};
+
+HEADLESS_ASYNC_DEVTOOLED_TEST_F(HeadlessWebContentsScreenshotTest);
+
 }  // namespace headless

headless/public/headless_browser.h

@@ -8,6 +8,7 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 #include "base/callback.h"
 #include "base/macros.h"
@@ -51,6 +52,10 @@ class HEADLESS_EXPORT HeadlessBrowser {
   virtual scoped_refptr<base::SingleThreadTaskRunner> BrowserMainThread()
       const = 0;
 
+  // Returns a task runner for submitting work to the browser file thread.
+  virtual scoped_refptr<base::SingleThreadTaskRunner> BrowserFileThread()
+      const = 0;
+
   // Requests browser to stop as soon as possible. |Run| will return as soon as
   // browser stops.
   virtual void Shutdown() = 0;

headless/test/headless_browser_test.cc

@@ -19,6 +19,8 @@
 #include "headless/public/headless_devtools_client.h"
 #include "headless/public/headless_devtools_target.h"
 #include "headless/public/headless_web_contents.h"
+#include "ui/gfx/geometry/size.h"
+#include "url/gurl.h"
 
 namespace headless {
 namespace {
@@ -173,4 +175,29 @@ void HeadlessBrowserTest::FinishAsynchronousTest() {
   run_loop_->Quit();
 }
 
+HeadlessAsyncDevTooledBrowserTest::HeadlessAsyncDevTooledBrowserTest()
+    : web_contents_(nullptr),
+      devtools_client_(HeadlessDevToolsClient::Create()) {}
+
+HeadlessAsyncDevTooledBrowserTest::~HeadlessAsyncDevTooledBrowserTest() {}
+
+void HeadlessAsyncDevTooledBrowserTest::DevToolsTargetReady() {
+  EXPECT_TRUE(web_contents_->GetDevToolsTarget());
+  web_contents_->GetDevToolsTarget()->AttachClient(devtools_client_.get());
+  RunDevTooledTest();
+}
+
+void HeadlessAsyncDevTooledBrowserTest::RunTest() {
+  web_contents_ =
+      browser()->CreateWebContents(GURL("about:blank"), gfx::Size(800, 600));
+  web_contents_->AddObserver(this);
+
+  RunAsynchronousTest();
+
+  web_contents_->GetDevToolsTarget()->DetachClient(devtools_client_.get());
+  web_contents_->RemoveObserver(this);
+  web_contents_->Close();
+  web_contents_ = nullptr;
+}
+
 }  // namespace headless

headless/test/headless_browser_test.h

@@ -6,8 +6,11 @@
 #define HEADLESS_TEST_HEADLESS_BROWSER_TEST_H_
 
 #include <memory>
+#include <string>
+
 #include "content/public/test/browser_test_base.h"
 #include "headless/public/headless_browser.h"
+#include "headless/public/headless_web_contents.h"
 
 namespace base {
 class RunLoop;
@@ -17,7 +20,7 @@ namespace headless {
 namespace runtime {
 class EvaluateResult;
 }
-class HeadlessWebContents;
+class HeadlessDevToolsClient;
 
 // Base class for tests which require a full instance of the headless browser.
 class HeadlessBrowserTest : public content::BrowserTestBase {
@@ -62,6 +65,34 @@ class HeadlessBrowserTest : public content::BrowserTestBase {
   DISALLOW_COPY_AND_ASSIGN(HeadlessBrowserTest);
 };
 
+#define HEADLESS_ASYNC_DEVTOOLED_TEST_F(TEST_FIXTURE_NAME)               \
+  IN_PROC_BROWSER_TEST_F(TEST_FIXTURE_NAME, RunAsyncTest) { RunTest(); } \
+  class AsyncHeadlessBrowserTestNeedsSemicolon##TEST_FIXTURE_NAME {}
+
+// Base class for tests that require access to a DevToolsClient. Subclasses
+// should override the RunDevTooledTest() method, which is called asynchronously
+// when the DevToolsClient is ready.
+class HeadlessAsyncDevTooledBrowserTest : public HeadlessBrowserTest,
+                                          public HeadlessWebContents::Observer {
+ public:
+  HeadlessAsyncDevTooledBrowserTest();
+  ~HeadlessAsyncDevTooledBrowserTest() override;
+
+  // HeadlessWebContentsObserver implementation:
+  void DevToolsTargetReady() override;
+
+  // Implemented by tests and used to send request(s) to DevTools. Subclasses
+  // need to ensure that FinishAsynchronousTest() is called after response(s)
+  // are processed (e.g. in a callback).
+  virtual void RunDevTooledTest() = 0;
+
+ protected:
+  void RunTest();
+
+  HeadlessWebContents* web_contents_;
+  std::unique_ptr<HeadlessDevToolsClient> devtools_client_;
+};
+
 }  // namespace headless
 
 #endif  // HEADLESS_TEST_HEADLESS_BROWSER_TEST_H_