From 4c8ffbab6a42f05e470fba4cc22924d55e340561 Mon Sep 17 00:00:00 2001
From: "Brian J. Tarricone" <brian@tarricone.org>
Date: Mon, 15 Aug 2022 22:54:54 -0700
Subject: [PATCH] Disable/reenable DPMS when inhibiting

---
 dbus-service/Cargo.toml  |   2 +-
 dbus-service/src/main.rs | 120 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 108 insertions(+), 14 deletions(-)

diff --git a/dbus-service/Cargo.toml b/dbus-service/Cargo.toml
index 8a0d183..3ac3679 100644
--- a/dbus-service/Cargo.toml
+++ b/dbus-service/Cargo.toml
@@ -11,5 +11,5 @@ bscreensaver-command = { path = "../command" }
 bscreensaver-util = { path = "../util" }
 futures = "0.3"
 log = "0.4"
-xcb = "1"
+xcb = { version = "1", features = ["dpms"] }
 zbus = "2"
diff --git a/dbus-service/src/main.rs b/dbus-service/src/main.rs
index 7ab2208..16f8647 100644
--- a/dbus-service/src/main.rs
+++ b/dbus-service/src/main.rs
@@ -79,6 +79,93 @@ impl ScreenSaver {
     }
 }
 
+#[derive(Clone, Copy, PartialEq)]
+enum DpmsState {
+    Unknown,
+    WasDisabled,
+    WeDisabled,
+    WeEnabled,
+}
+
+struct DpmsHandling {
+    conn: Option<xcb::Connection>,
+    state: DpmsState,
+}
+
+impl DpmsHandling {
+    fn init() -> Self {
+        xcb::Connection::connect_with_extensions(None, &[xcb::Extension::Dpms],&[])
+            .map(|(conn, _)| Some(conn))
+            .unwrap_or_else(|err| match err {
+                xcb::ConnError::ClosedExtNotSupported => None,
+                _ => {
+                    warn!("Failed to connect to X display; we will not handle DPMS on inhibit: {}", err);
+                    None
+                },
+            })
+        .and_then(|conn| {
+            let cookie = conn.send_request(&xcb::dpms::Capable {});
+            conn.wait_for_reply(cookie)
+                .map(|capable| {
+                    if capable.capable() {
+                        Some(Self {
+                            conn: Some(conn),
+                            state: DpmsState::Unknown,
+                        })
+                    } else {
+                        None
+                    }
+                })
+            .unwrap_or_else(|err| {
+                warn!("Failed to check if X server is DPMS-capable: {}", err);
+                None
+            })
+        })
+        .unwrap_or_else(|| Self {
+            conn: None,
+            state: DpmsState::Unknown,
+        })
+    }
+
+    fn maybe_disable_dpms(&mut self) {
+        if let Some(conn) = &self.conn {
+            let cookie = conn.send_request(&xcb::dpms::Info {});
+            match conn.wait_for_reply(cookie) {
+                Err(err) => warn!("Failed to query DPMS state: {}", err),
+                Ok(info) if info.state() => match conn.send_and_check_request(&xcb::dpms::Disable {}) {
+                    Err(err) => warn!("Failed to disable DPMS: {}", err),
+                    Ok(_) => {
+                        debug!("Successfully disabled DPMS");
+                        self.state = DpmsState::WeDisabled;
+                    },
+                }
+                Ok(_) => {
+                    debug!("DPMS was already disabled");
+                    if self.state != DpmsState::WeDisabled {
+                        self.state = DpmsState::WasDisabled;
+                    }
+                },
+            }
+        }
+    }
+
+    fn maybe_enable_dpms(&mut self) {
+        if let Some(conn) = &self.conn {
+            if self.state == DpmsState::WeDisabled {
+                match conn.send_and_check_request(&xcb::dpms::Enable {}) {
+                    Err(err) => warn!("Failed to enable DPMS: {}", err),
+                    Ok(_) => {
+                        debug!("Successfully enabled DPMS");
+                        self.state = DpmsState::WeEnabled;
+                    },
+                }
+            } else {
+                debug!("We didn't disable DPMS, so we're not going to re-enable it");
+            }
+        }
+    }
+}
+
 #[async_std::main]
 async fn main() {
     init_logging("BSCREENSAVER_DBUS_SERVICE_LOG");
@@ -183,6 +270,7 @@ async fn dbus_task(state: Arc<Mutex<State>>) -> anyhow::Result<()> {
 
 async fn heartbeat_task(state_mtx: Arc<Mutex<State>>) -> anyhow::Result<()> {
     let mut last_heartbeat: Option<Instant> = None;
+    let mut dpms_handling = task::block_on(async { DpmsHandling::init() });
 
     loop {
         let state = state_mtx.lock().await;
@@ -207,20 +295,26 @@ async fn heartbeat_task(state_mtx: Arc<Mutex<State>>) -> anyhow::Result<()> {
         debug!("Heartbeat timeout expired");
 
         let state = state_mtx.lock().await;
-        if !state.inhibitors.is_empty() && (last_heartbeat.is_none() || last_heartbeat.as_ref().filter(|lh| lh.elapsed() < HEARTBEAT_INTERVAL).is_none()) {
-            trace!("About to deactivate; active inhibitors:");
-            for inhibitor in &state.inhibitors {
-                trace!("    {}: {}", inhibitor.cookie, inhibitor.app_name);
-            }
-            drop(state);
-            task::block_on(async {
-                if let Err(err) = bscreensaver_command(BCommand::Deactivate, Some(Duration::from_secs(4))) {
-                    warn!("Failed to deactivate screen lock: {}", err);
-                } else {
-                    debug!("Successfully issued deactivate heartbeat");
-                    last_heartbeat = Some(Instant::now());
+        if !state.inhibitors.is_empty() {
+            task::block_on(async { dpms_handling.maybe_disable_dpms() });
+
+            if last_heartbeat.is_none() || last_heartbeat.as_ref().filter(|lh| lh.elapsed() < HEARTBEAT_INTERVAL).is_none() {
+                trace!("About to deactivate; active inhibitors:");
+                for inhibitor in &state.inhibitors {
+                    trace!("    {}: {}", inhibitor.cookie, inhibitor.app_name);
                 }
-            });
+                drop(state);
+                task::block_on(async {
+                    if let Err(err) = bscreensaver_command(BCommand::Deactivate, Some(Duration::from_secs(4))) {
+                        warn!("Failed to deactivate screen lock: {}", err);
+                    } else {
+                        debug!("Successfully issued deactivate heartbeat");
+                        last_heartbeat = Some(Instant::now());
+                    }
+                });
+            }
+        } else {
+            task::block_on(async { dpms_handling.maybe_enable_dpms() });
         }
     }
 }