diff --git a/src/command_runner.rs b/src/command_runner.rs index e82bba6..981e934 100644 --- a/src/command_runner.rs +++ b/src/command_runner.rs @@ -38,16 +38,31 @@ use std::os::unix::process::CommandExt; use std::env; use users::get_user_by_name; +struct TempSetEnv<'a> { name: &'a str, old_value: Option } + +impl<'a> TempSetEnv<'a> { + fn new(name: &'a str, new_value: String) -> TempSetEnv<'a> { + let old_value = env::var(name); + env::set_var(name, new_value); + TempSetEnv { name: name, old_value: old_value.ok() } + } +} + +impl<'a> Drop for TempSetEnv<'a> { + fn drop(&mut self) { + match self.old_value { + Some(ref val) => env::set_var(self.name, val), + None => env::remove_var(self.name) + } + } +} + impl<'a, C> CommandRunner for UserCommandRunner<'a, C> where C: 'a + CommandRunner { fn run_with_args(&self, program: &str, args: &[&str]) -> IoResult { let uid = get_user_by_name(self.user_name).unwrap().uid(); - let old_home = env::var("HOME"); - env::set_var("HOME", format!("/home/{}", self.user_name)); - let old_dbus = env::var("DBUS_SESSION_BUS_ADDRESS"); - env::set_var("DBUS_SESSION_BUS_ADDRESS", format!("unix:path=/run/user/{}/bus", uid)); + let set_home = TempSetEnv::new("HOME", format!("/home/{}", self.user_name)); + let set_dbus = TempSetEnv::new("DBUS_SESSION_BUS_ADDRESS", format!("unix:path=/run/user/{}/bus", uid)); let res = Command::new(program).uid(uid).gid(uid).args(args).output(); - if let Ok(dbus) = old_dbus { env::set_var("DBUS_SESSION_BUS_ADDRESS", dbus); } - if let Ok(home) = old_home { env::set_var("HOME", home); } res } }