From ce62521883e217aab4da305d71e40e4a50567dfb Mon Sep 17 00:00:00 2001
From: Ikalco <73481042+ikalco@users.noreply.github.com>
Date: Mon, 14 Oct 2024 05:37:42 -0500
Subject: [PATCH] drm-lease: fix crashes and implementation (#8116)

---
 src/protocols/DRMLease.cpp | 28 ++++++++++++++++++++--------
 src/protocols/DRMLease.hpp |  3 ++-
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/protocols/DRMLease.cpp b/src/protocols/DRMLease.cpp
index 37f2e2eb..d0114ce7 100644
--- a/src/protocols/DRMLease.cpp
+++ b/src/protocols/DRMLease.cpp
@@ -44,10 +44,6 @@ CDRMLeaseResource::CDRMLeaseResource(SP<CWpDrmLeaseV1> resource_, SP<CDRMLeaseRe
         return;
     }
 
-    LOGM(LOG, "Granting lease, sending fd {}", aqlease->leaseFD);
-
-    resource->sendLeaseFd(aqlease->leaseFD);
-
     lease = aqlease;
 
     for (auto const& m : requested) {
@@ -61,8 +57,13 @@ CDRMLeaseResource::CDRMLeaseResource(SP<CWpDrmLeaseV1> resource_, SP<CDRMLeaseRe
         }
 
         resource->sendFinished();
+        LOGM(LOG, "Revoking lease for fd {}", lease->leaseFD);
     });
 
+    LOGM(LOG, "Granting lease, sending fd {}", lease->leaseFD);
+
+    resource->sendLeaseFd(lease->leaseFD);
+
     close(lease->leaseFD);
 }
 
@@ -70,6 +71,12 @@ bool CDRMLeaseResource::good() {
     return resource->resource();
 }
 
+CDRMLeaseResource::~CDRMLeaseResource() {
+    // destroy in this order to ensure listener gets called
+    lease.reset();
+    listeners.destroyLease.reset();
+}
+
 CDRMLeaseRequestResource::CDRMLeaseRequestResource(SP<CWpDrmLeaseRequestV1> resource_) : resource(resource_) {
     if (!good())
         return;
@@ -100,7 +107,7 @@ CDRMLeaseRequestResource::CDRMLeaseRequestResource(SP<CWpDrmLeaseRequestV1> reso
             return;
         }
 
-        auto RESOURCE = makeShared<CDRMLeaseResource>(makeShared<CWpDrmLeaseV1>(resource->client(), resource->version(), -1), self.lock());
+        auto RESOURCE = makeShared<CDRMLeaseResource>(makeShared<CWpDrmLeaseV1>(resource->client(), resource->version(), id), self.lock());
         if (!RESOURCE) {
             resource->noMemory();
             return;
@@ -185,7 +192,8 @@ CDRMLeaseDeviceResource::CDRMLeaseDeviceResource(SP<CWpDrmLeaseDeviceV1> resourc
     close(fd);
 
     for (auto const& m : PROTO::lease->primaryDevice->offeredOutputs) {
-        sendConnector(m.lock());
+        if (m)
+            sendConnector(m.lock());
     }
 
     resource->sendDone();
@@ -196,10 +204,10 @@ bool CDRMLeaseDeviceResource::good() {
 }
 
 void CDRMLeaseDeviceResource::sendConnector(SP<CMonitor> monitor) {
-    if (std::find_if(connectorsSent.begin(), connectorsSent.end(), [monitor](const auto& e) { return e->monitor == monitor; }) != connectorsSent.end())
+    if (std::find_if(connectorsSent.begin(), connectorsSent.end(), [monitor](const auto& e) { return e && !e->dead && e->monitor == monitor; }) != connectorsSent.end())
         return;
 
-    auto RESOURCE = makeShared<CDRMLeaseConnectorResource>(makeShared<CWpDrmLeaseConnectorV1>(resource->client(), resource->version(), -1), monitor);
+    auto RESOURCE = makeShared<CDRMLeaseConnectorResource>(makeShared<CWpDrmLeaseConnectorV1>(resource->client(), resource->version(), 0), monitor);
     if (!RESOURCE) {
         resource->noMemory();
         return;
@@ -267,6 +275,9 @@ void CDRMLeaseProtocol::destroyResource(CDRMLeaseDeviceResource* resource) {
 }
 
 void CDRMLeaseProtocol::destroyResource(CDRMLeaseConnectorResource* resource) {
+    for (const auto& m : m_vManagers) {
+        std::erase_if(m->connectorsSent, [resource](const auto& e) { return e.expired() || e->dead || e.get() == resource; });
+    }
     std::erase_if(m_vConnectors, [resource](const auto& e) { return e.get() == resource; });
 }
 
@@ -279,6 +290,7 @@ void CDRMLeaseProtocol::destroyResource(CDRMLeaseResource* resource) {
 }
 
 void CDRMLeaseProtocol::offer(SP<CMonitor> monitor) {
+    std::erase_if(primaryDevice->offeredOutputs, [](const auto& e) { return e.expired(); });
     if (std::find(primaryDevice->offeredOutputs.begin(), primaryDevice->offeredOutputs.end(), monitor) != primaryDevice->offeredOutputs.end())
         return;
 
diff --git a/src/protocols/DRMLease.hpp b/src/protocols/DRMLease.hpp
index 56eaa3db..3671cfce 100644
--- a/src/protocols/DRMLease.hpp
+++ b/src/protocols/DRMLease.hpp
@@ -22,12 +22,13 @@ class CDRMLeaseRequestResource;
 class CDRMLeaseResource {
   public:
     CDRMLeaseResource(SP<CWpDrmLeaseV1> resource_, SP<CDRMLeaseRequestResource> request);
+    ~CDRMLeaseResource();
 
     bool                                        good();
 
     WP<CDRMLeaseDeviceResource>                 parent;
     std::vector<WP<CDRMLeaseConnectorResource>> requested;
-    WP<Aquamarine::CDRMLease>                   lease;
+    SP<Aquamarine::CDRMLease>                   lease;
 
     int                                         leaseFD = -1;