Add autotest private API for taking a screenshot

This is a temporary solution to taking screenshots for tast-tests that
will work properly on ARM devices as well as deal with hardware
overlays. Longer term the plan is to implement this via the DRM API with
a write-back connector and then use that from the CLI screenshot
program.

BUG=chromium:849438,chromium:880597,chromium:887016
TEST=tast run graphics.Screenshot (after updating it to use this)

Change-Id: I3919822599338e4786c715d7b0419742417b83c2
Reviewed-on: https://chromium-review.googlesource.com/1247181
Commit-Queue: Jeffrey Kardatzke <jkardatzke@google.com>
Reviewed-by: James Cook <jamescook@chromium.org>
Reviewed-by: Ben Wells <benwells@chromium.org>
Reviewed-by: Achuith Bhandarkar <achuith@chromium.org>
Cr-Commit-Position: refs/heads/master@{#595173}

chrome/browser/chromeos/extensions/autotest_private/DEPS

 include_rules = [
   "+chrome/browser/ui/views/crostini",
 ]
+
+specific_include_rules = {
+  "autotest_private_api\.cc": [
+    # TODO(mash): Remove the line below, http://crbug.com/557397
+    "+ash/shell.h",
+  ]
+}

chrome/browser/chromeos/extensions/autotest_private/autotest_private_api.cc

@@ -34,6 +34,8 @@
 #if defined(OS_CHROMEOS)
 #include "ash/public/interfaces/ash_message_center_controller.mojom.h"
 #include "ash/public/interfaces/constants.mojom.h"
+#include "ash/shell.h"
+#include "base/base64.h"
 #include "base/feature_list.h"
 #include "chrome/browser/chromeos/arc/arc_util.h"
 #include "chrome/browser/chromeos/crostini/crostini_manager.h"
@@ -842,6 +844,47 @@ void AutotestPrivateRunCrostiniUninstallerFunction::CrostiniRemoved(
 }
 #endif
 
+AutotestPrivateTakeScreenshotFunction::
+    ~AutotestPrivateTakeScreenshotFunction() = default;
+
+ExtensionFunction::ResponseAction AutotestPrivateTakeScreenshotFunction::Run() {
+  DVLOG(1) << "AutotestPrivateTakeScreenshotFunction";
+#if defined(OS_CHROMEOS)
+  auto grabber = std::make_unique<ui::ScreenshotGrabber>();
+  // TODO(mash): Fix for mash, http://crbug.com/557397
+  aura::Window* primary_root = ash::Shell::GetPrimaryRootWindow();
+  // Pass the ScreenshotGrabber to the callback so that it stays alive for the
+  // duration of the operation, it'll then get deallocated when the callback
+  // completes.
+  grabber->TakeScreenshot(
+      primary_root, primary_root->bounds(),
+      base::BindOnce(&AutotestPrivateTakeScreenshotFunction::ScreenshotTaken,
+                     this, base::Passed(&grabber)));
+  return RespondLater();
+#else
+  return RespondNow(Error(kOnlyAvailableOnChromeOSError));
+#endif
+}
+
+#if defined(OS_CHROMEOS)
+void AutotestPrivateTakeScreenshotFunction::ScreenshotTaken(
+    std::unique_ptr<ui::ScreenshotGrabber> grabber,
+    ui::ScreenshotResult screenshot_result,
+    scoped_refptr<base::RefCountedMemory> png_data) {
+  if (screenshot_result == ui::ScreenshotResult::SUCCESS) {
+    // Base64 encode the result so we can return it as a string.
+    std::string base64Png(png_data->front(),
+                          png_data->front() + png_data->size());
+    base::Base64Encode(base64Png, &base64Png);
+    Respond(OneArgument(std::make_unique<base::Value>(base64Png)));
+  } else {
+    Respond(Error(base::StrCat(
+        {"Error taking screenshot ",
+         base::NumberToString(static_cast<int>(screenshot_result))})));
+  }
+}
+#endif
+
 AutotestPrivateBootstrapMachineLearningServiceFunction::
     ~AutotestPrivateBootstrapMachineLearningServiceFunction() = default;
 

chrome/browser/chromeos/extensions/autotest_private/autotest_private_api.h

@@ -16,6 +16,7 @@
 #include "ash/public/interfaces/ash_message_center_controller.mojom.h"
 #include "chrome/browser/chromeos/printing/cups_printers_manager.h"
 #include "chromeos/services/machine_learning/public/mojom/machine_learning_service.mojom.h"
+#include "ui/snapshot/screenshot_grabber.h"
 #endif
 
 namespace message_center {
@@ -310,6 +311,23 @@ class AutotestPrivateRunCrostiniUninstallerFunction
   DISALLOW_COPY_AND_ASSIGN(AutotestPrivateRunCrostiniUninstallerFunction);
 };
 
+class AutotestPrivateTakeScreenshotFunction : public UIThreadExtensionFunction {
+ public:
+  AutotestPrivateTakeScreenshotFunction() = default;
+  DECLARE_EXTENSION_FUNCTION("autotestPrivate.takeScreenshot",
+                             AUTOTESTPRIVATE_TAKESCREENSHOT)
+
+ private:
+  ~AutotestPrivateTakeScreenshotFunction() override;
+  ResponseAction Run() override;
+#if defined(OS_CHROMEOS)
+  void ScreenshotTaken(std::unique_ptr<ui::ScreenshotGrabber> grabber,
+                       ui::ScreenshotResult screenshot_result,
+                       scoped_refptr<base::RefCountedMemory> png_data);
+#endif
+  DISALLOW_COPY_AND_ASSIGN(AutotestPrivateTakeScreenshotFunction);
+};
+
 class AutotestPrivateGetPrinterListFunction : public UIThreadExtensionFunction {
  public:
   AutotestPrivateGetPrinterListFunction() = default;

chrome/common/extensions/api/autotest_private.idl

@@ -104,6 +104,8 @@ namespace autotestPrivate {
 
   callback IsAppShownCallback = void (boolean appShown);
 
+  callback TakeScreenshotCallback = void (DOMString base64Png);
+
   callback VoidCallback = void ();
 
   interface Functions {
@@ -208,6 +210,9 @@ namespace autotestPrivate {
     // |callback|: Called when the operation has completed.
     static void setCrostiniEnabled(boolean enabled, VoidCallback callback);
 
+    // Takes a screenshot and returns the data in base64 encoded PNG format.
+    static void takeScreenshot(TakeScreenshotCallback callback);
+
     // Makes a basic request to ML Service, triggering 1. ML Service
     // daemon startup, and 2. the initial D-Bus -> Mojo IPC bootstrap.
     // |callback|: Called when the operation has completed.

chrome/test/data/extensions/api_test/autotest_private/test.js

@@ -193,6 +193,15 @@ chrome.test.runTests([
     chrome.autotestPrivate.runCrostiniUninstaller(chrome.test.callbackFail(
         'Crostini is not available for the current user'));
   },
+  function takeScreenshot() {
+    chrome.autotestPrivate.takeScreenshot(
+      function(base64Png) {
+        chrome.test.assertTrue(base64Png.length > 0);
+        chrome.test.assertNoLastError();
+        chrome.test.succeed();
+      }
+    )
+  },
   function getPrinterList() {
     chrome.autotestPrivate.getPrinterList(function(){
       chrome.test.succeed();

extensions/browser/extension_function_histogram_value.h

@@ -1343,6 +1343,7 @@ enum HistogramValue {
   AUTOTESTPRIVATE_LAUNCHAPP = 1280,
   AUTOTESTPRIVATE_BOOTSTRAPMACHINELEARNINGSERVICE = 1281,
   AUTOTESTPRIVATE_RUNCROSTINIUNINSTALLER = 1282,
+  AUTOTESTPRIVATE_TAKESCREENSHOT = 1283,
   // Last entry: Add new entries above, then run:
   // python tools/metrics/histograms/update_extension_histograms.py
   ENUM_BOUNDARY

tools/metrics/histograms/enums.xml

@@ -16880,6 +16880,7 @@ Called by update_net_error_codes.py.-->
   <int value="1280" label="AUTOTESTPRIVATE_LAUNCHAPP"/>
   <int value="1281" label="AUTOTESTPRIVATE_BOOTSTRAPMACHINELEARNINGSERVICE"/>
   <int value="1282" label="AUTOTESTPRIVATE_RUNCROSTINIUNINSTALLER"/>
+  <int value="1283" label="AUTOTESTPRIVATE_TAKESCREENSHOT"/>
 </enum>
 
 <enum name="ExtensionIconState">