Candidate fix for observe-window-after-free error in WindowTreeClientImpl::LocalSetCapture

components/mus/public/cpp/lib/window_tree_client_impl.cc

 // Copyright 2014 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 "components/mus/public/cpp/lib/window_tree_client_impl.h"
 
 #include <stddef.h>
 
 #include <string>
 #include <utility>
@@ skipping 102 matching lines @@
     mojo::InterfaceRequest<mojom::WindowTreeClient> request,
     CreateType create_type,
     WindowManagerDelegate* window_manager_delegate) {
   WindowTreeClientImpl* client = new WindowTreeClientImpl(
       delegate, window_manager_delegate, std::move(request));
   if (create_type == CreateType::WAIT_FOR_EMBED)
     client->WaitForEmbed();
   return client;
 }
 
+class WindowTreeClientImpl::MusWindowObserver : public mus::WindowObserver {
+ public:
+  explicit MusWindowObserver(Window* window)
+      : window_(window) {}
+
+  ~MusWindowObserver() override {
+    if (window_)
+      window_->RemoveObserver(this);
+  }
+
+  void SetWindow(Window* window) {
+    if (window_ == window)
+      return;
+
+    if (window_)
+      window_->RemoveObserver(this);
+    window_ = window;
+    if (window_)
+      window_->AddObserver(this);
+  }
+
+  // mus::WindowObserver:
+  void OnWindowDestroying(Window* window) override {
+    DCHECK_EQ(window_, window);
+    window_->RemoveObserver(this);
+    window_ = nullptr;
+  }
+
+  Window* window() { return window_; }
+
+ private:
+  Window* window_;
+
+  DISALLOW_COPY_AND_ASSIGN(MusWindowObserver);
+};
+
 WindowTreeClientImpl::WindowTreeClientImpl(
     WindowTreeDelegate* delegate,
     WindowManagerDelegate* window_manager_delegate,
     mojo::InterfaceRequest<mojom::WindowTreeClient> request)
     : connection_id_(0),
       next_window_id_(1),
       next_change_id_(1),
       delegate_(delegate),
       window_manager_delegate_(window_manager_delegate),
-      capture_window_(nullptr),
+      capture_window_observer_(new MusWindowObserver(nullptr)),
       focused_window_(nullptr),
       binding_(this),
       tree_(nullptr),
       delete_on_no_roots_(true),
       in_destructor_(false),
       cursor_location_memory_(nullptr),
       weak_factory_(this) {
   // Allow for a null request in tests.
   if (request.is_pending())
     binding_.Bind(std::move(request));
@@ skipping 47 matching lines @@
 
 void WindowTreeClientImpl::WaitForEmbed() {
   DCHECK(roots_.empty());
   // OnEmbed() is the first function called.
   binding_.WaitForIncomingMethodCall();
   // TODO(sky): deal with pipe being closed before we get OnEmbed().
 }
 
 void WindowTreeClientImpl::DestroyWindow(Window* window) {
   // TODO(jonross): Also clear the focused window (crbug.com/611983)
-  if (window == capture_window_) {
+  if (window == capture_window_observer_->window()) {
     InFlightCaptureChange reset_change(this, nullptr);
     ApplyServerChangeToExistingInFlightChange(reset_change);
     // Normally just updating the queued changes is sufficient. However since
     // |window| is being destroyed, it will not be possible to notify its
     // observers
     // of the lost capture. Update local state now.
     LocalSetCapture(nullptr);
   }
   DCHECK(tree_);
   const uint32_t change_id = ScheduleInFlightChange(base::WrapUnique(
@@ skipping 62 matching lines @@
   const uint32_t change_id = ScheduleInFlightChange(
       base::WrapUnique(new InFlightBoundsChange(window, old_bounds)));
   tree_->SetWindowBounds(change_id, server_id(window),
                          mojo::Rect::From(bounds));
 }
 
 void WindowTreeClientImpl::SetCapture(Window* window) {
   // In order for us to get here we had to have exposed a window, which implies
   // we got a connection.
   DCHECK(tree_);
-  if (capture_window_ == window)
+  if (capture_window_observer_->window() == window)
     return;
   const uint32_t change_id = ScheduleInFlightChange(
-      base::WrapUnique(new InFlightCaptureChange(this, capture_window_)));
+      base::WrapUnique(
+          new InFlightCaptureChange(this, capture_window_observer_->window())));
   tree_->SetCapture(change_id, server_id(window));
   LocalSetCapture(window);
 }
 
 void WindowTreeClientImpl::ReleaseCapture(Window* window) {
   // In order for us to get here we had to have exposed a window, which implies
   // we got a connection.
   DCHECK(tree_);
-  if (capture_window_ != window)
+  if (capture_window_observer_->window() != window)
     return;
   const uint32_t change_id = ScheduleInFlightChange(
       base::WrapUnique(new InFlightCaptureChange(this, window)));
   tree_->ReleaseCapture(change_id, server_id(window));
   LocalSetCapture(nullptr);
 }
 
 void WindowTreeClientImpl::SetClientArea(
     Id window_id,
     const gfx::Insets& client_area,
@@ skipping 92 matching lines @@
 void WindowTreeClientImpl::AttachSurface(
     Id window_id,
     mojom::SurfaceType type,
     mojo::InterfaceRequest<mojom::Surface> surface,
     mojom::SurfaceClientPtr client) {
   DCHECK(tree_);
   tree_->AttachSurface(window_id, type, std::move(surface), std::move(client));
 }
 
 void WindowTreeClientImpl::LocalSetCapture(Window* window) {
-  if (capture_window_ == window)
+  if (capture_window_observer_->window() == window)
     return;
-  Window* lost_capture = capture_window_;
-  capture_window_ = window;
+  Window* lost_capture = capture_window_observer_->window();
+  capture_window_observer_->SetWindow(window);
   if (lost_capture) {
     FOR_EACH_OBSERVER(WindowObserver, *WindowPrivate(lost_capture).observers(),
                       OnWindowLostCapture(lost_capture));
   }
 }
 
 void WindowTreeClientImpl::LocalSetFocus(Window* focused) {
   Window* blurred = focused_window_;
   // Update |focused_window_| before calling any of the observers, so that the
   // observers get the correct result from calling |Window::HasFocus()|,
@@ skipping 435 matching lines @@
     WindowPrivate(window).LocalReorder(relative_window, direction);
 }
 
 void WindowTreeClientImpl::OnWindowDeleted(Id window_id) {
   Window* window = GetWindowByServerId(window_id);
   if (window)
     WindowPrivate(window).LocalDestroy();
 }
 
 Window* WindowTreeClientImpl::GetCaptureWindow() {
-  return capture_window_;
+  return capture_window_observer_->window();
 }
 
 void WindowTreeClientImpl::OnWindowVisibilityChanged(Id window_id,
                                                      bool visible) {
   Window* window = GetWindowByServerId(window_id);
   if (!window)
     return;
 
   InFlightVisibleChange new_change(window, visible);
   if (ApplyServerChangeToExistingInFlightChange(new_change))
@@ skipping 231 matching lines @@
 
 void WindowTreeClientImpl::SetUnderlaySurfaceOffsetAndExtendedHitArea(
     Window* window,
     const gfx::Vector2d& offset,
     const gfx::Insets& hit_area) {
   window_manager_internal_client_->SetUnderlaySurfaceOffsetAndExtendedHitArea(
       server_id(window), offset.x(), offset.y(), mojo::Insets::From(hit_area));
 }
 
 }  // namespace mus