core: verify surface roles on creation of objects

5 files changed, 55 insertions, 9 deletions

diff --git a/src/protocols/LayerShell.cpp b/src/protocols/LayerShell.cpp
index 8fa6dd27..d4b105cb 100644
--- a/src/protocols/LayerShell.cpp
+++ b/src/protocols/LayerShell.cpp
@@ -159,6 +159,12 @@ CLayerShellResource::CLayerShellResource(SP<CZwlrLayerSurfaceV1> resource_, SP<C
 
 CLayerShellResource::~CLayerShellResource() {
     events.destroy.emit();
+    if (surface)
+        surface->resetRole();
+}
+
+eSurfaceRole CLayerShellResource::role() {
+    return SURFACE_ROLE_LAYER_SHELL;
 }
 
 bool CLayerShellResource::good() {
@@ -207,8 +213,19 @@ void CLayerShellProtocol::destroyResource(CLayerShellResource* surf) {
 void CLayerShellProtocol::onGetLayerSurface(CZwlrLayerShellV1* pMgr, uint32_t id, wl_resource* surface, wl_resource* output, zwlrLayerShellV1Layer layer, std::string namespace_) {
     const auto CLIENT   = pMgr->client();
     const auto PMONITOR = output ? CWLOutputResource::fromResource(output)->monitor.get() : nullptr;
-    const auto RESOURCE = m_vLayers.emplace_back(
-        makeShared<CLayerShellResource>(makeShared<CZwlrLayerSurfaceV1>(CLIENT, pMgr->version(), id), CWLSurfaceResource::fromResource(surface), namespace_, PMONITOR, layer));
+    auto       SURF     = CWLSurfaceResource::fromResource(surface);
+
+    if (!SURF) {
+        pMgr->error(-1, "Invalid surface");
+        return;
+    }
+
+    if (SURF->role->role() != SURFACE_ROLE_UNASSIGNED) {
+        pMgr->error(-1, "Surface already has a different role");
+        return;
+    }
+
+    const auto RESOURCE = m_vLayers.emplace_back(makeShared<CLayerShellResource>(makeShared<CZwlrLayerSurfaceV1>(CLIENT, pMgr->version(), id), SURF, namespace_, PMONITOR, layer));
 
     if (!RESOURCE->good()) {
         pMgr->noMemory();
@@ -216,6 +233,7 @@ void CLayerShellProtocol::onGetLayerSurface(CZwlrLayerShellV1* pMgr, uint32_t id
         return;
     }
 
+    SURF->role = RESOURCE;
     g_pCompositor->m_vLayers.emplace_back(CLayerSurface::create(RESOURCE));
 
     LOGM(LOG, "New wlr_layer_surface {:x}", (uintptr_t)RESOURCE.get());
diff --git a/src/protocols/LayerShell.hpp b/src/protocols/LayerShell.hpp
index f348e86b..7fa14447 100644
--- a/src/protocols/LayerShell.hpp
+++ b/src/protocols/LayerShell.hpp
@@ -8,18 +8,20 @@
 #include "wlr-layer-shell-unstable-v1.hpp"
 #include "../helpers/Vector2D.hpp"
 #include "../helpers/signal/Signal.hpp"
+#include "types/SurfaceRole.hpp"
 
 class CMonitor;
 class CWLSurfaceResource;
 
-class CLayerShellResource {
+class CLayerShellResource : public ISurfaceRole {
   public:
     CLayerShellResource(SP<CZwlrLayerSurfaceV1> resource_, SP<CWLSurfaceResource> surf_, std::string namespace_, CMonitor* pMonitor, zwlrLayerShellV1Layer layer);
     ~CLayerShellResource();
 
-    bool good();
-    void configure(const Vector2D& size);
-    void sendClosed();
+    bool                 good();
+    void                 configure(const Vector2D& size);
+    void                 sendClosed();
+    virtual eSurfaceRole role();
 
     enum eCommittedState {
         STATE_SIZE          = (1 << 0),
diff --git a/src/protocols/XDGShell.cpp b/src/protocols/XDGShell.cpp
index 191a1124..dcbc6162 100644
--- a/src/protocols/XDGShell.cpp
+++ b/src/protocols/XDGShell.cpp
@@ -419,6 +419,12 @@ CXDGSurfaceResource::~CXDGSurfaceResource() {
     events.destroy.emit();
     if (configureSource)
         wl_event_source_remove(configureSource);
+    if (surface)
+        surface->resetRole();
+}
+
+eSurfaceRole CXDGSurfaceResource::role() {
+    return SURFACE_ROLE_XDG_SHELL;
 }
 
 bool CXDGSurfaceResource::good() {
@@ -646,8 +652,19 @@ CXDGWMBase::CXDGWMBase(SP<CXdgWmBase> resource_) : resource(resource_) {
     });
 
     resource->setGetXdgSurface([this](CXdgWmBase* r, uint32_t id, wl_resource* surf) {
-        const auto RESOURCE = PROTO::xdgShell->m_vSurfaces.emplace_back(
-            makeShared<CXDGSurfaceResource>(makeShared<CXdgSurface>(r->client(), r->version(), id), self.lock(), CWLSurfaceResource::fromResource(surf)));
+        auto SURF = CWLSurfaceResource::fromResource(surf);
+
+        if (!SURF) {
+            r->error(-1, "Invalid surface passed");
+            return;
+        }
+
+        if (SURF->role->role() != SURFACE_ROLE_UNASSIGNED) {
+            r->error(-1, "Surface already has a different role");
+            return;
+        }
+
+        const auto RESOURCE = PROTO::xdgShell->m_vSurfaces.emplace_back(makeShared<CXDGSurfaceResource>(makeShared<CXdgSurface>(r->client(), r->version(), id), self.lock(), SURF));
 
         if (!RESOURCE->good()) {
             r->noMemory();
@@ -656,6 +673,7 @@ CXDGWMBase::CXDGWMBase(SP<CXdgWmBase> resource_) : resource(resource_) {
         }
 
         RESOURCE->self = RESOURCE;
+        SURF->role     = RESOURCE;
 
         surfaces.emplace_back(RESOURCE);
 
diff --git a/src/protocols/XDGShell.hpp b/src/protocols/XDGShell.hpp
index fe2517a2..d60db86a 100644
--- a/src/protocols/XDGShell.hpp
+++ b/src/protocols/XDGShell.hpp
@@ -9,6 +9,7 @@
 #include "../helpers/Vector2D.hpp"
 #include "../helpers/Box.hpp"
 #include "../helpers/signal/Signal.hpp"
+#include "types/SurfaceRole.hpp"
 
 class CXDGWMBase;
 class CXDGPositionerResource;
@@ -137,13 +138,15 @@ class CXDGToplevelResource {
     void             applyState();
 };
 
-class CXDGSurfaceResource {
+class CXDGSurfaceResource : public ISurfaceRole {
   public:
     CXDGSurfaceResource(SP<CXdgSurface> resource_, SP<CXDGWMBase> owner_, SP<CWLSurfaceResource> surface_);
     ~CXDGSurfaceResource();
 
     static SP<CXDGSurfaceResource> fromResource(wl_resource*);
 
+    virtual eSurfaceRole           role();
+
     bool                           good();
 
     WP<CXDGWMBase>                 owner;
diff --git a/src/protocols/core/Subcompositor.cpp b/src/protocols/core/Subcompositor.cpp
index 213fdc42..d407636c 100644
--- a/src/protocols/core/Subcompositor.cpp
+++ b/src/protocols/core/Subcompositor.cpp
@@ -133,6 +133,11 @@ CWLSubcompositorResource::CWLSubcompositorResource(SP<CWlSubcompositor> resource
             return;
         }
 
+        if (SURF->role->role() != SURFACE_ROLE_UNASSIGNED) {
+            r->error(-1, "Surface already has a different role");
+            return;
+        }
+
         SP<CWLSurfaceResource> t1Parent = nullptr;
 
         if (PARENT->role->role() == SURFACE_ROLE_SUBSURFACE) {