Refactored some functions and added a few tests.

src/greeter.rs

@@ -165,20 +165,17 @@ impl Greeter {
 
     if greeter.remember_session {
       if let Ok(session_path) = get_last_session_path() {
-        greeter.session_path = Some(session_path);
-      }
+        greeter.session_path = Some(session_path.clone());
 
-      if let Ok(command) = get_last_session() {
+        if let Some(session) = Session::from_path(&greeter, session_path) {
+          greeter.command = Some(session.command.clone());
+        }
+      } else if let Ok(command) = get_last_session() {
         greeter.command = Some(command.trim().to_string());
       }
     }
 
-    greeter.sessions.selected = greeter
-      .sessions
-      .options
-      .iter()
-      .position(|Session { path, .. }| path.as_deref() == greeter.session_path.as_deref())
-      .unwrap_or(0);
+    greeter.sessions.selected = greeter.sessions.options.iter().position(|Session { path, .. }| path == &greeter.session_path).unwrap_or(0);
 
     greeter
   }

src/info.rs

@@ -112,11 +112,11 @@ pub fn write_last_username(username: &str, name: Option<&str>) {
 }
 
 pub fn get_last_session_path() -> Result<PathBuf, io::Error> {
-  Ok(PathBuf::from(fs::read_to_string(LAST_SESSION_PATH)?))
+  Ok(PathBuf::from(fs::read_to_string(LAST_SESSION_PATH)?.trim()))
 }
 
 pub fn get_last_session() -> Result<String, io::Error> {
-  fs::read_to_string(LAST_SESSION)
+  Ok(fs::read_to_string(LAST_SESSION)?.trim().to_string())
 }
 
 pub fn write_last_session_path<P>(session: &P)
@@ -131,11 +131,11 @@ pub fn write_last_session(session: &str) {
 }
 
 pub fn get_last_user_session_path(username: &str) -> Result<PathBuf, io::Error> {
-  Ok(PathBuf::from(fs::read_to_string(format!("{LAST_SESSION_PATH}-{username}"))?))
+  Ok(PathBuf::from(fs::read_to_string(format!("{LAST_SESSION_PATH}-{username}"))?.trim()))
 }
 
 pub fn get_last_user_session(username: &str) -> Result<String, io::Error> {
-  fs::read_to_string(format!("{LAST_SESSION}-{username}"))
+  Ok(fs::read_to_string(format!("{LAST_SESSION}-{username}"))?.trim().to_string())
 }
 
 pub fn write_last_user_session_path<P>(username: &str, session: P)
@@ -153,6 +153,10 @@ pub fn write_last_user_session(username: &str, session: &str) {
   let _ = fs::write(format!("{LAST_SESSION}-{username}"), session);
 }
 
+pub fn delete_last_user_session_path(username: &str) {
+  let _ = fs::remove_file(format!("{LAST_SESSION_PATH}-{username}"));
+}
+
 pub fn get_users(min_uid: u16, max_uid: u16) -> Vec<User> {
   match File::open("/etc/passwd") {
     Err(_) => vec![],

src/ipc.rs

@@ -1,4 +1,4 @@
-use std::{error::Error, sync::Arc};
+use std::{borrow::Cow, error::Error, sync::Arc};
 
 use greetd_ipc::{codec::TokioCodec, AuthMessageType, ErrorType, Request, Response};
 use tokio::sync::{
@@ -7,7 +7,7 @@ use tokio::sync::{
 };
 
 use crate::{
-  info::{write_last_user_session, write_last_user_session_path, write_last_username},
+  info::{delete_last_user_session_path, write_last_user_session, write_last_user_session_path, write_last_username},
   ui::sessions::{Session, SessionType},
   AuthStatus, Greeter, Mode,
 };
@@ -113,6 +113,8 @@ impl Ipc {
 
               if let Some(session_path) = &greeter.session_path {
                 write_last_user_session_path(&greeter.username, session_path);
+              } else {
+                delete_last_user_session_path(&greeter.username);
               }
             }
           }
@@ -122,28 +124,11 @@ impl Ipc {
           greeter.done = true;
           greeter.mode = Mode::Processing;
 
-          let session = greeter.sessions.options.get(greeter.sessions.selected).filter(|s| &s.command == command);
-          let mut command = command.clone();
-          let mut env = vec![];
-
-          if let Some(Session { session_type, .. }) = session {
-            if *session_type != SessionType::None {
-              env.push(format!("XDG_SESSION_TYPE={}", session_type.as_xdg_session_type()));
-            }
-
-            if *session_type == SessionType::X11 {
-              if let Some(ref wrap) = greeter.xsession_wrapper {
-                command = format!("{} {}", wrap, command);
-              }
-            } else if let Some(ref wrap) = greeter.session_wrapper {
-              command = format!("{} {}", wrap, command);
-            }
-          } else if let Some(ref wrap) = greeter.session_wrapper {
-            command = format!("{} {}", wrap, command);
-          }
+          let session = Session::get_selected(greeter);
+          let (command, env) = wrap_session_command(greeter, session, command);
 
           #[cfg(not(debug_assertions))]
-          self.send(Request::StartSession { cmd: vec![command], env }).await;
+          self.send(Request::StartSession { cmd: vec![command.to_string()], env }).await;
 
           #[cfg(debug_assertions)]
           {
@@ -185,3 +170,92 @@ impl Ipc {
     let _ = Request::CancelSession.write_to(&mut *greeter.stream().await).await;
   }
 }
+
+fn wrap_session_command<'a>(greeter: &Greeter, session: Option<&Session>, command: &'a str) -> (Cow<'a, str>, Vec<String>) {
+  let mut env: Vec<String> = vec![];
+
+  if let Some(Session { session_type, .. }) = session {
+    if *session_type != SessionType::None {
+      env.push(format!("XDG_SESSION_TYPE={}", session_type.as_xdg_session_type()));
+    }
+
+    if *session_type == SessionType::X11 {
+      if let Some(ref wrap) = greeter.xsession_wrapper {
+        return (Cow::Owned(format!("{} {}", wrap, command)), env);
+      }
+    } else if let Some(ref wrap) = greeter.session_wrapper {
+      return (Cow::Owned(format!("{} {}", wrap, command)), env);
+    }
+  } else if let Some(ref wrap) = greeter.session_wrapper {
+    return (Cow::Owned(format!("{} {}", wrap, command)), env);
+  }
+
+  (Cow::Borrowed(command), env)
+}
+
+#[cfg(test)]
+mod test {
+  use std::path::PathBuf;
+
+  use crate::{
+    ui::sessions::{Session, SessionType},
+    Greeter,
+  };
+
+  use super::wrap_session_command;
+
+  #[test]
+  fn wayland_no_wrapper() {
+    let greeter = Greeter::default();
+
+    let session = Session {
+      name: "Session1".into(),
+      session_type: SessionType::Wayland,
+      command: "Session1Cmd".into(),
+      path: Some(PathBuf::from("/Session1Path")),
+    };
+
+    let (command, env) = wrap_session_command(&greeter, Some(&session), &session.command);
+
+    assert_eq!(command.as_ref(), "Session1Cmd");
+    assert_eq!(env, vec!["XDG_SESSION_TYPE=wayland"]);
+  }
+
+  #[test]
+  fn wayland_wrapper() {
+    let mut greeter = Greeter::default();
+    greeter.session_wrapper = Some("/wrapper.sh".into());
+
+    let session = Session {
+      name: "Session1".into(),
+      session_type: SessionType::Wayland,
+      command: "Session1Cmd".into(),
+      path: Some(PathBuf::from("/Session1Path")),
+    };
+
+    let (command, env) = wrap_session_command(&greeter, Some(&session), &session.command);
+
+    assert_eq!(command.as_ref(), "/wrapper.sh Session1Cmd");
+    assert_eq!(env, vec!["XDG_SESSION_TYPE=wayland"]);
+  }
+
+  #[test]
+  fn x11_wrapper() {
+    let mut greeter = Greeter::default();
+    greeter.xsession_wrapper = Some("startx /usr/bin/env".into());
+
+    println!("{:?}", greeter.xsession_wrapper);
+
+    let session = Session {
+      name: "Session1".into(),
+      session_type: SessionType::X11,
+      command: "Session1Cmd".into(),
+      path: Some(PathBuf::from("/Session1Path")),
+    };
+
+    let (command, env) = wrap_session_command(&greeter, Some(&session), &session.command);
+
+    assert_eq!(command.as_ref(), "startx /usr/bin/env Session1Cmd");
+    assert_eq!(env, vec!["XDG_SESSION_TYPE=x11"]);
+  }
+}

src/keyboard.rs

@@ -175,6 +175,7 @@ pub async fn handle(greeter: Arc<RwLock<Greeter>>, input: KeyEvent, ipc: Ipc) ->
 
       Mode::Command => {
         greeter.sessions.selected = 0;
+        greeter.session_path = None;
         greeter.command = Some(greeter.new_command.clone());
 
         if greeter.remember_session {
@@ -303,12 +304,11 @@ async fn validate_username(greeter: &mut Greeter, ipc: &Ipc) {
 
   if greeter.remember_user_session {
     if let Ok(last_session) = get_last_user_session_path(&greeter.username) {
-      greeter.sessions.selected = greeter
-        .sessions
-        .options
-        .iter()
-        .position(|Session { path, .. }| path.as_deref() == Some(last_session.as_ref()))
-        .unwrap_or(0);
+      if let Some(last_session) = Session::from_path(greeter, last_session).cloned() {
+        greeter.sessions.selected = greeter.sessions.options.iter().position(|sess| sess.path == last_session.path).unwrap_or(0);
+        greeter.session_path = last_session.path;
+        greeter.command = Some(last_session.command);
+      }
     }
 
     if let Ok(command) = get_last_user_session(&greeter.username) {

src/ui/sessions.rs

@@ -1,4 +1,6 @@
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
+
+use crate::Greeter;
 
 use super::common::menu::MenuItem;
 
@@ -35,3 +37,150 @@ impl MenuItem for Session {
     self.name.clone()
   }
 }
+
+impl Session {
+  pub fn from_path<P>(greeter: &Greeter, path: P) -> Option<&Session>
+  where
+    P: AsRef<Path>,
+  {
+    greeter.sessions.options.iter().find(|session| session.path.as_deref() == Some(path.as_ref()))
+  }
+
+  pub fn get_selected(greeter: &Greeter) -> Option<&Session> {
+    greeter.session_path.as_ref()?;
+    greeter.sessions.options.get(greeter.sessions.selected)
+  }
+}
+
+#[cfg(test)]
+mod test {
+  use crate::{
+    ui::{
+      common::menu::Menu,
+      sessions::{Session, SessionType},
+    },
+    Greeter,
+  };
+
+  #[test]
+  fn from_path_existing() {
+    let mut greeter = Greeter::default();
+    greeter.session_path = Some("/Session2Path".into());
+
+    greeter.sessions = Menu::<Session> {
+      title: "Sessions".into(),
+      selected: 1,
+      options: vec![
+        Session {
+          name: "Session1".into(),
+          command: "Session1Cmd".into(),
+          session_type: super::SessionType::Wayland,
+          path: Some("/Session1Path".into()),
+        },
+        Session {
+          name: "Session2".into(),
+          command: "Session2Cmd".into(),
+          session_type: super::SessionType::X11,
+          path: Some("/Session2Path".into()),
+        },
+      ],
+    };
+
+    let session = Session::from_path(&greeter, "/Session2Path");
+
+    assert!(matches!(session, Some(_)));
+    assert_eq!(session.unwrap().name, "Session2");
+    assert_eq!(session.unwrap().session_type, SessionType::X11);
+  }
+
+  #[test]
+  fn from_path_non_existing() {
+    let mut greeter = Greeter::default();
+    greeter.session_path = Some("/Session2Path".into());
+
+    greeter.sessions = Menu::<Session> {
+      title: "Sessions".into(),
+      selected: 1,
+      options: vec![Session {
+        name: "Session1".into(),
+        command: "Session1Cmd".into(),
+        session_type: super::SessionType::Wayland,
+        path: Some("/Session1Path".into()),
+      }],
+    };
+
+    let session = Session::from_path(&greeter, "/Session2Path");
+
+    assert!(matches!(session, None));
+  }
+
+  #[test]
+  fn no_session() {
+    let greeter = Greeter::default();
+
+    assert!(matches!(Session::get_selected(&greeter), None));
+  }
+
+  #[test]
+  fn distinct_session() {
+    let mut greeter = Greeter::default();
+    greeter.session_path = Some("/Session2Path".into());
+
+    greeter.sessions = Menu::<Session> {
+      title: "Sessions".into(),
+      selected: 1,
+      options: vec![
+        Session {
+          name: "Session1".into(),
+          command: "Session1Cmd".into(),
+          session_type: super::SessionType::Wayland,
+          path: Some("/Session1Path".into()),
+        },
+        Session {
+          name: "Session2".into(),
+          command: "Session2Cmd".into(),
+          session_type: super::SessionType::X11,
+          path: Some("/Session2Path".into()),
+        },
+      ],
+    };
+
+    let session = Session::get_selected(&greeter);
+
+    assert!(matches!(session, Some(_)));
+    assert_eq!(session.unwrap().name, "Session2");
+    assert_eq!(session.unwrap().session_type, SessionType::X11);
+  }
+
+  #[test]
+  fn same_name_session() {
+    let mut greeter = Greeter::default();
+    greeter.session_path = Some("/Session2Path".into());
+
+    greeter.sessions = Menu::<Session> {
+      title: "Sessions".into(),
+      selected: 1,
+      options: vec![
+        Session {
+          name: "Session".into(),
+          command: "Session1Cmd".into(),
+          session_type: super::SessionType::Wayland,
+          path: Some("/Session1Path".into()),
+        },
+        Session {
+          name: "Session".into(),
+          command: "Session2Cmd".into(),
+          session_type: super::SessionType::X11,
+          path: Some("/Session2Path".into()),
+        },
+      ],
+    };
+
+    let session = Session::get_selected(&greeter);
+
+    assert!(matches!(session, Some(_)));
+    assert_eq!(session.unwrap().name, "Session");
+    assert_eq!(session.unwrap().session_type, SessionType::X11);
+    assert_eq!(session.unwrap().command, "Session2Cmd");
+  }
+}