[exo] Integrate ClientTracker and IsClientDestroyed()

This CL adds the ClientTracker to the exo::wayland::Server instance
to ensure that the construction / destruction of clients can be
accurately observed for a given wl_display (which is 1:1 with a
Server instance).

A utility method is exposed to allow code to query whether a
wl_client has started destruction. This allows code to check whether
it is safe to iterate over / query the clients resources.

This replaces more bespoke methods used to achieve similar ends
for the WaylandDisplayObserver.

Bug: 1459067
Change-Id: Iec9bc424b7fdca756d935998f75947b67d628911
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4786809
Reviewed-by: Mitsuru Oshima <oshima@chromium.org>
Code-Coverage: findit-for-me@appspot.gserviceaccount.com <findit-for-me@appspot.gserviceaccount.com>
Commit-Queue: Thomas Lukaszewicz <tluk@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1188432}

components/exo/server/wayland_server_controller.cc

@@ -52,10 +52,26 @@ WaylandServerController::~WaylandServerController() {
   // TODO(https://crbug.com/1124106): Investigate if we can eliminate Shutdown
   // methods.
   display_->Shutdown();
+  wayland::Server::SetServerGetter(base::NullCallback());
   DCHECK_EQ(g_instance, this);
   g_instance = nullptr;
 }
 
+wayland::Server* WaylandServerController::GetServerForDisplay(
+    wl_display* display) {
+  if (default_server_ && default_server_->GetWaylandDisplay() == display) {
+    return default_server_.get();
+  }
+
+  for (const auto& pair : on_demand_servers_) {
+    if (pair.second->GetWaylandDisplay() == display) {
+      return pair.second.get();
+    }
+  }
+
+  return nullptr;
+}
+
 WaylandServerController::WaylandServerController(
     std::unique_ptr<DataExchangeDelegate> data_exchange_delegate,
     std::unique_ptr<NotificationSurfaceManager> notification_surface_manager,
@@ -74,6 +90,8 @@ WaylandServerController::WaylandServerController(
   default_server_->StartWithDefaultPath(base::BindOnce([](bool success) {
     DCHECK(success) << "Failed to start the default wayland server.";
   }));
+  wayland::Server::SetServerGetter(base::BindRepeating(
+      &WaylandServerController::GetServerForDisplay, base::Unretained(this)));
 }
 
 void WaylandServerController::ListenOnSocket(

components/exo/server/wayland_server_controller.h

@@ -14,6 +14,8 @@
 #include "components/exo/security_delegate.h"
 #include "components/exo/wayland/server.h"
 
+struct wl_display;
+
 namespace exo {
 
 namespace wayland {
@@ -49,6 +51,9 @@ class WaylandServerController {
 
   ~WaylandServerController();
 
+  // Gets the Server instance for the `display` if it exists.
+  wayland::Server* GetServerForDisplay(wl_display* display);
+
   InputMethodSurfaceManager* input_method_surface_manager() {
     return display_->input_method_surface_manager();
   }

components/exo/wayland/clients/security_delegate_binding_test.cc

@@ -30,7 +30,7 @@ class SecurityDelegateBindingTest : public test::WaylandServerTest {
   void SetUp() override {
     WaylandServerTest::SetUp();
     server_security_delegate_ =
-        GetSecurityDelegate(server_->GetWaylandDisplayForTesting());
+        GetSecurityDelegate(server_->GetWaylandDisplay());
     ASSERT_NE(server_security_delegate_, nullptr);
   }
 

components/exo/wayland/server.cc

@@ -59,6 +59,7 @@
 #include "build/build_config.h"
 #include "components/exo/display.h"
 #include "components/exo/security_delegate.h"
+#include "components/exo/wayland/client_tracker.h"
 #include "components/exo/wayland/content_type.h"
 #include "components/exo/wayland/overlay_prioritizer.h"
 #include "components/exo/wayland/serial_tracker.h"
@@ -135,6 +136,9 @@ const char kWaylandSocketGroup[] = "wayland";
 // (see `man 2 listen`).
 constexpr int kMaxPendingConnections = 128;
 
+// Callback used to find a Server instance for a given wl_display.
+Server::ServerGetter g_server_getter;
+
 bool IsDrmAtomicAvailable() {
 #if BUILDFLAG(IS_OZONE)
   auto& host_properties =
@@ -269,6 +273,8 @@ Server::Server(Display* display,
 
   wl_display_.reset(wl_display_create());
   SetSecurityDelegate(wl_display_.get(), security_delegate_.get());
+
+  client_tracker_ = std::make_unique<ClientTracker>(wl_display_.get());
 }
 
 void Server::Initialize() {
@@ -450,6 +456,17 @@ std::unique_ptr<Server> Server::Create(
   return server;
 }
 
+// static.
+Server* Server::GetServerForDisplay(wl_display* display) {
+  return g_server_getter ? g_server_getter.Run(display) : nullptr;
+}
+
+// static.
+void Server::SetServerGetter(Server::ServerGetter server_getter) {
+  CHECK(!server_getter || !g_server_getter);
+  g_server_getter = std::move(server_getter);
+}
+
 void Server::StartWithDefaultPath(StartCallback callback) {
   if (!Open()) {
     std::move(callback).Run(/*success=*/false);
@@ -513,6 +530,10 @@ wl_resource* Server::GetOutputResource(wl_client* client, int64_t display_id) {
   return iter->second.get()->GetOutputResourceForClient(client);
 }
 
+bool Server::IsClientDestroyed(wl_client* client) const {
+  return client_tracker_->IsClientDestroyed(client);
+}
+
 void Server::AddWaylandOutput(int64_t id,
                               std::unique_ptr<WaylandDisplayOutput> output) {
   outputs_.insert(std::make_pair(id, std::move(output)));

components/exo/wayland/server.h

@@ -27,6 +27,7 @@ class Display;
 
 namespace wayland {
 
+class ClientTracker;
 class SerialTracker;
 class UiControls;
 struct WaylandDataDeviceManager;
@@ -45,6 +46,7 @@ class WaylandWatcher;
 // requests are dispatched into the given Exosphere display.
 class Server : public display::DisplayObserver {
  public:
+  using ServerGetter = base::RepeatingCallback<Server*(wl_display*)>;
   using StartCallback = base::OnceCallback<void(bool)>;
 
   Server(Display* display, std::unique_ptr<SecurityDelegate> security_delegate);
@@ -63,6 +65,12 @@ class Server : public display::DisplayObserver {
       Display* display,
       std::unique_ptr<SecurityDelegate> security_delegate);
 
+  // Gets the Server instance for a given wl_display.
+  static Server* GetServerForDisplay(wl_display* display);
+
+  // Sets the callback used to find the Server instance for a given wl_display.
+  static void SetServerGetter(ServerGetter server_getter);
+
   void StartWithDefaultPath(StartCallback callback);
   void StartWithFdAsync(base::ScopedFD fd, StartCallback callback);
 
@@ -93,18 +101,17 @@ class Server : public display::DisplayObserver {
   wl_resource* GetOutputResource(wl_client* client, int64_t display_id);
 
   Display* GetDisplay() { return display_; }
+  wl_display* GetWaylandDisplay() { return wl_display_.get(); }
 
-  // Public version of the protected accessor below, to be used in tests.
-  wl_display* GetWaylandDisplayForTesting() const {
-    return GetWaylandDisplay();
-  }
+  // Returns whether a client associated with this server has started
+  // destruction.
+  bool IsClientDestroyed(wl_client* client) const;
 
  protected:
   friend class UiControls;
   friend class WestonTest;
   void AddWaylandOutput(int64_t id,
                         std::unique_ptr<WaylandDisplayOutput> output);
-  wl_display* GetWaylandDisplay() const { return wl_display_.get(); }
 
  private:
   friend class ScopedEventDispatchDisabler;
@@ -132,6 +139,7 @@ class Server : public display::DisplayObserver {
   std::unique_ptr<WaylandXdgShell> xdg_shell_data_;
   std::unique_ptr<WaylandRemoteShellData> remote_shell_data_;
   std::unique_ptr<UiControls> ui_controls_holder_;
+  std::unique_ptr<ClientTracker> client_tracker_;
 };
 
 }  // namespace wayland

components/exo/wayland/server_unittest.cc

@@ -68,7 +68,7 @@ TEST_F(ServerTest, SecurityDelegateAssociation) {
 
   auto server = CreateServer(std::move(security_delegate));
 
-  EXPECT_EQ(GetSecurityDelegate(server->GetWaylandDisplayForTesting()),
+  EXPECT_EQ(GetSecurityDelegate(server->GetWaylandDisplay()),
             security_delegate_ptr);
 }
 
@@ -101,7 +101,7 @@ TEST_F(ServerTest, StartFd) {
   EXPECT_NE(client_display, nullptr);
 
   wl_list* all_clients =
-      wl_display_get_client_list(server->GetWaylandDisplayForTesting());
+      wl_display_get_client_list(server->GetWaylandDisplay());
   ASSERT_FALSE(wl_list_empty(all_clients));
   wl_client* client = wl_client_from_link(all_clients->next);
 
@@ -126,7 +126,7 @@ TEST_F(ServerTest, Dispatch) {
   client_thread.Start();
 
   TestListener client_creation_listener;
-  wl_display_add_client_created_listener(server->GetWaylandDisplayForTesting(),
+  wl_display_add_client_created_listener(server->GetWaylandDisplay(),
                                          &client_creation_listener.listener);
 
   base::Lock lock;
@@ -154,7 +154,7 @@ TEST_F(ServerTest, Dispatch) {
   }
 
   wl_list* all_clients =
-      wl_display_get_client_list(server->GetWaylandDisplayForTesting());
+      wl_display_get_client_list(server->GetWaylandDisplay());
   ASSERT_FALSE(wl_list_empty(all_clients));
   wl_client* client = wl_client_from_link(all_clients->next);
 

components/exo/wayland/server_util.cc

@@ -9,6 +9,7 @@
 #include "base/containers/flat_map.h"
 #include "base/time/time.h"
 #include "components/exo/data_offer.h"
+#include "components/exo/wayland/server.h"
 #include "ui/display/display.h"
 #include "ui/gfx/geometry/insets.h"
 #include "ui/gfx/geometry/rect.h"
@@ -86,5 +87,17 @@ SecurityDelegate* GetSecurityDelegate(wl_client* client) {
   return GetSecurityDelegate(wl_client_get_display(client));
 }
 
+bool IsClientDestroyed(wl_client* client) {
+  CHECK(client);
+
+  // There should always be a display associated with a client and display will
+  // always outlive the wl_client object.
+  wl_display* display = wl_client_get_display(client);
+  CHECK(display);
+
+  Server* server = Server::GetServerForDisplay(display);
+  return server ? server->IsClientDestroyed(client) : true;
+}
+
 }  // namespace wayland
 }  // namespace exo

components/exo/wayland/server_util.h

@@ -77,6 +77,11 @@ SecurityDelegate* GetSecurityDelegate(wl_display* display);
 // connected to.
 SecurityDelegate* GetSecurityDelegate(wl_client* client);
 
+// Returns whether a client has initiated destruction. After destruction begins
+// resources associated with the client start to be freed and there is a risk of
+// UAFs in querying for client resources (see crbug.com/1433187).
+bool IsClientDestroyed(wl_client* client);
+
 }  // namespace wayland
 }  // namespace exo
 

components/exo/wayland/test/server_util.cc

@@ -21,7 +21,7 @@ wl_resource* LookUpResource(Server* server, const ResourceKey& key) {
 
   wl_client* client = nullptr;
   wl_list* all_clients =
-      wl_display_get_client_list(server->GetWaylandDisplayForTesting());
+      wl_display_get_client_list(server->GetWaylandDisplay());
 
   auto find_closure = [](struct wl_resource* resource, void* data) {
     IteratorData* iterator_data = static_cast<IteratorData*>(data);

components/exo/wayland/test/wayland_server_test.cc

@@ -12,13 +12,22 @@
 #include "base/run_loop.h"
 #include "base/test/bind.h"
 #include "components/exo/security_delegate.h"
+#include "components/exo/wayland/server.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace exo::wayland::test {
 
-WaylandServerTest::WaylandServerTest() = default;
+WaylandServerTest::WaylandServerTest() {
+  Server::SetServerGetter(base::BindLambdaForTesting([&](wl_display* display) {
+    // Currently tests run with a single Server instance.
+    EXPECT_EQ(display, server_->GetWaylandDisplay());
+    return server_.get();
+  }));
+}
 
-WaylandServerTest::~WaylandServerTest() = default;
+WaylandServerTest::~WaylandServerTest() {
+  Server::SetServerGetter(base::NullCallback());
+}
 
 void WaylandServerTest::SetUp() {
   WaylandServerTestBase::SetUp();

components/exo/wayland/test/wlcs/wlcs_helpers.cc

@@ -64,7 +64,7 @@ int ScopedWlcsServer::AddClient() {
   }
 
   struct wl_client* client =
-      wl_client_create(Get()->GetWaylandDisplayForTesting(), server_fd.get());
+      wl_client_create(Get()->GetWaylandDisplay(), server_fd.get());
   if (!client) {
     return -1;
   }

components/exo/wayland/wayland_display_observer.cc

@@ -29,22 +29,9 @@ WaylandDisplayObserver::~WaylandDisplayObserver() = default;
 
 WaylandDisplayHandler::WaylandDisplayHandler(WaylandDisplayOutput* output,
                                              wl_resource* output_resource)
-    : output_(output), output_resource_(output_resource) {
-  // At construction time the client object is guaranteed to exist.
-  wl_client* client = wl_resource_get_client(output_resource_);
-  CHECK(client);
-  client_destroy_listener_.listener.notify =
-      &WaylandDisplayHandler::OnClientDestroyed;
-  wl_client_add_destroy_listener(client, &client_destroy_listener_.listener);
-}
+    : output_(output), output_resource_(output_resource) {}
 
 WaylandDisplayHandler::~WaylandDisplayHandler() {
-  // Remove the listener to cover the case where the client outlives the
-  // handler.
-  if (!client_destroy_listener_.notified) {
-    wl_list_remove(&client_destroy_listener_.listener.link);
-  }
-
   ash::Shell::Get()->RemoveShellObserver(this);
   for (auto& obs : observers_) {
     obs.OnOutputDestroyed();
@@ -158,14 +145,6 @@ void WaylandDisplayHandler::UnsetXdgOutputResource() {
   xdg_output_resource_ = nullptr;
 }
 
-bool WaylandDisplayHandler::IsClientDestroyedForTesting() const {
-  return client_destroy_listener_.notified;
-}
-
-AuraOutputManager* WaylandDisplayHandler::GetAuraOutputManagerForTesting() {
-  return GetAuraOutputManager();
-}
-
 void WaylandDisplayHandler::XdgOutputSendLogicalPosition(
     const gfx::Point& position) {
   zxdg_output_v1_send_logical_position(xdg_output_resource_, position.x(),
@@ -185,15 +164,6 @@ void WaylandDisplayHandler::XdgOutputSendDescription(const std::string& desc) {
   zxdg_output_v1_send_description(xdg_output_resource_, desc.c_str());
 }
 
-// static.
-void WaylandDisplayHandler::OnClientDestroyed(struct wl_listener* listener,
-                                              void* data) {
-  ClientDestroyListener* client_destroy_listener = wl_container_of(
-      listener, /*sample=*/client_destroy_listener, /*member=*/listener);
-  client_destroy_listener->notified = true;
-  wl_list_remove(&client_destroy_listener->listener.link);
-}
-
 bool WaylandDisplayHandler::SendDisplayMetrics(const display::Display& display,
                                                uint32_t changed_metrics) {
   if (!output_resource_) {
@@ -278,14 +248,6 @@ void WaylandDisplayHandler::OnOutputDestroyed() {
 }
 
 AuraOutputManager* WaylandDisplayHandler::GetAuraOutputManager() {
-  // If the client has begun destruction avoid attempting to access the client's
-  // AuraOutputManager instance as libwayland may have freed the object's memory
-  // but not yet updated the data structures used to find the object (see
-  // crbug.com/1433187).
-  if (client_destroy_listener_.notified) {
-    return nullptr;
-  }
-
   wl_client* client = wl_resource_get_client(output_resource_);
   CHECK(client);
   return AuraOutputManager::Get(client);