Extract TempSetEnv

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<String> }
+
+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<Output> {
     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
   }
 }