Recreates the layer hierarchy modulo all layers that

belong to incognito windows, replacing them with a solid black
layer.

BUG=b/37755798
TEST=wm_unittests --gtest_filter=WindowUtilTest.RecreateLayersWithClosure

Review-Url: https://codereview.chromium.org/2854543002
Cr-Commit-Position: refs/heads/master@{#472678}

chrome/browser/chromeos/arc/voice_interaction/arc_voice_interaction_framework_service.cc

@@ -9,16 +9,29 @@
 
 #include "ash/accelerators/accelerator_controller.h"
 #include "ash/shell.h"
+#include "base/bind.h"
 #include "base/command_line.h"
+#include "base/containers/flat_set.h"
 #include "base/logging.h"
 #include "base/memory/ref_counted.h"
 #include "base/task_scheduler/post_task.h"
+#include "chrome/browser/ui/browser.h"
+#include "chrome/browser/ui/browser_list.h"
+#include "chrome/browser/ui/browser_window.h"
 #include "chromeos/chromeos_switches.h"
 #include "components/arc/arc_bridge_service.h"
 #include "components/arc/instance_holder.h"
 #include "content/public/browser/browser_thread.h"
 #include "ui/aura/window.h"
+#include "ui/compositor/layer.h"
+#include "ui/compositor/layer_owner.h"
+#include "ui/compositor/layer_tree_owner.h"
+#include "ui/gfx/image/image.h"
+#include "ui/gfx/image/image_util.h"
+#include "ui/gfx/native_widget_types.h"
 #include "ui/snapshot/snapshot.h"
+#include "ui/snapshot/snapshot_aura.h"
+#include "ui/wm/core/window_util.h"
 #include "ui/wm/public/activation_client.h"
 
 namespace arc {
@@ -37,6 +50,64 @@ void ScreenshotCallback(
   callback.Run(result);
 }
 
+std::unique_ptr<ui::LayerTreeOwner> CreateLayerTreeForSnapshot(
+    aura::Window* root_window) {
+  using LayerSet = base::flat_set<const ui::Layer*>;
+  LayerSet blocked_layers;
+  for (auto* browser : *BrowserList::GetInstance()) {
+    if (browser->profile()->IsOffTheRecord())
+      blocked_layers.insert(browser->window()->GetNativeWindow()->layer());
+  }
+
+  auto layer_tree_owner = ::wm::RecreateLayersWithClosure(
+      root_window, base::BindRepeating(
+                       [](LayerSet blocked_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 (blocked_layers.count(owner->layer()->parent()))
+                           return nullptr;
+                         if (blocked_layers.count(owner->layer())) {
+                           auto layer = base::MakeUnique<ui::Layer>(
+                               ui::LayerType::LAYER_SOLID_COLOR);
+                           layer->SetBounds(owner->layer()->bounds());
+                           layer->SetColor(SK_ColorBLACK);
+                           return layer;
+                         }
+                         return owner->RecreateLayer();
+                       },
+                       std::move(blocked_layers)));
+
+  // layer_tree_owner cannot be null since we are starting off from the root
+  // window, which could never be incognito.
+  DCHECK(layer_tree_owner);
+
+  auto* root_layer = layer_tree_owner->root();
+  root_window->layer()->Add(root_layer);
+  root_window->layer()->StackAtBottom(root_layer);
+  return layer_tree_owner;
+}
+
+void EncodeAndReturnImage(
+    const ArcVoiceInteractionFrameworkService::CaptureFullscreenCallback&
+        callback,
+    std::unique_ptr<ui::LayerTreeOwner> old_layer_owner,
+    const gfx::Image& image) {
+  old_layer_owner.reset();
+  base::PostTaskWithTraitsAndReplyWithResult(
+      FROM_HERE,
+      base::TaskTraits{base::MayBlock(), base::TaskPriority::USER_BLOCKING},
+      base::Bind(
+          [](const gfx::Image& image) -> std::vector<uint8_t> {
+            std::vector<uint8_t> res;
+            gfx::JPEG1xEncodedDataFromImage(image, 100, &res);
+            return res;
+          },
+          image),
+      callback);
+}
+
 }  // namespace
 
 // static
@@ -136,11 +207,12 @@ void ArcVoiceInteractionFrameworkService::CaptureFullscreen(
   // the screenshot to it.
   aura::Window* window = ash::Shell::GetPrimaryRootWindow();
   DCHECK(window);
-  ui::GrabWindowSnapshotAsyncJPEG(
-      window, gfx::Rect(window->bounds().size()),
-      base::CreateTaskRunnerWithTraits(
-          {base::MayBlock(), base::TaskPriority::USER_BLOCKING}),
-      base::Bind(&ScreenshotCallback, callback));
+
+  auto old_layer_owner = CreateLayerTreeForSnapshot(window);
+  ui::GrabLayerSnapshotAsync(
+      old_layer_owner->root(), gfx::Rect(window->bounds().size()),
+      base::Bind(&EncodeAndReturnImage, callback,
+                 base::Passed(std::move(old_layer_owner))));
 }
 
 void ArcVoiceInteractionFrameworkService::OnMetalayerClosed() {

ui/snapshot/snapshot_aura.cc

@@ -29,13 +29,13 @@ bool GrabWindowSnapshotAura(aura::Window* window,
 }
 
 static void MakeAsyncCopyRequest(
-    aura::Window* window,
+    Layer* layer,
     const gfx::Rect& source_rect,
     const cc::CopyOutputRequest::CopyOutputRequestCallback& callback) {
   std::unique_ptr<cc::CopyOutputRequest> request =
       cc::CopyOutputRequest::CreateBitmapRequest(callback);
   request->set_area(source_rect);
-  window->layer()->RequestCopyOfOutput(std::move(request));
+  layer->RequestCopyOfOutput(std::move(request));
 }
 
 static void FinishedAsyncCopyRequest(
@@ -55,7 +55,7 @@ static void FinishedAsyncCopyRequest(
     // dereference.
     aura::Window* window = tracker->windows()[0];
     MakeAsyncCopyRequest(
-        window, source_rect,
+        window->layer(), source_rect,
         base::Bind(&FinishedAsyncCopyRequest, base::Passed(&tracker),
                    source_rect, callback, retry_count + 1));
     return;
@@ -71,7 +71,7 @@ static void MakeInitialAsyncCopyRequest(
   auto tracker = base::MakeUnique<aura::WindowTracker>();
   tracker->Add(window);
   MakeAsyncCopyRequest(
-      window, source_rect,
+      window->layer(), source_rect,
       base::Bind(&FinishedAsyncCopyRequest, base::Passed(&tracker), source_rect,
                  callback, 0));
 }
@@ -133,6 +133,14 @@ void GrabViewSnapshotAsync(gfx::NativeView view,
   GrabWindowSnapshotAsyncAura(view, source_rect, callback);
 }
 
+void GrabLayerSnapshotAsync(ui::Layer* layer,
+                            const gfx::Rect& source_rect,
+                            const GrabLayerSnapshotCallback& callback) {
+  MakeAsyncCopyRequest(
+      layer, source_rect,
+      base::Bind(&SnapshotAsync::RunCallbackWithCopyOutputResult, callback));
+}
+
 #endif
 
 }  // namespace ui

ui/snapshot/snapshot_aura.h

@@ -9,6 +9,8 @@
 
 namespace ui {
 
+class Layer;
+
 // These functions are identical to those in snapshot.h, except they're
 // guaranteed to read the frame using an Aura CopyOutputRequest and not the
 // native windowing system. source_rect and target_size are in DIP.
@@ -25,6 +27,16 @@ SNAPSHOT_EXPORT void GrabWindowSnapshotAsyncAura(
     const gfx::Rect& source_rect,
     const GrabWindowSnapshotAsyncCallback& callback);
 
+using GrabLayerSnapshotCallback =
+    base::Callback<void(const gfx::Image& snapshot)>;
+
+// Grabs a snapshot of a |layer| and all its descendants.
+// |source_rect| is the bounds of the snapshot content relative to |layer|.
+SNAPSHOT_EXPORT void GrabLayerSnapshotAsync(
+    Layer* layer,
+    const gfx::Rect& source_rect,
+    const GrabLayerSnapshotCallback& callback);
+
 }  // namespace ui
 
 #endif  // UI_SNAPSHOT_SNAPSHOT_AURA_H_

ui/wm/core/window_util.cc

@@ -14,22 +14,25 @@
 
 namespace {
 
-// Invokes RecreateLayer() on all the children of |to_clone|, adding the newly
-// cloned children to |parent|.
+// Invokes |map_func| on all the children of |to_clone|, adding the newly
+// cloned children to |parent|. If |map_func| returns nullptr on
+// the layer owner, all its layer's children will not be cloned.
 //
 // WARNING: It is assumed that |parent| is ultimately owned by a LayerTreeOwner.
-void CloneChildren(ui::Layer* to_clone, ui::Layer* parent) {
+void CloneChildren(ui::Layer* to_clone,
+                   ui::Layer* parent,
+                   const wm::MapLayerFunc& map_func) {
   typedef std::vector<ui::Layer*> Layers;
   // Make a copy of the children since RecreateLayer() mutates it.
   Layers children(to_clone->children());
   for (Layers::const_iterator i = children.begin(); i != children.end(); ++i) {
     ui::LayerOwner* owner = (*i)->owner();
-    ui::Layer* old_layer = owner ? owner->RecreateLayer().release() : NULL;
+    ui::Layer* old_layer = owner ? map_func.Run(owner).release() : NULL;
     if (old_layer) {
       parent->Add(old_layer);
       // RecreateLayer() moves the existing children to the new layer. Create a
       // copy of those.
-      CloneChildren(owner->layer(), old_layer);
+      CloneChildren(owner->layer(), old_layer, map_func);
     }
   }
 }
@@ -132,8 +135,20 @@ aura::Window* GetToplevelWindow(aura::Window* window) {
 
 std::unique_ptr<ui::LayerTreeOwner> RecreateLayers(ui::LayerOwner* root) {
   DCHECK(root->OwnsLayer());
-  auto old_layer = base::MakeUnique<ui::LayerTreeOwner>(root->RecreateLayer());
-  CloneChildren(root->layer(), old_layer->root());
+  return RecreateLayersWithClosure(root, base::Bind([](ui::LayerOwner* owner) {
+                                     return owner->RecreateLayer();
+                                   }));
+}
+
+std::unique_ptr<ui::LayerTreeOwner> RecreateLayersWithClosure(
+    ui::LayerOwner* root,
+    const MapLayerFunc& map_func) {
+  DCHECK(root->OwnsLayer());
+  auto layer = map_func.Run(root);
+  if (!layer)
+    return nullptr;
+  auto old_layer = base::MakeUnique<ui::LayerTreeOwner>(std::move(layer));
+  CloneChildren(root->layer(), old_layer->root(), map_func);
   return old_layer;
 }
 

ui/wm/core/window_util.h

@@ -6,8 +6,10 @@
 #define UI_WM_CORE_WINDOW_UTIL_H_
 
 #include <memory>
+#include <utility>
 #include <vector>
 
+#include "base/callback_forward.h"
 #include "base/compiler_specific.h"
 #include "ui/wm/wm_export.h"
 
@@ -46,6 +48,17 @@ WM_EXPORT aura::Window* GetToplevelWindow(aura::Window* window);
 WM_EXPORT std::unique_ptr<ui::LayerTreeOwner> RecreateLayers(
     ui::LayerOwner* root);
 
+using MapLayerFunc =
+    base::RepeatingCallback<std::unique_ptr<ui::Layer>(ui::LayerOwner*)>;
+
+// Maps |map_func| over each layer of the layer tree and returns a copy of the
+// layer tree. The recursion stops at the level when |map_func| returns nullptr
+// on the owner's layer. MapLayers might return nullptr when |map_func| returns
+// nullptr on the root layer's owner.
+WM_EXPORT std::unique_ptr<ui::LayerTreeOwner> RecreateLayersWithClosure(
+    ui::LayerOwner* root,
+    const MapLayerFunc& map_func);
+
 // Returns a layer tree that mirrors |root|. Used for live window previews. If
 // |sync_bounds| is true, the bounds of all mirror layers except the root are
 // synchronized. See |sync_bounds_| in ui::Layer.

ui/wm/core/window_util_unittest.cc

@@ -6,6 +6,7 @@
 
 #include <memory>
 
+#include "base/bind.h"
 #include "ui/aura/test/aura_test_base.h"
 #include "ui/aura/test/test_windows.h"
 #include "ui/aura/window.h"
@@ -50,4 +51,52 @@ TEST_F(WindowUtilTest, RecreateLayers) {
   window11.reset();
 }
 
+// Test if map_func is correctly executed in RecreateLayerWithClosure.
+TEST_F(WindowUtilTest, RecreateLayersWithClosure) {
+  std::unique_ptr<aura::Window> window1(
+      aura::test::CreateTestWindowWithId(0, NULL));
+  std::unique_ptr<aura::Window> window11(
+      aura::test::CreateTestWindowWithId(1, window1.get()));
+  std::unique_ptr<aura::Window> window12(
+      aura::test::CreateTestWindowWithId(2, window1.get()));
+
+  ASSERT_EQ(2u, window1->layer()->children().size());
+
+  auto tree_empty = wm::RecreateLayersWithClosure(
+      window1.get(),
+      base::BindRepeating(
+          [](const aura::Window* to_filter,
+             ui::LayerOwner* owner) -> std::unique_ptr<ui::Layer> {
+            if (owner->layer() == to_filter->layer())
+              return nullptr;
+            return owner->RecreateLayer();
+          },
+          window1.get()));
+
+  // The root is filtered. RecreateLayersWithClosure should return nullptr.
+  ASSERT_FALSE(tree_empty);
+
+  auto tree = wm::RecreateLayersWithClosure(
+      window1.get(),
+      base::BindRepeating(
+          [](const aura::Window* to_filter,
+             ui::LayerOwner* owner) -> std::unique_ptr<ui::Layer> {
+            if (owner->layer() == to_filter->layer())
+              return nullptr;
+            return owner->RecreateLayer();
+          },
+          window12.get()));
+
+  ASSERT_TRUE(tree);
+  // window12 is filtered out in the above recreation logic.
+  ASSERT_EQ(1u, tree->root()->children().size());
+  // Child layer is new instance.
+  EXPECT_NE(window11->layer(), tree->root()->children()[0]);
+
+  // The original window should have both.
+  ASSERT_EQ(2u, window1->layer()->children().size());
+  EXPECT_EQ(window11->layer(), window1->layer()->children()[0]);
+  EXPECT_EQ(window12->layer(), window1->layer()->children()[1]);
+}
+
 }  // namespace wm