headless: Implement screenshot capturing

headless/app/headless_shell.cc

 // Copyright 2016 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.
 
 #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"
 #include "base/numerics/safe_conversions.h"
 #include "base/strings/string_number_conversions.h"
 #include "content/public/common/content_switches.h"
 #include "headless/app/headless_shell_switches.h"
 #include "headless/public/domains/page.h"
 #include "headless/public/domains/runtime.h"
 #include "headless/public/headless_browser.h"
 #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;
 using headless::HeadlessDevToolsClient;
 using headless::HeadlessWebContents;
 namespace page = headless::page;
 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.
 class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
  public:
   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) {
     browser_ = browser;
 
     base::CommandLine::StringVector args =
         base::CommandLine::ForCurrentProcess()->GetArgs();
 
     const char kDefaultUrl[] = "about:blank";
     if (args.empty() || args[0].empty()) {
@@ skipping 73 matching lines @@
 
     if (base::CommandLine::ForCurrentProcess()->HasSwitch(
             headless::switches::kDumpDom)) {
       FetchDom();
     } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
                    headless::switches::kRepl)) {
       std::cout
           << "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();
     }
   }
 
   void FetchDom() {
     devtools_client_->GetRuntime()->Evaluate(
         "document.body.innerHTML",
         base::Bind(&HeadlessShell::OnDomFetched, base::Unretained(this)));
   }
@@ skipping 26 matching lines @@
   }
 
   void OnExpressionResult(std::unique_ptr<runtime::EvaluateResult> result) {
     std::unique_ptr<base::Value> value = result->Serialize();
     std::string result_json;
     base::JSONWriter::Write(*value, &result_json);
     std::cout << result_json << std::endl;
     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();
     return command_line.HasSwitch(switches::kRemoteDebuggingPort);
   }
 
  private:
   GURL url_;
   HeadlessBrowser* browser_;  // Not owned.
   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);
 };
 
 int main(int argc, const char** argv) {
   HeadlessShell shell;
   HeadlessBrowser::Options::Builder builder(argc, argv);
 
   // Enable devtools if requested.
   base::CommandLine command_line(argc, argv);
@@ skipping 37 matching lines @@
 
   if (command_line.HasSwitch(switches::kHostResolverRules)) {
     builder.SetHostResolverRules(
         command_line.GetSwitchValueASCII(switches::kHostResolverRules));
   }
 
   return HeadlessBrowserMain(
       builder.Build(),
       base::Bind(&HeadlessShell::OnStart, base::Unretained(&shell)));
 }