Reland "assistant: filter out incognito windows in assistant screenshot."

TBR=dcheng@chromium.org,jamescook@chromium.org,msw@chromium.org,xiaohuic@chromium.org

Original change's description:
> assistant: filter out incognito windows in assistant screenshot.
>
> Per privacy requirement, incognito windows should not be included
> in screenshot sent to assistant server.
>
> Bug: b/78193324
> Test:
>   ash_unittests --gtest_filter=AssistantControllerTest.Screenshot
>   browser_tests --gtest_filter=BrowserNonClientFrameViewAshTest.IncognitoMarkedAsAssistantBlocked*
>
> Change-Id: Idf55135959721d69414a1ec6fae396f9763ebe49
> Reviewed-on: https://chromium-review.googlesource.com/1114293
> Commit-Queue: Muyuan Li <muyuanli@chromium.org>
> Reviewed-by: James Cook <jamescook@chromium.org>
> Reviewed-by: Daniel Cheng <dcheng@chromium.org>
> Reviewed-by: Xiaohui Chen <xiaohuic@chromium.org>
> Reviewed-by: Michael Wasserman <msw@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#571385}

Change-Id: Ieecafffe5f1e14971b3c07dc799323a29885de16
Reviewed-on: https://chromium-review.googlesource.com/1119766Reviewed-by: Muyuan Li <muyuanli@chromium.org>
Commit-Queue: Muyuan Li <muyuanli@chromium.org>
Cr-Commit-Position: refs/heads/master@{#571407}

ash/assistant/assistant_controller.cc

@@ -9,15 +9,104 @@
 #include "ash/assistant/assistant_ui_controller.h"
 #include "ash/assistant/util/deep_link_util.h"
 #include "ash/new_window_controller.h"
+#include "ash/public/cpp/shell_window_ids.h"
+#include "ash/public/cpp/window_properties.h"
 #include "ash/session/session_controller.h"
 #include "ash/shell.h"
+#include "ash/wm/mru_window_tracker.h"
 #include "base/bind.h"
 #include "base/memory/scoped_refptr.h"
+#include "base/stl_util.h"
+#include "base/task_scheduler/post_task.h"
 #include "base/unguessable_token.h"
+#include "ui/aura/client/aura_constants.h"
+#include "ui/compositor/layer_tree_owner.h"
+#include "ui/gfx/codec/jpeg_codec.h"
+#include "ui/gfx/skbitmap_operations.h"
 #include "ui/snapshot/snapshot.h"
+#include "ui/snapshot/snapshot_aura.h"
+#include "ui/wm/core/window_util.h"
 
 namespace ash {
 
+namespace {
+
+// When screenshot's width or height is smaller than this size, we will stop
+// downsampling.
+constexpr int kScreenshotMaxWidth = 1366;
+constexpr int kScreenshotMaxHeight = 768;
+
+std::vector<uint8_t> DownsampleAndEncodeImage(gfx::Image image) {
+  std::vector<uint8_t> res;
+  // We'll downsample the screenshot to avoid exceeding max allowed size on
+  // assistant server side if we are taking screenshot from high-res screen.
+  gfx::JPEGCodec::Encode(
+      SkBitmapOperations::DownsampleByTwoUntilSize(
+          image.AsBitmap(), kScreenshotMaxWidth, kScreenshotMaxHeight),
+      100 /* quality */, &res);
+  return res;
+}
+
+void EncodeScreenshotAndRunCallback(
+    AssistantController::RequestScreenshotCallback callback,
+    std::unique_ptr<ui::LayerTreeOwner> layer_owner,
+    gfx::Image image) {
+  base::PostTaskWithTraitsAndReplyWithResult(
+      FROM_HERE,
+      base::TaskTraits{base::MayBlock(), base::TaskPriority::USER_BLOCKING},
+      base::BindOnce(&DownsampleAndEncodeImage, std::move(image)),
+      std::move(callback));
+}
+
+std::unique_ptr<ui::LayerTreeOwner> CreateLayerForAssistantSnapshot(
+    aura::Window* root_window) {
+  using LayerSet = base::flat_set<const ui::Layer*>;
+  LayerSet excluded_layers;
+  LayerSet blocked_layers;
+
+  aura::Window* overlay_container =
+      ash::Shell::GetContainer(root_window, kShellWindowId_OverlayContainer);
+  if (overlay_container)
+    excluded_layers.insert(overlay_container->layer());
+
+  MruWindowTracker::WindowList windows =
+      Shell::Get()->mru_window_tracker()->BuildMruWindowList();
+
+  for (aura::Window* window : windows) {
+    if (window->GetProperty(kBlockedForAssistantSnapshotKey))
+      blocked_layers.insert(window->layer());
+  }
+
+  return ::wm::RecreateLayersWithClosure(
+      root_window,
+      base::BindRepeating(
+          [](LayerSet blocked_layers, LayerSet excluded_layers,
+             ui::LayerOwner* owner) -> std::unique_ptr<ui::Layer> {
+            // Parent layer is excluded meaning that it's pointless to clone
+            // current child and all its descendants. This reduces the number
+            // of layers to create.
+            if (base::ContainsKey(blocked_layers, owner->layer()->parent()))
+              return nullptr;
+            if (base::ContainsKey(blocked_layers, owner->layer())) {
+              // Blocked layers are replaced with solid black layers so that
+              // they won't be included in the screenshot but still preserve
+              // the window stacking.
+              auto layer =
+                  std::make_unique<ui::Layer>(ui::LayerType::LAYER_SOLID_COLOR);
+              layer->SetBounds(owner->layer()->bounds());
+              layer->SetColor(SK_ColorBLACK);
+              return layer;
+            }
+            if (excluded_layers.count(owner->layer()))
+              return nullptr;
+            return owner->RecreateLayer();
+
+          },
+          std::move(blocked_layers), std::move(excluded_layers)));
+}
+
+}  // namespace
+
 AssistantController::AssistantController()
     : assistant_interaction_controller_(
           std::make_unique<AssistantInteractionController>(this)),
@@ -86,18 +175,29 @@ void AssistantController::RequestScreenshot(
     const gfx::Rect& rect,
     RequestScreenshotCallback callback) {
   // TODO(muyuanli): handle multi-display when assistant's behavior is defined.
-  auto* root_window = Shell::GetPrimaryRootWindow();
+  aura::Window* root_window = Shell::GetPrimaryRootWindow();
+
+  std::unique_ptr<ui::LayerTreeOwner> layer_owner =
+      CreateLayerForAssistantSnapshot(root_window);
+
+  ui::Layer* root_layer = layer_owner->root();
+
   gfx::Rect source_rect =
       rect.IsEmpty() ? gfx::Rect(root_window->bounds().size()) : rect;
-  ui::GrabWindowSnapshotAsyncJPEG(
-      root_window, source_rect,
-      base::BindRepeating(
-          [](RequestScreenshotCallback callback,
-             scoped_refptr<base::RefCountedMemory> data) {
-            std::move(callback).Run(std::vector<uint8_t>(
-                data->front(), data->front() + data->size()));
-          },
-          base::Passed(&callback)));
+
+  // The root layer might have a scaling transform applied (if the user has
+  // changed the UI scale via Ctrl-Shift-Plus/Minus).
+  // Clear the transform so that the snapshot is taken at 1:1 scale relative
+  // to screen pixels.
+  root_layer->SetTransform(gfx::Transform());
+  root_window->layer()->Add(root_layer);
+  root_window->layer()->StackAtBottom(root_layer);
+
+  ui::GrabLayerSnapshotAsync(
+      root_layer, source_rect,
+      base::BindRepeating(&EncodeScreenshotAndRunCallback,
+                          base::Passed(std::move(callback)),
+                          base::Passed(std::move(layer_owner))));
 }
 
 void AssistantController::ManageWebContents(
@@ -179,4 +279,10 @@ void AssistantController::OpenUrl(const GURL& url) {
     observer.OnUrlOpened(url);
 }
 
+std::unique_ptr<ui::LayerTreeOwner>
+AssistantController::CreateLayerForAssistantSnapshotForTest() {
+  aura::Window* root_window = Shell::GetPrimaryRootWindow();
+  return CreateLayerForAssistantSnapshot(root_window);
+}
+
 }  // namespace ash

ash/assistant/assistant_controller.h

@@ -9,6 +9,7 @@
 #include <string>
 #include <vector>
 
+#include "ash/ash_export.h"
 #include "ash/highlighter/highlighter_controller.h"
 #include "ash/public/interfaces/assistant_controller.mojom.h"
 #include "ash/public/interfaces/assistant_image_downloader.mojom.h"
@@ -24,15 +25,20 @@ namespace base {
 class UnguessableToken;
 }  // namespace base
 
+namespace ui {
+class LayerTreeOwner;
+}  // namespace ui
+
 namespace ash {
 
 class AssistantControllerObserver;
 class AssistantInteractionController;
 class AssistantUiController;
 
-class AssistantController : public mojom::AssistantController,
-                            public HighlighterController::Observer,
-                            public mojom::ManagedWebContentsOpenUrlDelegate {
+class ASH_EXPORT AssistantController
+    : public mojom::AssistantController,
+      public HighlighterController::Observer,
+      public mojom::ManagedWebContentsOpenUrlDelegate {
  public:
   AssistantController();
   ~AssistantController() override;
@@ -99,6 +105,8 @@ class AssistantController : public mojom::AssistantController,
     return assistant_ui_controller_.get();
   }
 
+  std::unique_ptr<ui::LayerTreeOwner> CreateLayerForAssistantSnapshotForTest();
+
  private:
   // The observer list should be initialized early so that sub-controllers may
   // register as observers during their construction.

ash/assistant/assistant_controller_unittest.cc

@@ -7,6 +7,8 @@
 #include <memory>
 
 #include "ash/assistant/assistant_ui_controller.h"
+#include "ash/public/cpp/shell_window_ids.h"
+#include "ash/public/cpp/window_properties.h"
 #include "ash/shell.h"
 #include "ash/test/ash_test_base.h"
 #include "base/macros.h"
@@ -14,9 +16,30 @@
 #include "chromeos/chromeos_switches.h"
 #include "chromeos/services/assistant/test_support/mock_assistant.h"
 #include "mojo/public/cpp/bindings/binding.h"
+#include "ui/aura/window.h"
+#include "ui/compositor/layer.h"
+#include "ui/compositor/layer_tree_owner.h"
+#include "ui/compositor/layer_type.h"
 
 namespace ash {
 
+namespace {
+
+ui::Layer* FindLayerWithClosure(
+    ui::Layer* root,
+    const base::RepeatingCallback<bool(ui::Layer*)>& callback) {
+  if (callback.Run(root))
+    return root;
+  for (ui::Layer* child : root->children()) {
+    ui::Layer* result = FindLayerWithClosure(child, callback);
+    if (result)
+      return result;
+  }
+  return nullptr;
+}
+
+}  // namespace
+
 class AssistantControllerTest : public AshTestBase {
  protected:
   AssistantControllerTest() = default;
@@ -48,6 +71,8 @@ class AssistantControllerTest : public AshTestBase {
     return controller_->ui_controller()->model()->visible();
   }
 
+  ash::AssistantController* controller() { return controller_; }
+
  private:
   base::test::ScopedFeatureList scoped_feature_list_;
 
@@ -61,4 +86,38 @@ class AssistantControllerTest : public AshTestBase {
   DISALLOW_COPY_AND_ASSIGN(AssistantControllerTest);
 };
 
+// Verify that incognito windows are blocked in screenshot.
+TEST_F(AssistantControllerTest, Screenshot) {
+  std::unique_ptr<aura::Window> window1 = CreateToplevelTestWindow(
+      gfx::Rect(0, 0, 200, 200), kShellWindowId_DefaultContainer);
+  std::unique_ptr<aura::Window> window2 = CreateToplevelTestWindow(
+      gfx::Rect(30, 30, 100, 100), kShellWindowId_DefaultContainer);
+
+  ui::Layer* window1_layer = window1->layer();
+  ui::Layer* window2_layer = window2->layer();
+
+  window1->SetProperty(kBlockedForAssistantSnapshotKey, true);
+
+  std::unique_ptr<ui::LayerTreeOwner> layer_owner =
+      controller()->CreateLayerForAssistantSnapshotForTest();
+
+  // Test that windows marked as blocked for assistant snapshot is not included.
+  EXPECT_FALSE(FindLayerWithClosure(
+      layer_owner->root(),
+      base::BindRepeating(
+          [](ui::Layer* target, ui::Layer* layer) { return layer == target; },
+          window1_layer)));
+  EXPECT_TRUE(FindLayerWithClosure(
+      layer_owner->root(),
+      base::BindRepeating(
+          [](ui::Layer* target, ui::Layer* layer) { return layer == target; },
+          window2_layer)));
+
+  // Test that a solid black layer is inserted.
+  EXPECT_TRUE(FindLayerWithClosure(
+      layer_owner->root(), base::BindRepeating([](ui::Layer* layer) {
+        return layer->type() == ui::LayerType::LAYER_SOLID_COLOR;
+      })));
+}
+
 }  // namespace ash

ash/public/cpp/mus_property_mirror_ash.cc

@@ -31,7 +31,11 @@ void MusPropertyMirrorAsh::MirrorPropertyFromWidgetWindowToRootWindow(
     aura::Window* window,
     aura::Window* root_window,
     const void* key) {
-  if (key == kPanelAttachedKey) {
+  if (key == kBlockedForAssistantSnapshotKey) {
+    root_window->SetProperty(
+        kBlockedForAssistantSnapshotKey,
+        window->GetProperty(kBlockedForAssistantSnapshotKey));
+  } else if (key == kPanelAttachedKey) {
     bool value = window->GetProperty(kPanelAttachedKey);
     root_window->SetProperty(kPanelAttachedKey, value);
   } else if (key == kShelfItemTypeKey) {

ash/public/cpp/window_properties.cc

@@ -32,6 +32,10 @@ DEFINE_EXPORTED_UI_CLASS_PROPERTY_TYPE(ASH_PUBLIC_EXPORT,
 namespace ash {
 
 void RegisterWindowProperties(aura::PropertyConverter* property_converter) {
+  property_converter->RegisterPrimitiveProperty(
+      kBlockedForAssistantSnapshotKey,
+      mojom::kBlockedForAssistantSnapshot_Property,
+      aura::PropertyConverter::CreateAcceptAnyValueCallback());
   property_converter->RegisterPrimitiveProperty(
       kCanConsumeSystemKeysKey, mojom::kCanConsumeSystemKeys_Property,
       aura::PropertyConverter::CreateAcceptAnyValueCallback());
@@ -103,6 +107,7 @@ void RegisterWindowProperties(aura::PropertyConverter* property_converter) {
 DEFINE_UI_CLASS_PROPERTY_KEY(BackdropWindowMode,
                              kBackdropWindowMode,
                              BackdropWindowMode::kAuto);
+DEFINE_UI_CLASS_PROPERTY_KEY(bool, kBlockedForAssistantSnapshotKey, false);
 DEFINE_UI_CLASS_PROPERTY_KEY(bool, kCanConsumeSystemKeysKey, false);
 DEFINE_UI_CLASS_PROPERTY_KEY(FrameBackButtonState,
                              kFrameBackButtonStateKey,

ash/public/cpp/window_properties.h

@@ -63,6 +63,11 @@ ASH_PUBLIC_EXPORT void RegisterWindowProperties(
 ASH_PUBLIC_EXPORT extern const aura::WindowProperty<BackdropWindowMode>* const
     kBackdropWindowMode;
 
+// If set to true, the window will be replaced by a black rectangle when taking
+// screenshot for assistant. Used to preserve privacy for incognito windows.
+ASH_PUBLIC_EXPORT extern const aura::WindowProperty<bool>* const
+    kBlockedForAssistantSnapshotKey;
+
 // If true, will send system keys to the window for dispatch.
 ASH_PUBLIC_EXPORT extern const aura::WindowProperty<bool>* const
     kCanConsumeSystemKeysKey;

ash/public/interfaces/window_properties.mojom

@@ -4,6 +4,11 @@
 
 module ash.mojom;
 
+// A bool to indicate whether this window should be replaced by a black
+// rectangle in assistant screenshot for privacy purpose.
+const string kBlockedForAssistantSnapshot_Property =
+  "ash:blocked-for-assistant-snapshot";
+
 // V1 apps can intercept system keys. This will let the app handle F-keys instead
 // of the window manager.
 const string kCanConsumeSystemKeys_Property =

chrome/browser/ui/views/DEPS

@@ -26,5 +26,6 @@ specific_include_rules = {
   ],
   ".*test.*": [
    "!ash",
+   "+ash/public",
   ],
 }

chrome/browser/ui/views/frame/browser_non_client_frame_view_ash.cc

@@ -192,6 +192,13 @@ void BrowserNonClientFrameViewAsh::Init() {
       static_cast<int>(browser->is_app() ? ash::AppType::CHROME_APP
                                          : ash::AppType::BROWSER));
 
+  // To preserve privacy, tag incognito windows so that they won't be included
+  // in screenshot sent to assistant server.
+  if (browser->profile()->IsOffTheRecord()) {
+    frame()->GetNativeWindow()->SetProperty(
+        ash::kBlockedForAssistantSnapshotKey, true);
+  }
+
   // TODO(estade): how much of the rest of this needs porting to Mash?
   if (IsMash()) {
     OnThemeChanged();

chrome/browser/ui/views/frame/browser_non_client_frame_view_ash_browsertest.cc

@@ -16,6 +16,7 @@
 #include "ash/public/cpp/ash_switches.h"
 #include "ash/public/cpp/immersive/immersive_fullscreen_controller_test_api.h"
 #include "ash/public/cpp/vector_icons/vector_icons.h"
+#include "ash/public/cpp/window_properties.h"
 #include "ash/shell.h"
 #include "ash/wm/overview/window_selector_controller.h"
 #include "ash/wm/splitview/split_view_controller.h"
@@ -373,6 +374,13 @@ IN_PROC_BROWSER_TEST_P(BrowserNonClientFrameViewAshTest, IncognitoAvatar) {
   EXPECT_EQ(should_have_avatar, has_avatar);
 }
 
+IN_PROC_BROWSER_TEST_P(BrowserNonClientFrameViewAshTest,
+                       IncognitoMarkedAsAssistantBlocked) {
+  Browser* incognito_browser = CreateIncognitoBrowser();
+  EXPECT_TRUE(incognito_browser->window()->GetNativeWindow()->GetProperty(
+      ash::kBlockedForAssistantSnapshotKey));
+}
+
 // Tests that FrameCaptionButtonContainer has been relaid out in response to
 // tablet mode being toggled.
 IN_PROC_BROWSER_TEST_P(BrowserNonClientFrameViewAshTest,