Skip closing ShellSurfaceBase if capture mode is active.

It also updates `ash::~Shell` to fix a pre-existing issue that global
instance of `CaptureModeController` sometimes get accessed its
destruction.

Without this change what can happen is that:
- `views::MenuHost` has a capture.
- Chrome starts shutdown and `ash::~Shell` starts. First
  `CaptureModeController` is destroyed.
= Then `ash::Shell::CloseAllRootWindowChildWindows` is called which
causes the `MenuHost` to lose capture and call
`MenuHost::OnMouseCaptureLost` which calls
`ViewsDelegate::ShouldCloseMenuIfMouseCaptureLost()` which then tries to
access `CaptureModeController` after its destruction.

Test: exo_unittests, ash_unittests
Bug: 1383272
Change-Id: I149c51a6b0bd2f95014f043a9f59aa9a03711289
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4914800
Reviewed-by: Ahmed Fakhry <afakhry@chromium.org>
Reviewed-by: Eliot Courtney <edcourtney@chromium.org>
Commit-Queue: Yuta Hijikata <ythjkt@chromium.org>
Reviewed-by: Mitsuru Oshima <oshima@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1216759}

ash/capture_mode/test_capture_mode_delegate.cc

@@ -147,6 +147,8 @@ void TestCaptureModeDelegate::BindAudioStreamFactory(
     mojo::PendingReceiver<media::mojom::AudioStreamFactory> receiver) {}
 
 void TestCaptureModeDelegate::OnSessionStateChanged(bool started) {
+  is_session_active_ = started;
+
   if (on_session_state_changed_callback_)
     std::move(on_session_state_changed_callback_).Run();
 }

ash/capture_mode/test_capture_mode_delegate.h

@@ -32,6 +32,8 @@ class TestCaptureModeDelegate : public CaptureModeDelegate {
   TestCaptureModeDelegate& operator=(const TestCaptureModeDelegate&) = delete;
   ~TestCaptureModeDelegate() override;
 
+  bool is_session_active() const { return is_session_active_; }
+
   recording::RecordingServiceTestApi* recording_service() const {
     return recording_service_.get();
   }
@@ -139,6 +141,7 @@ class TestCaptureModeDelegate : public CaptureModeDelegate {
   base::ScopedTempDir fake_downloads_dir_;
   base::OnceClosure on_session_state_changed_callback_;
   base::OnceClosure on_recording_started_callback_;
+  bool is_session_active_ = false;
   bool is_allowed_by_dlp_ = true;
   bool is_allowed_by_policy_ = true;
   bool should_save_after_dlp_check_ = true;

ash/shell.cc

@@ -282,6 +282,7 @@
 #include "ui/views/widget/native_widget_aura.h"
 #include "ui/views/widget/widget.h"
 #include "ui/wm/core/accelerator_filter.h"
+#include "ui/wm/core/capture_controller.h"
 #include "ui/wm/core/compound_event_filter.h"
 #include "ui/wm/core/focus_controller.h"
 #include "ui/wm/core/shadow_controller.h"
@@ -865,6 +866,17 @@ Shell::~Shell() {
   projector_controller_.reset();
   game_dashboard_controller_.reset();
 
+  // This must be called before `capture_mode_controller_` is destroyed. Note
+  // that 'capture' in `CaptureModeController` means 'screenshot capture, while
+  // 'capture' in `wm::CaptureController` means 'input capture'. Some windows
+  // like popup windows close themselves when losing 'input capture' but if
+  // 'screenshot capture' is in progress, they do not close themselves. For this
+  // reason, change of 'input capture' can cause a call to
+  // `CaptureModeController::Get()`. By calling `PrepareForShutdown()` here, it
+  // prevents `CaptureModeController::Get()` from being called after the object
+  // is destroyed.
+  wm::CaptureController::Get()->PrepareForShutdown();
+
   // This must be destroyed before deleting all the windows below in
   // `CloseAllRootWindowChildWindows()`, since shutting down the session will
   // need to access those windows and it will be a UAF.

ash/test/ash_test_views_delegate.cc

@@ -5,6 +5,7 @@
 #include "ash/test/ash_test_views_delegate.h"
 
 #include "ash/accelerators/accelerator_controller_impl.h"
+#include "ash/capture_mode/capture_mode_test_util.h"
 #include "ash/shell.h"
 #include "base/task/single_thread_task_runner.h"
 #include "chromeos/ui/frame/frame_utils.h"
@@ -48,4 +49,11 @@ AshTestViewsDelegate::ProcessAcceleratorWhileMenuShowing(
   return views::ViewsDelegate::ProcessMenuAcceleratorResult::LEAVE_MENU_OPEN;
 }
 
+bool AshTestViewsDelegate::ShouldCloseMenuIfMouseCaptureLost() const {
+  // This is the same behaviour as `ChromeViewsDelegate`.
+  auto* capture_mode_test_delegate = GetTestDelegate();
+  CHECK(capture_mode_test_delegate);
+  return !capture_mode_test_delegate->is_session_active();
+}
+
 }  // namespace ash

ash/test/ash_test_views_delegate.h

@@ -21,14 +21,17 @@ class AshTestViewsDelegate : public views::TestViewsDelegate {
 
   ~AshTestViewsDelegate() override;
 
-  // Overriden from TestViewsDelegate.
+  // views::TestViewsDelegate:
   void OnBeforeWidgetInit(
       views::Widget::InitParams* params,
       views::internal::NativeWidgetDelegate* delegate) override;
   views::TestViewsDelegate::ProcessMenuAcceleratorResult
   ProcessAcceleratorWhileMenuShowing(
       const ui::Accelerator& accelerator) override;
 
+  // views::ViewsDelegate:
+  bool ShouldCloseMenuIfMouseCaptureLost() const override;
+
   void set_close_menu_accelerator(const ui::Accelerator& accelerator) {
     close_menu_accelerator_ = accelerator;
   }

components/exo/shell_surface_base.cc

@@ -66,6 +66,7 @@
 #include "ui/gfx/geometry/vector2d_conversions.h"
 #include "ui/views/accessibility/view_accessibility.h"
 #include "ui/views/controls/menu/menu_config.h"
+#include "ui/views/views_delegate.h"
 #include "ui/views/widget/tooltip_manager.h"
 #include "ui/views/widget/widget.h"
 #include "ui/wm/core/shadow_controller.h"
@@ -1325,6 +1326,13 @@ void ShellSurfaceBase::OnCaptureChanged(aura::Window* lost_capture,
   if (lost_capture != widget_->GetNativeWindow() || !is_popup_)
     return;
 
+  // If the capture mode is active, do not close the popup to include it in a
+  // screenshot.
+  if (!views::ViewsDelegate::GetInstance()
+           ->ShouldCloseMenuIfMouseCaptureLost()) {
+    return;
+  }
+
   WMHelper::GetInstance()->GetCaptureClient()->RemoveObserver(this);
 
   // Fast return for a simple case: if `lost_capture` is the parent of

components/exo/shell_surface_unittest.cc

@@ -7,6 +7,7 @@
 #include <sstream>
 #include <vector>
 
+#include "ash/capture_mode/capture_mode_test_util.h"
 #include "ash/constants/app_types.h"
 #include "ash/frame/non_client_frame_view_ash.h"
 #include "ash/public/cpp/test/shell_test_api.h"
@@ -2329,6 +2330,43 @@ TEST_F(ShellSurfaceTest, PopupWithInputRegion) {
   }
 }
 
+// Test that popup does not close when trying to take a screenshot.
+TEST_F(ShellSurfaceTest, PopupWithCaptureMode) {
+  // Setup popup_shell_surface.
+  constexpr gfx::Size kBufferSize(256, 256);
+  auto shell_surface =
+      test::ShellSurfaceBuilder(kBufferSize).BuildShellSurface();
+  auto popup_buffer = std::make_unique<Buffer>(
+      exo_test_helper()->CreateGpuMemoryBuffer(kBufferSize));
+  auto popup_surface = std::make_unique<Surface>();
+  popup_surface->Attach(popup_buffer.get());
+  std::unique_ptr<ShellSurface> popup_shell_surface(CreatePopupShellSurface(
+      popup_surface.get(), shell_surface.get(), gfx::Point(50, 50)));
+  popup_shell_surface->Grab();
+  popup_surface->Commit();
+
+  bool closed = false;
+  auto callback =
+      base::BindRepeating([](bool* closed) { *closed = true; }, &closed);
+  popup_shell_surface->set_close_callback(callback);
+
+  // This simulates enabling (screenshot) capture mode.
+  ash::GetTestDelegate()->OnSessionStateChanged(true);
+  popup_shell_surface->OnCaptureChanged(
+      popup_shell_surface->GetWidget()->GetNativeWindow(), nullptr);
+  // With (screenshot) capture mode on, losing capture should not close the
+  // shell surface.
+  EXPECT_FALSE(closed);
+
+  // This simulates ending (screenshot) capture mode.
+  ash::GetTestDelegate()->OnSessionStateChanged(false);
+  popup_shell_surface->OnCaptureChanged(
+      popup_shell_surface->GetWidget()->GetNativeWindow(), nullptr);
+  // With (screenshot) capture mode off, losing capture should close the shell
+  // surface.
+  EXPECT_TRUE(closed);
+}
+
 TEST_F(ShellSurfaceTest, PopupWithInvisibleParent) {
   // Invisible main window.
   std::unique_ptr<ShellSurface> root_shell_surface =

ui/wm/core/capture_controller.cc

@@ -42,10 +42,20 @@ void CaptureController::Detach(aura::Window* root) {
   aura::client::SetCaptureClient(root, nullptr);
 }
 
+void CaptureController::PrepareForShutdown() {
+  DCHECK(!destroying_);
+  SetCapture(nullptr);
+  destroying_ = true;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // CaptureController, aura::client::CaptureClient implementation:
 
 void CaptureController::SetCapture(aura::Window* new_capture_window) {
+  if (destroying_) {
+    return;
+  }
+
   if (capture_window_ == new_capture_window)
     return;
 

ui/wm/core/capture_controller.h

@@ -40,6 +40,9 @@ class COMPONENT_EXPORT(UI_WM) CaptureController
   // Removes |root| from the list of root windows notified when capture changes.
   void Detach(aura::Window* root);
 
+  // Resets the current capture window and prevents it from being set again.
+  void PrepareForShutdown();
+
   // Returns true if this CaptureController is installed on at least one
   // root window.
   bool is_active() const { return !delegates_.empty(); }
@@ -57,6 +60,10 @@ class COMPONENT_EXPORT(UI_WM) CaptureController
 
   static CaptureController* instance_;
 
+  // Set to true by `PrepareForShutdown()`. If this is true, `SetCapture()` will
+  // have no effect.
+  bool destroying_ = false;
+
   // The current capture window. NULL if there is no capture window.
   raw_ptr<aura::Window> capture_window_;
 

ui/wm/core/capture_controller_unittest.cc

@@ -219,6 +219,21 @@ TEST_F(CaptureControllerTest, ReparentedWhileCaptured) {
   EXPECT_FALSE(delegate2->HasNativeCapture());
 }
 
+// Test that a call to `PrepareForShutdown()` releases capture from capture
+// window and also prevents capture window from being set afterwards.
+TEST_F(CaptureControllerTest, PrepareForShutdown) {
+  aura::Window* w = CreateNormalWindow(1, root_window(), nullptr);
+  w->SetCapture();
+  EXPECT_EQ(CaptureController::Get()->GetCaptureWindow(), w);
+
+  CaptureController::Get()->PrepareForShutdown();
+  EXPECT_EQ(CaptureController::Get()->GetCaptureWindow(), nullptr);
+
+  w->SetCapture();
+  // Once `PrepareForShutdown()` is called, capture window cannot be set.
+  EXPECT_EQ(CaptureController::Get()->GetCaptureWindow(), nullptr);
+}
+
 // A delegate that deletes a window on scroll cancel gesture event.
 class GestureEventDeleteWindowOnScrollEnd
     : public aura::test::TestWindowDelegate {