From 1cdce673df3d19ac1bc9ad0797ab613921aedc40 Mon Sep 17 00:00:00 2001
From: Ian Douglas Scott <idscott@system76.com>
Date: Mon, 28 Oct 2024 17:01:52 -0700
Subject: [PATCH] Remove screensaver inhibitors for disconnected clients

---
 src/freedesktop_screensaver.rs | 61 +++++++++++++++++++++++++---------
 src/main.rs                    |  1 -
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/src/freedesktop_screensaver.rs b/src/freedesktop_screensaver.rs
index 2e989db..fb7c2da 100644
--- a/src/freedesktop_screensaver.rs
+++ b/src/freedesktop_screensaver.rs
@@ -1,18 +1,23 @@
 // https://specifications.freedesktop.org/idle-inhibit-spec/latest
+// https://invent.kde.org/plasma/kscreenlocker/-/blob/master/dbus/org.freedesktop.ScreenSaver.xml
 
-use std::collections::HashMap;
+use futures_lite::StreamExt;
+use std::sync::{
+    atomic::{AtomicU32, Ordering},
+    Arc, Mutex,
+};
 
 #[derive(Debug)]
 struct Inhibitor {
+    cookie: u32,
     application_name: String,
     reason_for_inhibit: String,
     client: zbus::names::UniqueName<'static>,
 }
 
-#[derive(Default)]
 pub struct Screensaver {
-    inhibitors: HashMap<u32, Inhibitor>,
-    last_cookie: u32,
+    inhibitors: Arc<Mutex<Vec<Inhibitor>>>,
+    last_cookie: AtomicU32,
 }
 
 #[zbus::interface(name = "org.freedesktop.ScreenSaver")]
@@ -23,33 +28,57 @@ impl Screensaver {
         reason_for_inhibit: String,
         #[zbus(header)] header: zbus::message::Header<'_>,
     ) -> u32 {
-        self.last_cookie += 1;
+        let cookie = self.last_cookie.fetch_add(1, Ordering::Relaxed) + 1;
         if let Some(sender) = header.sender() {
-            self.inhibitors.insert(
-                self.last_cookie,
-                Inhibitor {
-                    application_name,
-                    reason_for_inhibit,
-                    client: sender.to_owned(),
-                },
-            );
+            self.inhibitors.lock().unwrap().push(Inhibitor {
+                cookie,
+                application_name,
+                reason_for_inhibit,
+                client: sender.to_owned(),
+            });
         }
-        self.last_cookie
+        cookie
     }
 
     fn un_inhibit(&mut self, cookie: u32) {
-        self.inhibitors.remove(&cookie);
+        let mut inhibitors = self.inhibitors.lock().unwrap();
+        if let Some(idx) = inhibitors.iter().position(|x| x.cookie == cookie) {
+            inhibitors.remove(idx);
+        }
     }
 }
 
 pub async fn serve(conn: &zbus::Connection) -> zbus::Result<()> {
+    let inhibitors = Arc::new(Mutex::new(Vec::new()));
+
     conn.request_name_with_flags(
         "org.freedesktop.ScreenSaver",
         zbus::fdo::RequestNameFlags::ReplaceExisting.into(),
     )
     .await?;
     conn.object_server()
-        .at("/org/freedesktop/ScreenSaver", Screensaver::default())
+        .at(
+            "/org/freedesktop/ScreenSaver",
+            Screensaver {
+                inhibitors: inhibitors.clone(),
+                last_cookie: AtomicU32::new(0),
+            },
+        )
         .await?;
+
+    let dbus = zbus::fdo::DBusProxy::new(conn).await?;
+    let mut name_owner_stream = dbus.receive_name_owner_changed().await?;
+    while let Some(event) = name_owner_stream.next().await {
+        let args = event.args()?;
+        if args.new_owner.is_none() {
+            if let zbus::names::BusName::Unique(name) = args.name {
+                inhibitors
+                    .lock()
+                    .unwrap()
+                    .retain(|inhibitor| inhibitor.client != name);
+            }
+        }
+    }
+
     Ok(())
 }
diff --git a/src/main.rs b/src/main.rs
index 34599eb..9b37655 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -338,7 +338,6 @@ fn main() {
                 if let Err(err) = freedesktop_screensaver::serve(&connection).await {
                     log::error!("failed to serve FreeDesktop screensaver interface: {}", err);
                 }
-                std::future::pending::<()>().await;
             }
         })
         .unwrap();