headless: Implement screenshot capturing

headless/app/headless_shell.cc

Index: headless/app/headless_shell.cc
diff --git a/headless/app/headless_shell.cc b/headless/app/headless_shell.cc
index 25c823d541cb196bba2579c004e3722c57d7da3a..38fbc60055b3c9b33306702a11d9b57ba2fcff4f 100644
--- a/headless/app/headless_shell.cc
+++ b/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,7 +49,9 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
   HeadlessShell()
       : browser_(nullptr),
         devtools_client_(HeadlessDevToolsClient::Create()),
-        processed_page_ready_(false) {}
+        web_contents_(nullptr),
+        processed_page_ready_(false),
+        screenshot_file_stream_(nullptr) {}

[Sami, 2016/06/03 11:46:24]

No need to initialize a unique_ptr. The web_contents_ fix is a good find though.

[Eric Seckler, 2016/06/03 13:40:38]

Ack, removed :)

   ~HeadlessShell() override {}
 
   void OnStart(HeadlessBrowser* browser) {
@@ -135,6 +145,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 +194,76 @@ 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) {
+      OnScreenshotFileOpened(std::move(result), file_name, open_result);

[Sami, 2016/06/03 11:46:24]

Does this work since you already moved |result| a few lines ago? Might need pass
in a nullptr instead.

[Eric Seckler, 2016/06/03 13:40:38]

Passing a nullptr now. From the comments in FileStream, it seems that Open() and
Close() can only execute asynchronously, so this should be fine. Interestingly,
Write() could execute synchronously (if we trust the comment). The FileStream
implementation seems to always return ERR_IO_PENDING anyway.

+    }
+  }
+
+  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);

[Sami, 2016/06/03 11:46:24]

nit: you could early-out here to reduce the indentation in the rest of the
function.

[Eric Seckler, 2016/06/03 13:40:38]

Done.

+    } else {
+      std::string decoded_png;
+      base::Base64Decode(result->GetData(), &decoded_png);
+      net::IOBufferWithSize* buf =

[Sami, 2016/06/03 11:46:24]

Who owns |buf|?

[Eric Seckler, 2016/06/03 13:40:38]

IOBuffers are refcounted, but passed as plain pointers to the functions (see
comment for IOBuffer class). So effectively we transfer ownership to the
FileStream. I've added a scoped_refptr to make this a bit clearer (that seems to
be the way that this is handled elsewhere in chromium).

+          new net::IOBufferWithSize(decoded_png.size());
+      memcpy(buf->data(), decoded_png.data(), decoded_png.size());
+      const int write_result = screenshot_file_stream_->Write(
+          buf, buf->size(), base::Bind(&HeadlessShell::OnScreenshotFileWritten,
+                                       base::Unretained(this), file_name));
+      if (write_result != net::ERR_IO_PENDING) {
+        OnScreenshotFileWritten(file_name, open_result);

[Sami, 2016/06/03 11:46:24]

write_result?

[Eric Seckler, 2016/06/03 13:40:38]

Good catch, done.

+      }
+    }
+  }
+
+  void OnScreenshotFileWritten(const base::FilePath file_name,
+                               const int write_result) {
+    if (write_result < 0) {
+      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) {
+      OnScreenshotFileClosed(close_result);
+    }
+  }
+
+  void OnScreenshotFileClosed(const int close_result) { Shutdown(); }
+
   bool RemoteDebuggingEnabled() const {
     const base::CommandLine& command_line =
         *base::CommandLine::ForCurrentProcess();
@@ -193,6 +276,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);
 };