From a759d5b516505c7bc062ef74303cbdb2fc875df8 Mon Sep 17 00:00:00 2001
From: Lae Chen <laec@proton.me>
Date: Sun, 18 Dec 2022 10:00:17 +1300
Subject: [PATCH] Reconnect to X11 when existing connection stops working.
 (#212)

---
 src/client/x11_client.rs | 114 +++++++++++++++++++++++----------------
 1 file changed, 69 insertions(+), 45 deletions(-)

diff --git a/src/client/x11_client.rs b/src/client/x11_client.rs
index c68a8a7..285ab90 100644
--- a/src/client/x11_client.rs
+++ b/src/client/x11_client.rs
@@ -1,7 +1,11 @@
 use crate::client::Client;
+use anyhow::bail;
 use std::env;
+use x11rb::cookie::Cookie;
 use x11rb::protocol::xproto::{self};
 use x11rb::protocol::xproto::{AtomEnum, Window};
+use x11rb::rust_connection::ConnectionError;
+use x11rb::x11_utils::TryParse;
 use x11rb::{protocol::xproto::get_property, rust_connection::RustConnection};
 
 pub struct X11Client {
@@ -22,6 +26,10 @@ impl X11Client {
             println!("$DISPLAY is not set. Defaulting to DISPLAY=:0");
             env::set_var("DISPLAY", ":0");
         }
+        self.reconnect();
+    }
+
+    fn reconnect(&mut self) {
         match x11rb::connect(None) {
             Ok((connection, _)) => self.connection = Some(connection),
             Err(error) => {
@@ -43,57 +51,49 @@ impl Client for X11Client {
 
     fn current_application(&mut self) -> Option<String> {
         self.connect();
-        if let Some(conn) = &self.connection {
-            let mut window = get_focus_window(conn)?;
-            loop {
-                if let Some(wm_class) = get_wm_class(conn, window) {
-                    // Workaround: https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/classes/sun/awt/X11/XFocusProxyWindow.java#L35
-                    if &wm_class != "Focus-Proxy-Window.FocusProxy" {
-                        return Some(wm_class);
-                    }
+        let mut window = get_focus_window(self)?;
+        loop {
+            if let Some(wm_class) = get_wm_class(self, window) {
+                // Workaround: https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/classes/sun/awt/X11/XFocusProxyWindow.java#L35
+                if &wm_class != "Focus-Proxy-Window.FocusProxy" {
+                    return Some(wm_class);
                 }
-
-                window = get_parent_window(conn, window)?;
             }
+
+            window = get_parent_window(self, window)?;
         }
+    }
+}
+
+fn get_focus_window(client: &mut X11Client) -> Option<Window> {
+    get_cookie_reply_with_reconnect(client, xproto::get_input_focus)
+        .map(|reply| reply.focus)
+        .ok()
+}
+
+fn get_parent_window(client: &mut X11Client, window: Window) -> Option<Window> {
+    get_cookie_reply_with_reconnect(client, |conn| xproto::query_tree(conn, window))
+        .map(|reply| reply.parent)
+        .ok()
+}
+
+fn get_wm_class(client: &mut X11Client, window: Window) -> Option<String> {
+    let reply = get_cookie_reply_with_reconnect(client, |conn| {
+        get_property(conn, false, window, AtomEnum::WM_CLASS, AtomEnum::STRING, 0, 1024)
+    })
+    .ok()?;
+
+    if reply.value.is_empty() {
         return None;
     }
-}
 
-fn get_focus_window(conn: &RustConnection) -> Option<Window> {
-    if let Ok(cookie) = xproto::get_input_focus(conn) {
-        if let Ok(reply) = cookie.reply() {
-            return Some(reply.focus);
-        }
-    }
-    return None;
-}
-
-fn get_parent_window(conn: &RustConnection, window: Window) -> Option<Window> {
-    if let Ok(cookie) = xproto::query_tree(conn, window) {
-        if let Ok(reply) = cookie.reply() {
-            return Some(reply.parent);
-        }
-    }
-    return None;
-}
-
-fn get_wm_class(conn: &RustConnection, window: Window) -> Option<String> {
-    if let Ok(cookie) = get_property(conn, false, window, AtomEnum::WM_CLASS, AtomEnum::STRING, 0, 1024) {
-        if let Ok(reply) = cookie.reply() {
-            if reply.value.is_empty() {
-                return None;
-            }
-
-            if let Some(delimiter) = reply.value.iter().position(|byte| *byte == '\0' as u8) {
-                if let Ok(prefix) = String::from_utf8(reply.value[..delimiter].to_vec()) {
-                    let name = reply.value[(delimiter + 1)..].to_vec();
-                    if let Some(end) = name.iter().position(|byte| *byte == '\0' as u8) {
-                        if end == name.len() - 1 {
-                            if let Ok(name) = String::from_utf8(name[..end].to_vec()) {
-                                return Some(format!("{prefix}.{name}"));
-                            }
-                        }
+    if let Some(delimiter) = reply.value.iter().position(|byte| *byte == '\0' as u8) {
+        if let Ok(prefix) = String::from_utf8(reply.value[..delimiter].to_vec()) {
+            let name = reply.value[(delimiter + 1)..].to_vec();
+            if let Some(end) = name.iter().position(|byte| *byte == '\0' as u8) {
+                if end == name.len() - 1 {
+                    if let Ok(name) = String::from_utf8(name[..end].to_vec()) {
+                        return Some(format!("{prefix}.{name}"));
                     }
                 }
             }
@@ -101,3 +101,27 @@ fn get_wm_class(conn: &RustConnection, window: Window) -> Option<String> {
     }
     return None;
 }
+
+fn get_cookie_reply_with_reconnect<T: TryParse>(
+    client: &mut X11Client,
+    get_cookie: impl Fn(&RustConnection) -> Result<Cookie<RustConnection, T>, ConnectionError>,
+) -> anyhow::Result<T> {
+    return match get_cookie_reply(client, &get_cookie) {
+        Err(e) => {
+            println!("Reconnecting to X11 due to error: {}", e);
+            client.reconnect();
+            get_cookie_reply(client, &get_cookie)
+        }
+        x => x,
+    };
+}
+
+fn get_cookie_reply<T: TryParse>(
+    client: &X11Client,
+    get_cookie: &impl Fn(&RustConnection) -> Result<Cookie<RustConnection, T>, ConnectionError>,
+) -> anyhow::Result<T> {
+    match &client.connection {
+        Some(conn) => return Ok(get_cookie(conn)?.reply()?),
+        None => bail!("No connection to X11"),
+    }
+}