Add Page.stopLoading to devtools and --timeout switch to headless_shell

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/memory/weak_ptr.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/devtools/domains/emulation.h"
 #include "headless/public/devtools/domains/inspector.h"
 #include "headless/public/devtools/domains/page.h"
 #include "headless/public/devtools/domains/runtime.h"
 #include "headless/public/headless_browser.h"
 #include "headless/public/headless_devtools_client.h"
@@ skipping 30 matching lines @@
 class HeadlessShell : public HeadlessWebContents::Observer,
                       emulation::ExperimentalObserver,
                       inspector::ExperimentalObserver,
                       page::Observer {
  public:
   HeadlessShell()
       : browser_(nullptr),
         devtools_client_(HeadlessDevToolsClient::Create()),
         web_contents_(nullptr),
         processed_page_ready_(false),
-        browser_context_(nullptr) {}
+        browser_context_(nullptr),
+        weak_factory_(this) {}
   ~HeadlessShell() override {}
 
   void OnStart(HeadlessBrowser* browser) {
     browser_ = browser;
 
     HeadlessBrowserContext::Builder context_builder =
         browser_->CreateBrowserContextBuilder();
     if (base::CommandLine::ForCurrentProcess()->HasSwitch(
             switches::kDeterministicFetch)) {
       deterministic_dispatcher_.reset(
@@ skipping 78 matching lines @@
           << "Expected an integer value for --virtual-time-budget=";
       devtools_client_->GetEmulation()->GetExperimental()->SetVirtualTimePolicy(
           emulation::SetVirtualTimePolicyParams::Builder()
               .SetPolicy(emulation::VirtualTimePolicy::
                              PAUSE_IF_NETWORK_FETCHES_PENDING)
               .SetBudget(budget_ms)
               .Build());
     } else {
       PollReadyState();
     }
+
+    if (base::CommandLine::ForCurrentProcess()->HasSwitch(switches::kTimeout)) {
+      std::string timeout_ms_ascii =
+          base::CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
+              switches::kTimeout);
+      int timeout_ms;
+      CHECK(base::StringToInt(timeout_ms_ascii, &timeout_ms))
+          << "Expected an integer value for --timeout=";
+      browser_->BrowserMainThread()->PostDelayedTask(
+          FROM_HERE,
+          base::Bind(&HeadlessShell::FetchTimeout, weak_factory_.GetWeakPtr()),
+          base::TimeDelta::FromMilliseconds(timeout_ms));
+    }
+
     // TODO(skyostil): Implement more features to demonstrate the devtools API.
   }
 
+  void FetchTimeout() {
+    LOG(INFO) << "Timeout.";
+    devtools_client_->GetPage()->GetExperimental()->StopLoading(
+        page::StopLoadingParams::Builder().Build());
+  }
+
   void OnTargetCrashed(const inspector::TargetCrashedParams& params) override {
     LOG(ERROR) << "Abnormal renderer termination.";
     // NB this never gets called if remote debugging is enabled.
     Shutdown();
   }
 
   void PollReadyState() {
     // We need to check the current location in addition to the ready state to
     // be sure the expected page is ready.
     devtools_client_->GetRuntime()->Evaluate(
         "document.readyState + ' ' + document.location.href",
-        base::Bind(&HeadlessShell::OnReadyState, base::Unretained(this)));
+        base::Bind(&HeadlessShell::OnReadyState, weak_factory_.GetWeakPtr()));
   }
 
   void OnReadyState(std::unique_ptr<runtime::EvaluateResult> result) {
     std::string ready_state_and_url;
     if (result->GetResult()->GetValue()->GetAsString(&ready_state_and_url)) {
       std::stringstream stream(ready_state_and_url);
       std::string ready_state;
       std::string url;
       stream >> ready_state;
       stream >> url;
@@ skipping 38 matching lines @@
                    switches::kScreenshot)) {
       CaptureScreenshot();
     } else {
       Shutdown();
     }
   }
 
   void FetchDom() {
     devtools_client_->GetRuntime()->Evaluate(
         "document.body.innerHTML",
-        base::Bind(&HeadlessShell::OnDomFetched, base::Unretained(this)));
+        base::Bind(&HeadlessShell::OnDomFetched, weak_factory_.GetWeakPtr()));
   }
 
   void OnDomFetched(std::unique_ptr<runtime::EvaluateResult> result) {
     if (result->HasExceptionDetails()) {
       LOG(ERROR) << "Failed to evaluate document.body.innerHTML: "
                  << result->GetExceptionDetails()->GetText();
     } else {
       std::string dom;
       if (result->GetResult()->GetValue()->GetAsString(&dom)) {
         std::cout << dom << std::endl;
       }
     }
     Shutdown();
   }
 
   void InputExpression() {
     // Note that a real system should read user input asynchronously, because
     // otherwise all other browser activity is suspended (e.g., page loading).
     std::string expression;
     std::cout << ">>> ";
     std::getline(std::cin, expression);
     if (std::cin.bad() || std::cin.eof() || expression == "quit") {
       Shutdown();
       return;
     }
     devtools_client_->GetRuntime()->Evaluate(
-        expression,
-        base::Bind(&HeadlessShell::OnExpressionResult, base::Unretained(this)));
+        expression, base::Bind(&HeadlessShell::OnExpressionResult,
+                               weak_factory_.GetWeakPtr()));
   }
 
   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)));
+                   weak_factory_.GetWeakPtr()));
   }
 
   void OnScreenshotCaptured(
       std::unique_ptr<page::CaptureScreenshotResult> result) {
     base::FilePath file_name =
         base::CommandLine::ForCurrentProcess()->GetSwitchValuePath(
             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)),
+                   weak_factory_.GetWeakPtr(), 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()));
+                   weak_factory_.GetWeakPtr(), 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)));
+        &HeadlessShell::OnScreenshotFileClosed, weak_factory_.GetWeakPtr()));
     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();
     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_;
   HeadlessBrowserContext* browser_context_;
   std::unique_ptr<DeterministicDispatcher> deterministic_dispatcher_;
+  base::WeakPtrFactory<HeadlessShell> weak_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(HeadlessShell);
 };
 
 int HeadlessShellMain(int argc, const char** argv) {
   RunChildProcessIfNeeded(argc, argv);
   HeadlessShell shell;
   HeadlessBrowser::Options::Builder builder(argc, argv);
 
   // Enable devtools if requested.
@@ skipping 62 matching lines @@
     }
     builder.SetWindowSize(parsed_window_size);
   }
 
   return HeadlessBrowserMain(
       builder.Build(),
       base::Bind(&HeadlessShell::OnStart, base::Unretained(&shell)));
 }
 
 }  // namespace headless