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

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

Index: components/mus/public/cpp/lib/window_tree_client_impl.cc
diff --git a/components/mus/public/cpp/lib/window_tree_client_impl.cc b/components/mus/public/cpp/lib/window_tree_client_impl.cc
index 2fb60f3fe58f958ca93f24e48cd7cc3c30ee923b..f28f40dbea70ada3b36f87ac1c162f9ed8cb9b7f 100644
--- a/components/mus/public/cpp/lib/window_tree_client_impl.cc
+++ b/components/mus/public/cpp/lib/window_tree_client_impl.cc
@@ -120,6 +120,42 @@ WindowTreeConnection* WindowTreeConnection::CreateForWindowManager(
   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,
@@ -129,7 +165,7 @@ WindowTreeClientImpl::WindowTreeClientImpl(
       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),
@@ -197,7 +233,7 @@ void WindowTreeClientImpl::WaitForEmbed() {
 
 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
@@ -280,10 +316,11 @@ 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);
 }
@@ -292,7 +329,7 @@ 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)));
@@ -405,10 +442,10 @@ void WindowTreeClientImpl::AttachSurface(
 }
 
 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));
@@ -864,7 +901,7 @@ void WindowTreeClientImpl::OnWindowDeleted(Id window_id) {
 }
 
 Window* WindowTreeClientImpl::GetCaptureWindow() {
-  return capture_window_;
+  return capture_window_observer_->window();
 }
 
 void WindowTreeClientImpl::OnWindowVisibilityChanged(Id window_id,