[refactor] environment for swhkd (#246)

newtoallofthis123 · web-flow · commit 5e8ebe3b5fea · 2024-03-26T21:38:45.000+05:30
* refactor: New Env for swhkd

* Refactor Functions

* Fix Spelling in environ

* Fix spelling in daemon

* Add Path Validations

* Run formatting
diff --git a/swhkd/src/daemon.rs b/swhkd/src/daemon.rs
@@ -26,6 +26,7 @@ use tokio_stream::{StreamExt, StreamMap};
 use tokio_udev::{AsyncMonitorSocket, EventType, MonitorBuilder};
 
 mod config;
+mod environ;
 mod perms;
 mod uinput;
 
@@ -55,20 +56,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
     env_logger::init();
     log::trace!("Logger initialized.");
 
-    let invoking_uid = match env::var("PKEXEC_UID") {
-        Ok(uid) => {
-            let uid = uid.parse::<u32>().unwrap();
-            log::trace!("Invoking UID: {}", uid);
-            uid
-        }
-        Err(_) => {
-            log::error!("Failed to launch swhkd!!!");
-            log::error!("Make sure to launch the binary with pkexec.");
-            exit(1);
-        }
-    };
+    let env = environ::Env::construct();
+    log::trace!("Environment Aquired");
+
+    let invoking_uid = env.pkexec_id;
 
-    setup_swhkd(invoking_uid);
+    setup_swhkd(invoking_uid, env.xdg_runtime_dir.clone().to_string_lossy().to_string());
 
     let load_config = || {
         // Drop privileges to the invoking user.
@@ -77,7 +70,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
         let config_file_path: PathBuf = if args.is_present("config") {
             Path::new(args.value_of("config").unwrap()).to_path_buf()
         } else {
-            fetch_xdg_config_path()
+            env.fetch_xdg_config_path()
         };
 
         log::debug!("Using config file path: {:#?}", config_file_path);
@@ -231,7 +224,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     tokio::pin!(hotkey_repeat_timer);
 
     // The socket we're sending the commands to.
-    let socket_file_path = fetch_xdg_runtime_socket_path();
+    let socket_file_path = env.fetch_xdg_runtime_socket_path();
     loop {
         select! {
             _ = &mut hotkey_repeat_timer, if &last_hotkey.is_some() => {
@@ -501,50 +494,12 @@ pub fn set_command_line_args() -> Command<'static> {
     app
 }
 
-pub fn fetch_xdg_config_path() -> PathBuf {
-    let config_file_path: PathBuf = match env::var("XDG_CONFIG_HOME") {
-        Ok(val) => {
-            log::debug!("XDG_CONFIG_HOME exists: {:#?}", val);
-            Path::new(&val).join("swhkd/swhkdrc")
-        }
-        Err(_) => {
-            log::error!("XDG_CONFIG_HOME has not been set.");
-            Path::new("/etc/swhkd/swhkdrc").to_path_buf()
-        }
-    };
-    config_file_path
-}
-
-pub fn fetch_xdg_runtime_socket_path() -> PathBuf {
-    match env::var("XDG_RUNTIME_DIR") {
-        Ok(val) => {
-            log::debug!("XDG_RUNTIME_DIR exists: {:#?}", val);
-            Path::new(&val).join("swhkd.sock")
-        }
-        Err(_) => {
-            log::error!("XDG_RUNTIME_DIR has not been set.");
-            Path::new(&format!("/run/user/{}/swhkd.sock", env::var("PKEXEC_UID").unwrap()))
-                .to_path_buf()
-        }
-    }
-}
-
-pub fn setup_swhkd(invoking_uid: u32) {
+pub fn setup_swhkd(invoking_uid: u32, runtime_path: String) {
     // Set a sane process umask.
     log::trace!("Setting process umask.");
     umask(Mode::S_IWGRP | Mode::S_IWOTH);
 
     // Get the runtime path and create it if needed.
-    let runtime_path: String = match env::var("XDG_RUNTIME_DIR") {
-        Ok(runtime_path) => {
-            log::debug!("XDG_RUNTIME_DIR exists: {:#?}", runtime_path);
-            Path::new(&runtime_path).join("swhkd").to_str().unwrap().to_owned()
-        }
-        Err(_) => {
-            log::error!("XDG_RUNTIME_DIR has not been set.");
-            String::from("/run/swhkd/")
-        }
-    };
     if !Path::new(&runtime_path).exists() {
         match fs::create_dir_all(Path::new(&runtime_path)) {
             Ok(_) => {
diff --git a/swhkd/src/environ.rs b/swhkd/src/environ.rs
@@ -0,0 +1,141 @@
+use std::{
+    env::VarError,
+    path::{Path, PathBuf},
+};
+
+pub struct Env {
+    pub pkexec_id: u32,
+    pub xdg_config_home: PathBuf,
+    pub xdg_runtime_socket: PathBuf,
+    pub xdg_runtime_dir: PathBuf,
+}
+
+#[derive(Debug)]
+pub enum EnvError {
+    PkexecNotFound,
+    XdgConfigNotFound,
+    XdgRuntimeNotFound,
+    PathNotFound,
+    GenericError(String),
+}
+
+impl Env {
+    pub fn construct() -> Self {
+        let pkexec_id = match Self::get_env("PKEXEC_UID") {
+            Ok(val) => match val.parse::<u32>() {
+                Ok(val) => val,
+                Err(_) => {
+                    log::error!("Failed to launch swhkd!!!");
+                    log::error!("Make sure to launch the binary with pkexec.");
+                    std::process::exit(1);
+                }
+            },
+            Err(_) => {
+                log::error!("Failed to launch swhkd!!!");
+                log::error!("Make sure to launch the binary with pkexec.");
+                std::process::exit(1);
+            }
+        };
+
+        let xdg_config_home = match Self::get_env("XDG_CONFIG_HOME") {
+            Ok(val) => match validate_path(&PathBuf::from(val)) {
+                Ok(val) => val,
+                Err(e) => match e {
+                    EnvError::PathNotFound => {
+                        log::warn!("XDG_CONFIG_HOME does not exist, using hardcoded /etc");
+                        PathBuf::from("/etc")
+                    }
+                    _ => {
+                        eprintln!("Failed to get XDG_CONFIG_HOME: {:?}", e);
+                        std::process::exit(1);
+                    }
+                },
+            },
+            Err(e) => match e {
+                EnvError::XdgConfigNotFound => {
+                    log::warn!("XDG_CONFIG_HOME not found, using hardcoded /etc");
+                    PathBuf::from("/etc")
+                }
+                _ => {
+                    eprintln!("Failed to get XDG_CONFIG_HOME: {:?}", e);
+                    std::process::exit(1);
+                }
+            },
+        };
+
+        let xdg_runtime_socket = match Self::get_env("XDG_RUNTIME_DIR") {
+            Ok(val) => match validate_path(&PathBuf::from(val).join("swhkd.sock")) {
+                Ok(val) => val,
+                Err(e) => match e {
+                    EnvError::PathNotFound => {
+                        log::warn!("XDG_RUNTIME_DIR does not exist, using hardcoded /run/user");
+                        PathBuf::from(format!("/run/user/{}", pkexec_id))
+                    }
+                    _ => {
+                        eprintln!("Failed to get XDG_RUNTIME_DIR: {:?}", e);
+                        std::process::exit(1);
+                    }
+                },
+            },
+            Err(e) => match e {
+                EnvError::XdgRuntimeNotFound => {
+                    log::warn!("XDG_RUNTIME_DIR not found, using hardcoded /run/user");
+                    PathBuf::from(format!("/run/user/{}", pkexec_id))
+                }
+                _ => {
+                    eprintln!("Failed to get XDG_RUNTIME_DIR: {:?}", e);
+                    std::process::exit(1);
+                }
+            },
+        };
+
+        let xdg_runtime_dir = match Self::get_env("XDG_RUNTIME_DIR") {
+            Ok(val) => PathBuf::from(val),
+            Err(e) => match e {
+                EnvError::XdgRuntimeNotFound => {
+                    log::warn!("XDG_RUNTIME_DIR not found, using hardcoded /run/swhkd");
+                    PathBuf::from("/run/swhkd")
+                }
+                _ => {
+                    eprintln!("Failed to get XDG_RUNTIME_DIR: {:?}", e);
+                    std::process::exit(1);
+                }
+            },
+        };
+
+        Self { pkexec_id, xdg_config_home, xdg_runtime_dir, xdg_runtime_socket }
+    }
+
+    fn get_env(name: &str) -> Result<String, EnvError> {
+        match std::env::var(name) {
+            Ok(val) => Ok(val),
+            Err(e) => match e {
+                VarError::NotPresent => match name {
+                    "PKEXEC_UID" => Err(EnvError::PkexecNotFound),
+                    "XDG_CONFIG_HOME" => Err(EnvError::XdgConfigNotFound),
+                    "XDG_RUNTIME_DIR" => Err(EnvError::XdgRuntimeNotFound),
+                    _ => Err(EnvError::GenericError(e.to_string())),
+                },
+                VarError::NotUnicode(_) => {
+                    Err(EnvError::GenericError("Not a valid unicode".to_string()))
+                }
+            },
+        }
+    }
+
+    pub fn fetch_xdg_config_path(&self) -> PathBuf {
+        PathBuf::from(&self.xdg_config_home).join("swhkd/swhkdrc")
+    }
+
+    pub fn fetch_xdg_runtime_socket_path(&self) -> PathBuf {
+        PathBuf::from(&self.xdg_runtime_dir).join("swhkd.sock")
+    }
+}
+
+fn validate_path(path: &Path) -> Result<PathBuf, EnvError> {
+    if path.exists() {
+        Ok(path.to_path_buf())
+    } else {
+        Err(EnvError::PathNotFound)
+    }
+}
diff --git a/swhks/src/environ.rs b/swhks/src/environ.rs
@@ -20,6 +20,7 @@ pub enum EnvError {
     DataHomeNotSet,
     HomeNotSet,
     RuntimeDirNotSet,
+    PathNotFound,
     GenericError(String),
 }
 
@@ -30,15 +31,15 @@ impl Env {
         let home = match Self::get_env("HOME") {
             Ok(val) => val,
             Err(_) => {
-                eprintln!("HOME Variable is not set, cannot fall back on hardcoded path for XDG_DATA_HOME.");
+                eprintln!("HOME Variable is not set/found, cannot fall back on hardcoded path for XDG_DATA_HOME.");
                 std::process::exit(1);
             }
         };
 
         let data_home = match Self::get_env("XDG_DATA_HOME") {
             Ok(val) => val,
             Err(e) => match e {
-                EnvError::DataHomeNotSet => {
+                EnvError::DataHomeNotSet | EnvError::PathNotFound => {
                     log::warn!(
                         "XDG_DATA_HOME Variable is not set, falling back on hardcoded path."
                     );
@@ -51,7 +52,7 @@ impl Env {
         let runtime_dir = match Self::get_env("XDG_RUNTIME_DIR") {
             Ok(val) => val,
             Err(e) => match e {
-                EnvError::RuntimeDirNotSet => {
+                EnvError::RuntimeDirNotSet | EnvError::PathNotFound => {
                     log::warn!(
                         "XDG_RUNTIME_DIR Variable is not set, falling back on hardcoded path."
                     );
@@ -67,7 +68,10 @@ impl Env {
     /// Actual interface to get the environment variable.
     fn get_env(name: &str) -> Result<PathBuf, EnvError> {
         match std::env::var(name) {
-            Ok(val) => Ok(PathBuf::from(val)),
+            Ok(val) => match PathBuf::from(&val).exists() {
+                true => Ok(PathBuf::from(val)),
+                false => Err(EnvError::PathNotFound),
+            },
             Err(e) => match e {
                 VarError::NotPresent => match name {
                     "XDG_DATA_HOME" => Err(EnvError::DataHomeNotSet),
