screencopy: improve restore data handling

src/portals/Screencopy.cpp

@@ -84,6 +84,7 @@ void CScreencopyPortal::onSelectSources(sdbus::MethodCall& call) {
         uint64_t    windowHandle;
         bool        withCursor;
         uint64_t    timeIssued;
+        std::string windowClass;
     } restoreData;
 
     for (auto& [key, val] : options) {
@@ -133,55 +134,29 @@ void CScreencopyPortal::onSelectSources(sdbus::MethodCall& call) {
                 auto     sv = data.get<std::unordered_map<std::string, sdbus::Variant>>();
 
                 uint64_t windowHandle = 0;
-                std::string windowClass;
 
                 restoreData.windowHandle = 0;
                 restoreData.exists       = true;
 
                 for (auto& [tkkey, tkval] : sv) {
-                    if (tkkey == "output") {
+                    if (tkkey == "output")
                         restoreData.output = tkval.get<std::string>();
-                    } else if (tkkey == "windowHandle") {
+                    else if (tkkey == "windowHandle")
                         windowHandle = tkval.get<uint64_t>();
-                    } else if (tkkey == "windowClass") {
+                    else if (tkkey == "windowClass")
-                        windowClass = tkval.get<std::string>();
+                        restoreData.windowClass = tkval.get<std::string>();
-                    } else if (tkkey == "withCursor") {
+                    else if (tkkey == "withCursor")
                         restoreData.withCursor = (bool)tkval.get<uint32_t>();
-                    } else if (tkkey == "timeIssued") {
+                    else if (tkkey == "timeIssued")
                         restoreData.timeIssued = tkval.get<uint64_t>();
-                    } else if (tkkey == "token") {
+                    else if (tkkey == "token")
                         restoreData.token = tkval.get<std::string>();
-                    } else {
+                    else
                         Debug::log(LOG, "[screencopy] restore token v3, unknown prop {}", tkkey);
                 }
-                }
 
-                Debug::log(LOG, "[screencopy] Restore token v3 {} with data: {} {} {} {} {}", restoreData.token, windowHandle, windowClass, restoreData.output,
+                Debug::log(LOG, "[screencopy] Restore token v3 {} with data: {} {} {} {} {}", restoreData.token, windowHandle, restoreData.windowClass, restoreData.output,
                            restoreData.withCursor, restoreData.timeIssued);
-
-                // find window
-                if (windowHandle != 0 || !windowClass.empty()) {
-                    if (windowHandle != 0) {
-                        for (auto& h : g_pPortalManager->m_sHelpers.toplevel->m_vToplevels) {
-                            if ((uint64_t)h->handle == windowHandle) {
-                                restoreData.windowHandle = (uint64_t)h->handle;
-                                Debug::log(LOG, "[screencopy] token v3 window found by handle {}", (void*)windowHandle);
-                                break;
-                            }
-                        }
-                    }
-
-                    if (restoreData.windowHandle == 0 && !windowClass.empty()) {
-                        // try class
-                        for (auto& h : g_pPortalManager->m_sHelpers.toplevel->m_vToplevels) {
-                            if (h->windowClass == windowClass) {
-                                restoreData.windowHandle = (uint64_t)h->handle;
-                                Debug::log(LOG, "[screencopy] token v3 window found by class {}", windowClass);
-                                break;
-                            }
-                        }
-                    }
-                }
             }
 
         } else if (key == "persist_mode") {
@@ -192,14 +167,21 @@ void CScreencopyPortal::onSelectSources(sdbus::MethodCall& call) {
         }
     }
 
-    const bool     RESTOREDATAVALID = restoreData.exists && g_pPortalManager->getOutputFromName(restoreData.output);
+    // clang-format off
+    const bool     RESTOREDATAVALID = restoreData.exists &&
+    (
+        (!restoreData.output.empty() && g_pPortalManager->getOutputFromName(restoreData.output)) || // output exists
+        (!restoreData.windowClass.empty() && g_pPortalManager->m_sHelpers.toplevel->handleFromClass(restoreData.windowClass)) // window exists
+    );
+    // clang-format on
 
     SSelectionData SHAREDATA;
     if (RESTOREDATAVALID) {
         Debug::log(LOG, "[screencopy] restore data valid, not prompting");
 
         SHAREDATA.output       = restoreData.output;
-        SHAREDATA.type       = restoreData.windowHandle ? TYPE_WINDOW : TYPE_OUTPUT;
+        SHAREDATA.type         = !restoreData.windowClass.empty() ? TYPE_WINDOW : TYPE_OUTPUT;
+        SHAREDATA.windowHandle = !restoreData.windowClass.empty() ? g_pPortalManager->m_sHelpers.toplevel->handleFromClass(restoreData.windowClass)->handle : nullptr;
         SHAREDATA.allowToken   = true; // user allowed token before
         PSESSION->cursorMode   = restoreData.withCursor;
     } else {
@@ -275,13 +257,8 @@ void CScreencopyPortal::onStart(sdbus::MethodCall& call) {
             case TYPE_GEOMETRY:
             case TYPE_OUTPUT: mapData["output"] = PSESSION->selection.output; break;
             case TYPE_WINDOW:
-                mapData["windowHandle"] = (uint64_t)PSESSION->selection.windowHandle;
+                mapData["windowHandle"] = (uint64_t)PSESSION->selection.windowHandle->resource();
-                for (auto& w : g_pPortalManager->m_sHelpers.toplevel->m_vToplevels) {
+                mapData["windowClass"]  = PSESSION->selection.windowClass;
-                    if (w->handle == PSESSION->selection.windowHandle) {
-                        mapData["windowClass"] = w->windowClass;
-                        break;
-                    }
-                }
                 break;
             default: Debug::log(ERR, "[screencopy] wonk selection in token saving"); break;
         }

src/shared/ScreencopyShared.cpp

@@ -28,7 +28,6 @@ std::string buildWindowList() {
         return result;
 
     for (auto& e : g_pPortalManager->m_sHelpers.toplevel->m_vToplevels) {
-
         result += std::format("{}[HC>]{}[HT>]{}[HE>]", (uint32_t)(((uint64_t)e->handle->resource()) & 0xFFFFFFFF), sanitizeNameForWindowList(e->windowClass),
                               sanitizeNameForWindowList(e->windowTitle));
     }
@@ -87,15 +86,11 @@ SSelectionData promptForScreencopySelection() {
         uint32_t handleLo = std::stoull(SEL.substr(7));
         data.windowHandle = nullptr;
 
-        for (auto& e : g_pPortalManager->m_sHelpers.toplevel->m_vToplevels) {
+        const auto HANDLE = g_pPortalManager->m_sHelpers.toplevel->handleFromHandleLower(handleLo);
-            uint32_t handleLoE = (uint32_t)(((uint64_t)e->handle->resource()) & 0xFFFFFFFF);
+        if (HANDLE) {
-
+            data.windowHandle = HANDLE->handle;
-            if (handleLoE == handleLo) {
+            data.windowClass  = HANDLE->windowClass;
-                data.windowHandle = e->handle;
-                break;
         }
-        }
-
     } else if (SEL.find("region:") == 0) {
         std::string running = SEL;
         running             = running.substr(7);

src/shared/ScreencopyShared.hpp

@@ -36,6 +36,9 @@ struct SSelectionData {
     SP<CCZwlrForeignToplevelHandleV1> windowHandle = nullptr;
     uint32_t                          x = 0, y = 0, w = 0, h = 0; // for TYPE_GEOMETRY
     bool                              allowToken = false;
+
+    // for restoring
+    std::string windowClass;
 };
 
 struct wl_buffer;

src/shared/ToplevelManager.cpp

@@ -41,7 +41,7 @@ void CToplevelManager::activate() {
     m_pManager->setToplevel([this](CCZwlrForeignToplevelManagerV1* r, wl_proxy* newHandle) {
         Debug::log(TRACE, "[toplevel] New toplevel at {}", (void*)newHandle);
 
-        m_vToplevels.emplace_back(std::make_unique<SToplevelHandle>(makeShared<CCZwlrForeignToplevelHandleV1>(newHandle)));
+        m_vToplevels.emplace_back(makeShared<SToplevelHandle>(makeShared<CCZwlrForeignToplevelHandleV1>(newHandle)));
     });
     m_pManager->setFinished([this](CCZwlrForeignToplevelManagerV1* r) { m_vToplevels.clear(); });
 
@@ -63,3 +63,21 @@ void CToplevelManager::deactivate() {
 
     Debug::log(LOG, "[toplevel] unbound manager");
 }
+
+SP<SToplevelHandle> CToplevelManager::handleFromClass(const std::string& windowClass) {
+    for (auto& tl : m_vToplevels) {
+        if (tl->windowClass == windowClass)
+            return tl;
+    }
+
+    return nullptr;
+}
+
+SP<SToplevelHandle> CToplevelManager::handleFromHandleLower(uint32_t handle) {
+    for (auto& tl : m_vToplevels) {
+        if (((uint64_t)tl->handle->resource() & 0xFFFFFFFF) == handle)
+            return tl;
+    }
+
+    return nullptr;
+}

src/shared/ToplevelManager.hpp

@@ -23,8 +23,10 @@ class CToplevelManager {
 
     void                             activate();
     void                             deactivate();
+    SP<SToplevelHandle>              handleFromClass(const std::string& windowClass);
+    SP<SToplevelHandle>              handleFromHandleLower(uint32_t handle);
 
-    std::vector<std::unique_ptr<SToplevelHandle>> m_vToplevels;
+    std::vector<SP<SToplevelHandle>> m_vToplevels;
 
   private:
     SP<CCZwlrForeignToplevelManagerV1> m_pManager = nullptr;
@@ -35,4 +37,6 @@ class CToplevelManager {
         uint32_t name    = 0;
         uint32_t version = 0;
     } m_sWaylandConnection;
+
+    friend struct SToplevelHandle;
 };