Introduce the chsh hook and the clean up of /etc/environment

distrod/distrod/src/main.rs

@@ -215,11 +215,15 @@ fn enable_wsl_exec_hook(opts: EnableOpts) -> Result<()> {
 fn disable_wsl_exec_hook(_opts: DisableOpts) -> Result<()> {
     shell_hook::disable_default_shell_hook()
         .with_context(|| "Failed to disable the hook to the default shell.")?;
+    if let Err(e) = distro::cleanup_distro_rootfs("/") {
+        log::warn!("Failed to clean up the rootfs: {:?}", e);
+    }
     log::info!("Distrod has been disabled. Now systemd will not start automatically.");
-    autostart::disable_autostart_on_windows_boot(
+    if let Err(e) = autostart::disable_autostart_on_windows_boot(
         &wsl_interop::get_distro_name().with_context(|| "Failed to get the distro name.")?,
-    )
-    .with_context(|| "Failed to disable the autostart on Windows boot.")?;
+    ) {
+        log::warn!("Failed to disable the autostart on Windows boot.: {:?}", e);
+    }
     Ok(())
 }
 

distrod/libs/src/container.rs

@@ -3,14 +3,14 @@ use nix::sched::CloneFlags;
 use nix::unistd::{chown, Gid, Uid};
 use nix::NixPath;
 use passfd::FdPassingExt;
-use std::ffi::{OsStr, OsString};
+use std::ffi::OsString;
 use std::fs::{self, File};
 use std::io::Write;
 use std::os::unix::io::AsRawFd;
 use std::os::unix::net::UnixStream;
 use std::os::unix::prelude::OsStrExt;
-use std::os::unix::process::CommandExt;
 use std::path::{Path, PathBuf};
+use std::process::Command;
 
 use crate::mount_info::{get_mount_entries, MountEntry};
 use crate::multifork::{CommandByMultiFork, Waiter};
@@ -62,8 +62,9 @@ impl Container {
 
         let (fd_channel_host, fd_channel_child) = UnixStream::pair()?;
         {
-            let mut command = CommandByMultiFork::new(&init[0]);
+            let mut command = Command::new(&init[0]);
             command.args(&init[1..]);
+            let mut command = CommandByMultiFork::new(command);
             let fds_to_keep = vec![fd_channel_child.as_raw_fd()];
             command.pre_second_fork(move || {
                 daemonize(&fds_to_keep)
@@ -111,34 +112,13 @@ impl Container {
         Ok(())
     }
 
-    pub fn exec_command<I, S, T1, T2, P>(
-        &self,
-        program: S,
-        args: I,
-        wd: Option<P>,
-        arg0: Option<T2>,
-        cred: Option<&Credential>,
-    ) -> Result<Waiter>
-    where
-        I: IntoIterator<Item = T1>,
-        S: AsRef<OsStr>,
-        T1: AsRef<OsStr>,
-        T2: AsRef<OsStr>,
-        P: AsRef<Path>,
-    {
+    pub fn exec_command(&self, command: Command, cred: Option<&Credential>) -> Result<Waiter> {
         log::debug!("Container::exec_command.");
         if self.init_pid.is_none() {
             bail!("This container is not launched yet.");
         }
 
-        let mut command = CommandByMultiFork::new(&program);
-        command.args(args);
-        if let Some(arg0) = arg0 {
-            command.arg0(arg0);
-        }
-        if let Some(wd) = wd {
-            command.current_dir(wd);
-        }
+        let mut command = CommandByMultiFork::new(command);
         command.pre_second_fork(|| {
             enter_namespace(self.init_procfile.as_ref().unwrap())
                 .with_context(|| "Failed to enter the init's namespace")?;
@@ -153,7 +133,7 @@ impl Container {
             .with_context(|| "Failed to request a proxy process.")?;
         command
             .spawn()
-            .with_context(|| format!("Container::exec_command failed: {:?}", &program.as_ref()))?;
+            .with_context(|| "Container::exec_command failed")?;
         log::debug!("Double fork done.");
         Ok(waiter)
     }

distrod/libs/src/distro.rs

@@ -1,12 +1,14 @@
 use anyhow::{anyhow, bail, Context, Result};
-use std::ffi::OsStr;
+use std::ffi::{OsStr, OsString};
 use std::fs::{self, File};
 use std::io::{BufReader, BufWriter, Write};
 use std::os::linux::fs::MetadataExt;
+use std::os::unix::prelude::CommandExt;
 use std::path::{Path, PathBuf};
+use std::process::Command;
 
 use crate::container::Container;
-use crate::distrod_config::DistrodConfig;
+use crate::distrod_config::{self, DistrodConfig};
 use crate::envfile::EnvFile;
 use crate::mount_info::get_mount_entries;
 pub use crate::multifork::Waiter;
@@ -106,8 +108,23 @@ impl Distro {
         P: AsRef<Path>,
     {
         log::debug!("Distro::exec_command.");
+        let mut command = Command::new(command.as_ref());
+        command
+            .args(args)
+            // Adding the path to distrod bin allows us to hook "chsh" command to show the message
+            // to ask users to run "distrod enable" command.
+            .env(
+                "PATH",
+                add_distrod_bin_to_path(std::env::var("PATH").unwrap_or_default()),
+            );
+        if let Some(wd) = wd {
+            command.current_dir(wd.as_ref());
+        }
+        if let Some(arg0) = arg0 {
+            command.arg0(arg0.as_ref());
+        }
         self.container
-            .exec_command(command, args, wd, arg0, cred)
+            .exec_command(command, cred)
             .with_context(|| "Failed to exec command in the container")
     }
 
@@ -198,6 +215,18 @@ pub fn initialize_distro_rootfs<P: AsRef<Path>>(
     Ok(())
 }
 
+pub fn cleanup_distro_rootfs<P: AsRef<Path>>(path: P) -> Result<()> {
+    let metadata = fs::metadata(path.as_ref())?;
+    if !metadata.is_dir() {
+        bail!("The given path is not a directory: '{:?}'", path.as_ref());
+    }
+
+    cleanup_etc_environment_file(path.as_ref())
+        .with_context(|| "Failed to cleanup /etc/environment")?;
+
+    Ok(())
+}
+
 fn setup_etc_environment_file<P: AsRef<Path>>(rootfs_path: P) -> Result<()> {
     let env_file_path = rootfs_path.as_ref().join("etc/environment");
     let mut env_file = EnvFile::open(&env_file_path)
@@ -208,12 +237,75 @@ fn setup_etc_environment_file<P: AsRef<Path>>(rootfs_path: P) -> Result<()> {
     for (key, value) in &wsl_envs {
         env_file.put(&key.to_string_lossy(), value.to_string_lossy().to_string());
     }
+
+    // Put the Distrod's bin dir in PATH
+    // This allows us to hook "chsh" command to show the message to ask users to run "distrod enable" command
+    if env_file
+        .get("PATH")
+        .map(|path| path.contains(distrod_config::get_distrod_bin_dir_path()))
+        != Some(true)
+    {
+        env_file.put(
+            "PATH",
+            add_distrod_bin_to_path(env_file.get("PATH").unwrap_or(""))
+                .to_string_lossy()
+                .to_string(),
+        );
+    }
+
+    env_file
+        .save()
+        .with_context(|| "Failed to save the environment file.")?;
+    Ok(())
+}
+
+fn cleanup_etc_environment_file<P: AsRef<Path>>(rootfs_path: P) -> Result<()> {
+    let env_file_path = rootfs_path.as_ref().join("etc/environment");
+    let mut env_file = EnvFile::open(&env_file_path)
+        .with_context(|| format!("Failed to open '{:?}'.", &&env_file_path))?;
+
+    // Set the WSL envs in the default environment variables
+    let wsl_envs = collect_wsl_env_vars().with_context(|| "Failed to collect WSL envs.")?;
+    let keys: Vec<_> = wsl_envs.keys().collect();
+    for key in keys {
+        env_file.remove(&key.to_string_lossy());
+    }
+
+    // Put the Distrod's bin dir in PATH
+    // This allows us to hook "chsh" command to show the message to ask users to run "distrod enable" command
+    if env_file
+        .get("PATH")
+        .map(|path| path.contains(distrod_config::get_distrod_bin_dir_path()))
+        == Some(true)
+    {
+        env_file.put(
+            "PATH",
+            remove_distrod_bin_from_path(env_file.get("PATH").unwrap_or("")),
+        );
+    }
+
     env_file
         .save()
         .with_context(|| "Failed to save the environment file.")?;
     Ok(())
 }
 
+fn add_distrod_bin_to_path<S: AsRef<OsStr>>(path: S) -> OsString {
+    let mut result = OsString::from(distrod_config::get_distrod_bin_dir_path());
+    result.push(":");
+    result.push(path);
+    result
+}
+
+fn remove_distrod_bin_from_path(path: &str) -> String {
+    let mut distrod_bin_path = distrod_config::get_distrod_bin_dir_path().to_owned();
+    distrod_bin_path.push(':');
+    let result = path.replace(&distrod_bin_path, "");
+    distrod_bin_path.pop();
+    distrod_bin_path.insert(0, ':');
+    result.replace(&distrod_bin_path, "")
+}
+
 fn get_distro_run_info_file(create: bool, write: bool) -> Result<Option<File>> {
     let mut json = fs::OpenOptions::new();
     json.read(true);

distrod/libs/src/distrod_config.rs

@@ -69,15 +69,23 @@ pub fn get_alias_dir() -> &'static str {
     DISTROD_ALIAS_DIR.as_str()
 }
 
-static DISTROD_BIN_PATH: Lazy<String> = Lazy::new(|| format!("{}/{}", DISTROD_ROOT_DIR, "distrod"));
+static DISTROD_BIN_DIR: Lazy<String> = Lazy::new(|| format!("{}/{}", DISTROD_ROOT_DIR, "bin"));
+
+/// The path to the distrod binary.
+pub fn get_distrod_bin_dir_path() -> &'static str {
+    DISTROD_BIN_DIR.as_str()
+}
+
+static DISTROD_BIN_PATH: Lazy<String> =
+    Lazy::new(|| format!("{}/{}", DISTROD_BIN_DIR.as_str(), "distrod"));
 
 /// The path to the distrod binary.
 pub fn get_distrod_bin_path() -> &'static str {
     DISTROD_BIN_PATH.as_str()
 }
 
 static DISTROD_EXEC_BIN_PATH: Lazy<String> =
-    Lazy::new(|| format!("{}/{}", DISTROD_ROOT_DIR, "distrod-exec"));
+    Lazy::new(|| format!("{}/{}", DISTROD_BIN_DIR.as_str(), "distrod-exec"));
 
 /// The path to the distrod-exec binary.
 pub fn get_distrod_exec_bin_path() -> &'static str {

distrod/libs/src/envfile.rs

@@ -67,6 +67,10 @@ impl EnvFile {
         self.envs.insert(key.to_owned(), vec![(usize::MAX, value)]);
     }
 
+    pub fn remove(&mut self, key: &str) {
+        self.envs.remove(key);
+    }
+
     pub fn save(&mut self) -> Result<()> {
         let mut lines = self
             .envs

distrod/libs/src/multifork.rs

@@ -3,11 +3,11 @@ use nix::fcntl::OFlag;
 use nix::libc::c_int;
 use nix::sys::signal;
 use std::convert::From;
-use std::ffi::OsStr;
 use std::fs::File;
 use std::io::{Read, Write};
-use std::ops::{Deref, DerefMut};
+use std::ops::Deref;
 use std::os::unix::io::FromRawFd;
+use std::os::unix::prelude::CommandExt;
 use std::process::Command;
 
 pub struct CommandByMultiFork<'a> {
@@ -18,9 +18,9 @@ pub struct CommandByMultiFork<'a> {
 }
 
 impl<'a> CommandByMultiFork<'a> {
-    pub fn new<S: AsRef<OsStr>>(program: S) -> CommandByMultiFork<'a> {
+    pub fn new(command: Command) -> CommandByMultiFork<'a> {
         CommandByMultiFork {
-            command: Command::new(program),
+            command,
             pre_second_fork: None,
             proxy_process: None,
             does_triple_fork: false,
@@ -32,6 +32,15 @@ impl<'a> CommandByMultiFork<'a> {
         self
     }
 
+    // Define proxy function to allow it to be called before pre_second_fork for readability.
+    pub unsafe fn pre_exec<F>(&mut self, f: F) -> &mut CommandByMultiFork<'a>
+    where
+        F: FnMut() -> std::io::Result<()> + Send + Sync + 'static,
+    {
+        self.command.pre_exec(f);
+        self
+    }
+
     pub fn pre_second_fork<F>(&mut self, f: F) -> &mut CommandByMultiFork<'a>
     where
         F: FnMut() -> Result<()> + 'a,
@@ -97,12 +106,6 @@ impl<'a> Deref for CommandByMultiFork<'a> {
     }
 }
 
-impl<'a> DerefMut for CommandByMultiFork<'a> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.command
-    }
-}
-
 impl<'a> From<Command> for CommandByMultiFork<'a> {
     fn from(command: Command) -> Self {
         CommandByMultiFork {
@@ -198,8 +201,9 @@ mod tests {
 
     #[test]
     fn test_insert_proxy() {
-        let mut doublefork = CommandByMultiFork::new("/bin/bash");
-        doublefork.args(&["-c", "sleep 1; exit 42"]);
+        let mut command = Command::new("/bin/bash");
+        command.args(&["-c", "sleep 1; exit 42"]);
+        let mut doublefork = CommandByMultiFork::new(command);
         let mut waiter = doublefork.insert_waiter_proxy().unwrap();
         let _ = doublefork.spawn().unwrap();
         let exit_code = waiter.wait();
@@ -208,11 +212,12 @@ mod tests {
 
     #[test]
     fn test_inserted_proxy_ignore_signal() {
-        let mut doublefork = CommandByMultiFork::new("/bin/bash");
-        doublefork.args(&[
+        let mut command = Command::new("/bin/bash");
+        command.args(&[
             "-c",
             "trap '' SIGINT; kill -SIGINT $PPID; sleep 1; exit 42;",
         ]);
+        let mut doublefork = CommandByMultiFork::new(command);
         let mut waiter = doublefork.insert_waiter_proxy().unwrap();
         let _ = doublefork.spawn().unwrap();
         let exit_code = waiter.wait();