add print to pdf for headless

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 <memory>
 #include <sstream>
 #include <string>
+#include <utility>
 
 #include "base/base64.h"
 #include "base/base_switches.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/weak_ptr.h"
 #include "base/numerics/safe_conversions.h"
 #include "base/strings/string_number_conversions.h"
 #include "headless/app/headless_shell.h"
 #include "headless/app/headless_shell_switches.h"
 #include "headless/public/headless_devtools_target.h"
 #include "headless/public/util/deterministic_http_protocol_handler.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"
 
 namespace headless {
 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";
+// Default file name for pdf. Can be overriden by "--print-to-pdf" switch.
+const char kDefaultPDFFileName[] = "output.pdf";
 
 bool ParseWindowSize(std::string window_size, gfx::Size* parsed_window_size) {
   int width, height = 0;
   if (sscanf(window_size.c_str(), "%d%*[x,]%d", &width, &height) >= 2 &&
       width >= 0 && height >= 0) {
     parsed_window_size->set_width(width);
     parsed_window_size->set_height(height);
     return true;
   }
   return false;
@@ skipping 203 matching lines @@
   if (base::CommandLine::ForCurrentProcess()->HasSwitch(switches::kDumpDom)) {
     FetchDom();
   } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
                  switches::kRepl)) {
     LOG(INFO)
         << "Type a Javascript expression to evaluate or \"quit\" to exit.";
     InputExpression();
   } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
                  switches::kScreenshot)) {
     CaptureScreenshot();
+  } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+                 switches::kPrintToPDF)) {
+    PrintToPDF();
   } else {
     Shutdown();
   }
 }
 
 void HeadlessShell::FetchDom() {
   devtools_client_->GetRuntime()->Evaluate(
       "document.body.outerHTML",
       base::Bind(&HeadlessShell::OnDomFetched, weak_factory_.GetWeakPtr()));
 }
@@ skipping 44 matching lines @@
 
 void HeadlessShell::CaptureScreenshot() {
   devtools_client_->GetPage()->GetExperimental()->CaptureScreenshot(
       page::CaptureScreenshotParams::Builder().SetFromSurface(true).Build(),
       base::Bind(&HeadlessShell::OnScreenshotCaptured,
                  weak_factory_.GetWeakPtr()));
 }
 
 void HeadlessShell::OnScreenshotCaptured(
     std::unique_ptr<page::CaptureScreenshotResult> result) {
+  if (!result) {
+    LOG(ERROR) << "Capture screenshot failed";
+    Shutdown();
+    return;
+  }
+  WriteFile(switches::kScreenshot, kDefaultScreenshotFileName,
+            result->GetData());
+}
+
+void HeadlessShell::PrintToPDF() {
+  devtools_client_->GetPage()->GetExperimental()->PrintToPDF(
+      page::PrintToPDFParams::Builder().Build(),
+      base::Bind(&HeadlessShell::OnPDFCreated, weak_factory_.GetWeakPtr()));
+}
+
+void HeadlessShell::OnPDFCreated(
+    std::unique_ptr<page::PrintToPDFResult> result) {
+  if (!result) {
+    LOG(ERROR) << "Print to PDF failed";
+    Shutdown();
+    return;
+  }
+  WriteFile(switches::kPrintToPDF, kDefaultPDFFileName, result->GetData());
+}
+
+void HeadlessShell::WriteFile(const std::string& file_path_switch,
+                              const std::string& default_file_name,
+                              const std::string& base64_data) {
   base::FilePath file_name =
       base::CommandLine::ForCurrentProcess()->GetSwitchValuePath(
-          switches::kScreenshot);
-  if (file_name.empty()) {
-    file_name = base::FilePath().AppendASCII(kDefaultScreenshotFileName);
-  }
+          file_path_switch);
+  if (file_name.empty())
+    file_name = base::FilePath().AppendASCII(default_file_name);
 
-  screenshot_file_proxy_.reset(
-      new base::FileProxy(browser_->BrowserFileThread().get()));
-  if (!screenshot_file_proxy_->CreateOrOpen(
+  file_proxy_.reset(new base::FileProxy(browser_->BrowserFileThread().get()));
+  if (!file_proxy_->CreateOrOpen(
           file_name, base::File::FLAG_CREATE_ALWAYS | base::File::FLAG_WRITE,
-          base::Bind(&HeadlessShell::OnScreenshotFileOpened,
-                     weak_factory_.GetWeakPtr(),
-                     base::Passed(std::move(result)), file_name))) {
+          base::Bind(&HeadlessShell::OnFileOpened, weak_factory_.GetWeakPtr(),
+                     base64_data, file_name))) {
     // Operation could not be started.
-    OnScreenshotFileOpened(nullptr, file_name, base::File::FILE_ERROR_FAILED);
+    OnFileOpened(std::string(), file_name, base::File::FILE_ERROR_FAILED);
   }
 }
 
-void HeadlessShell::OnScreenshotFileOpened(
-    std::unique_ptr<page::CaptureScreenshotResult> result,
-    const base::FilePath file_name,
-    base::File::Error error_code) {
-  if (!screenshot_file_proxy_->IsValid()) {
-    LOG(ERROR) << "Writing screenshot to file " << file_name.value()
+void HeadlessShell::OnFileOpened(const std::string& base64_data,
+                                 const base::FilePath file_name,
+                                 base::File::Error error_code) {
+  if (!file_proxy_->IsValid()) {
+    LOG(ERROR) << "Writing to file " << file_name.value()
                << " was unsuccessful, could not open file: "
                << base::File::ErrorToString(error_code);
     return;
   }
 
-  std::string decoded_png;
-  base::Base64Decode(result->GetData(), &decoded_png);
+  std::string decoded_data;
+  if (!base::Base64Decode(base64_data, &decoded_data)) {
+    LOG(ERROR) << "Failed to decode base64 data";
+    return;

[Eric Seckler, 2017/04/03 10:34:34]

Should also close the file in this case. Maybe call into OnFileWritten:

OnFileWritten(file_name, decoded_data.size(), base::File::FILE_ERROR_FAILED, 0);

[jzfeng, 2017/04/04 05:05:19]

Done.

+  }
+
   scoped_refptr<net::IOBufferWithSize> buf =
-      new net::IOBufferWithSize(decoded_png.size());
-  memcpy(buf->data(), decoded_png.data(), decoded_png.size());
+      new net::IOBufferWithSize(decoded_data.size());
+  memcpy(buf->data(), decoded_data.data(), decoded_data.size());
 
-  if (!screenshot_file_proxy_->Write(
+  if (!file_proxy_->Write(
           0, buf->data(), buf->size(),
-          base::Bind(&HeadlessShell::OnScreenshotFileWritten,
-                     weak_factory_.GetWeakPtr(), file_name, buf->size()))) {
+          base::Bind(&HeadlessShell::OnFileWritten, weak_factory_.GetWeakPtr(),
+                     file_name, buf->size()))) {
     // Operation may have completed successfully or failed.
-    OnScreenshotFileWritten(file_name, buf->size(),
-                            base::File::FILE_ERROR_FAILED, 0);
+    OnFileWritten(file_name, buf->size(), base::File::FILE_ERROR_FAILED, 0);
   }
 }
 
-void HeadlessShell::OnScreenshotFileWritten(const base::FilePath file_name,
-                                            const int length,
-                                            base::File::Error error_code,
-                                            int write_result) {
+void HeadlessShell::OnFileWritten(const base::FilePath file_name,
+                                  const int length,
+                                  base::File::Error error_code,
+                                  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);
+    LOG(ERROR) << "Writing to file " << file_name.value()
+               << " was unsuccessful: "
+               << base::File::ErrorToString(error_code);
   } else {
-    LOG(INFO) << "Screenshot written to file " << file_name.value() << "."
-              << std::endl;
+    LOG(INFO) << "Written to file " << file_name.value() << ".";
   }
-  if (!screenshot_file_proxy_->Close(
-          base::Bind(&HeadlessShell::OnScreenshotFileClosed,
-                     weak_factory_.GetWeakPtr()))) {
+  if (!file_proxy_->Close(base::Bind(&HeadlessShell::OnFileClosed,
+                                     weak_factory_.GetWeakPtr()))) {
     // Operation could not be started.
-    OnScreenshotFileClosed(base::File::FILE_ERROR_FAILED);
+    OnFileClosed(base::File::FILE_ERROR_FAILED);
   }
 }
 
-void HeadlessShell::OnScreenshotFileClosed(base::File::Error error_code) {
+void HeadlessShell::OnFileClosed(base::File::Error error_code) {
   Shutdown();
 }
 
 bool HeadlessShell::RemoteDebuggingEnabled() const {
   const base::CommandLine& command_line =
       *base::CommandLine::ForCurrentProcess();
   return command_line.HasSwitch(switches::kRemoteDebuggingPort);
 }
 
 bool ValidateCommandLine(const base::CommandLine& command_line) {
@@ skipping 11 matching lines @@
   if (command_line.HasSwitch(switches::kRepl)) {
     LOG(ERROR) << "Evaluate Javascript is disabled "
                << "when remote debugging is enabled.";
     return false;
   }
   if (command_line.HasSwitch(switches::kScreenshot)) {
     LOG(ERROR) << "Capture screenshot is disabled "
                << "when remote debugging is enabled.";
     return false;
   }
+  if (command_line.HasSwitch(switches::kPrintToPDF)) {
+    LOG(ERROR) << "Print to PDF is disabled "
+               << "when remote debugging is enabled.";
+    return false;
+  }
   if (command_line.HasSwitch(switches::kTimeout)) {
     LOG(ERROR) << "Navigation timeout is disabled "
                << "when remote debugging is enabled.";
     return false;
   }
   if (command_line.HasSwitch(switches::kVirtualTimeBudget)) {
     LOG(ERROR) << "Virtual time budget is disabled "
                << "when remote debugging is enabled.";
     return false;
   }
@@ skipping 88 matching lines @@
     builder.SetOverrideWebPreferencesCallback(base::Bind([](
         WebPreferences* preferences) { preferences->hide_scrollbars = true; }));
   }
 
   return HeadlessBrowserMain(
       builder.Build(),
       base::Bind(&HeadlessShell::OnStart, base::Unretained(&shell)));
 }
 
 }  // namespace headless