xdg-foreign-v2: Keep invalid imported objects alive (#14166)

* xdg-foreign-v2: validate toplevel surfaces correctly before exporting or parenting

* xdg-foreign-v2: Keep invalid imported objects alive

2 files changed, 52 insertions, 25 deletions

diff --git a/src/protocols/XDGForeignV2.cpp b/src/protocols/XDGForeignV2.cpp
index f67cb86c..198fceb3 100644
--- a/src/protocols/XDGForeignV2.cpp
+++ b/src/protocols/XDGForeignV2.cpp
@@ -46,6 +46,19 @@ std::string_view CXDGExportedResourceV2::handle() const {
 
 CXDGForeignExporterProtocolV2::CXDGForeignExporterProtocolV2(const wl_interface* iface, const int& ver, const std::string& name) : IWaylandProtocol(iface, ver, name) {}
 
+static SP<CXDGToplevelResource> xdgToplevelFromWlSurface(wl_resource* surface) {
+    const auto wlSurf = CWLSurfaceResource::fromResource(surface);
+
+    if (!wlSurf || !wlSurf->m_role || wlSurf->m_role->role() != SURFACE_ROLE_XDG_SHELL)
+        return nullptr;
+
+    const auto xdgSurfResource = sc<CXDGSurfaceRole*>(wlSurf->m_role.get())->m_xdgSurface.lock();
+    if (!xdgSurfResource || xdgSurfResource->m_toplevel.expired())
+        return nullptr;
+
+    return xdgSurfResource->m_toplevel.lock();
+}
+
 void CXDGForeignExporterProtocolV2::bindManager(wl_client* client, void* data, uint32_t ver, uint32_t id) {
     const auto RESOURCE = m_exporters.emplace_back(makeUnique<CZxdgExporterV2>(client, ver, id)).get();
 
@@ -56,21 +69,16 @@ void CXDGForeignExporterProtocolV2::bindManager(wl_client* client, void* data, u
     }
 
     RESOURCE->setExportToplevel([this](CZxdgExporterV2* exporter, uint32_t id, wl_resource* surface) {
-        auto wlSurf = CWLSurfaceResource::fromResource(surface);
+        auto TOPLEVEL = xdgToplevelFromWlSurface(surface);
 
-        if (wlSurf->m_role != SURFACE_ROLE_XDG_SHELL) {
+        if (!TOPLEVEL) {
             exporter->error(zxdgExporterV2Error::ZXDG_EXPORTER_V2_ERROR_INVALID_SURFACE, "surface must be an xdg_toplevel");
             return;
         }
 
-        auto xdgSurfResource = sc<CXDGSurfaceRole*>(wlSurf->m_role.get())->m_xdgSurface.lock();
-        if (xdgSurfResource->m_toplevel.expired())
-            return;
-
-        auto              xdgSurf = xdgSurfResource->m_toplevel.lock();
-        const std::string HANDLE  = g_pTokenManager->getRandomUUID();
+        const std::string HANDLE = g_pTokenManager->getRandomUUID();
         const auto [ELM, EMPLACED] =
-            m_exported.emplace(HANDLE, makeShared<CXDGExportedResourceV2>(makeShared<CZxdgExportedV2>(exporter->client(), exporter->version(), id), xdgSurf, HANDLE));
+            m_exported.emplace(HANDLE, makeShared<CXDGExportedResourceV2>(makeShared<CZxdgExportedV2>(exporter->client(), exporter->version(), id), TOPLEVEL, HANDLE));
 
         // This should only happen if we have our generated handles collide.
         if UNLIKELY (!EMPLACED) {
@@ -102,34 +110,52 @@ void CXDGForeignExporterProtocolV2::destroyExported(CXDGExportedResourceV2* r) {
 
 CXDGImportedResourceV2::CXDGImportedResourceV2(SP<CZxdgImportedV2> imported, SP<CXDGExportedResourceV2> exported, const std::string& handle) :
     m_resource(imported), m_exported(exported), m_handle(handle) {
-    if UNLIKELY (!m_resource->resource() || m_exported.expired())
+    if UNLIKELY (!good())
         return;
 
     m_resource->setData(this);
 
     m_resource->setSetParentOf([this](CZxdgImportedV2* r, wl_resource* surf) {
-        const auto CHILDSURF = CWLSurfaceResource::fromResource(surf);
+        if (m_invalid)
+            return;
 
-        if (CHILDSURF->m_role != SURFACE_ROLE_XDG_SHELL) {
+        const auto CHILDTOPLEVEL = xdgToplevelFromWlSurface(surf);
+
+        if (!CHILDTOPLEVEL) {
             m_resource->error(zxdgImportedV2Error::ZXDG_IMPORTED_V2_ERROR_INVALID_SURFACE, "surface must be an xdg_toplevel");
             return;
         }
 
-        const auto CHILDXDGSURF = sc<CXDGSurfaceRole*>(CHILDSURF->m_role.get())->m_xdgSurface.lock();
-        if (CHILDXDGSURF->m_toplevel.expired())
-            return;
-
         if LIKELY (auto exportedTopLevel = m_exported->xdgSurf(); !exportedTopLevel.expired())
-            CHILDXDGSURF->m_toplevel->setNewParent(exportedTopLevel.lock());
+            CHILDTOPLEVEL->setNewParent(exportedTopLevel.lock());
     });
 
-    m_listeners.exportedDestroyed = m_exported->m_events.destroy.listen([this]() { PROTO::xdgForeignImporter->destroyImported(this); });
+    if (exported)
+        m_listeners.exportedDestroyed = exported->m_events.destroy.listen([this]() { invalidate(); });
+
     m_resource->setDestroy([this](CZxdgImportedV2*) { PROTO::xdgForeignImporter->destroyImported(this); });
     m_resource->setOnDestroy([this](CZxdgImportedV2*) { PROTO::xdgForeignImporter->destroyImported(this); });
+
+    if (m_exported.expired())
+        invalidate();
 }
 
-CXDGImportedResourceV2::~CXDGImportedResourceV2() {
-    m_resource->sendDestroyed();
+CXDGImportedResourceV2::~CXDGImportedResourceV2() {}
+
+bool CXDGImportedResourceV2::good() const {
+    return m_resource->resource();
+}
+
+void CXDGImportedResourceV2::invalidate() {
+    if (m_invalid)
+        return;
+
+    m_invalid = true;
+
+    if (!m_destroyedSent && m_resource && m_resource->resource()) {
+        m_destroyedSent = true;
+        m_resource->sendDestroyed();
+    }
 }
 
 CXDGForeignImporterProtocolV2::CXDGForeignImporterProtocolV2(const wl_interface* iface, const int& ver, const std::string& name) : IWaylandProtocol(iface, ver, name) {
@@ -151,15 +177,11 @@ void CXDGForeignImporterProtocolV2::bindManager(wl_client* client, void* data, u
         auto              imported =
             m_imports.emplace_back(makeUnique<CXDGImportedResourceV2>(makeShared<CZxdgImportedV2>(importer->client(), importer->version(), id), exported, HANDLE)).get();
 
-        if UNLIKELY (!imported->m_resource->resource()) {
+        if UNLIKELY (!imported->good()) {
             wl_client_post_no_memory(importer->client());
             m_imports.pop_back();
             return;
         }
-
-        // Couldn't find the handle.
-        if UNLIKELY (imported->m_exported.expired())
-            destroyImported(imported);
     });
 
     RESOURCE->setDestroy([this](CZxdgImporterV2* r) { onImporterDestroyed(r); });
diff --git a/src/protocols/XDGForeignV2.hpp b/src/protocols/XDGForeignV2.hpp
index 78c8dbf7..092945d4 100644
--- a/src/protocols/XDGForeignV2.hpp
+++ b/src/protocols/XDGForeignV2.hpp
@@ -42,10 +42,15 @@ class CXDGImportedResourceV2 {
     CXDGImportedResourceV2(SP<CZxdgImportedV2> resource, SP<CXDGExportedResourceV2> exported, const std::string& handle);
     ~CXDGImportedResourceV2();
 
+    bool good() const;
+    void invalidate();
+
   private:
     SP<CZxdgImportedV2>        m_resource;
     WP<CXDGExportedResourceV2> m_exported;
     std::string                m_handle;
+    bool                       m_invalid       = false;
+    bool                       m_destroyedSent = false;
 
     struct {
         CHyprSignalListener exportedDestroyed;