Add new setup_envs method for the Executor trait (#2820)

* Add set_envs to the Executor trait Signed-off-by: Kotaro Inoue <[email protected]> * Use set_envs in container_init_process Signed-off-by: Kotaro Inoue <[email protected]> * Fix linter issues Signed-off-by: Kotaro Inoue <[email protected]> * Add test of default set_envs implementation Signed-off-by: Kotaro Inoue <[email protected]> * Fix set_envs to return ExecuteSetEnvsError Signed-off-by: Kotaro Inoue <[email protected]> * Rename set_envs to setup_envs Signed-off-by: Kotaro Inoue <[email protected]> * Add detailed description for the setup_envs Signed-off-by: Kotaro Inoue <[email protected]> --------- Signed-off-by: Kotaro Inoue <[email protected]>
youki-dev · Jun 27, 2024 · bdd5aab · bdd5aab
1 parent 11ca60d
commit bdd5aab
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 12 deletions.
diff --git a/crates/libcontainer/src/process/container_init_process.rs b/crates/libcontainer/src/process/container_init_process.rs
@@ -73,6 +73,8 @@ pub enum InitProcessError {
     Workload(#[from] workload::ExecutorError),
     #[error(transparent)]
     WorkloadValidation(#[from] workload::ExecutorValidationError),
+    #[error(transparent)]
+    WorkloadSetEnvs(#[from] workload::ExecutorSetEnvsError),
     #[error("invalid io priority class: {0}")]
     IoPriorityClass(String),
     #[error("call exec sched_setattr error: {0}")]
@@ -558,18 +560,6 @@ pub fn container_init_process(
         err
     })?;
 
-    // add HOME into envs if not exists
-    if !envs.contains_key("HOME") {
-        if let Some(dir_home) = utils::get_user_home(proc.user().uid()) {
-            envs.insert("HOME".to_owned(), dir_home.to_string_lossy().to_string());
-        }
-    }
-
-    // Reset the process env based on oci spec.
-    env::vars().for_each(|(key, _value)| env::remove_var(key));
-    envs.iter()
-        .for_each(|(key, value)| env::set_var(key, value));
-
     // Initialize seccomp profile right before we are ready to execute the
     // payload so as few syscalls will happen between here and payload exec. The
     // notify socket will still need network related syscalls.
@@ -591,7 +581,15 @@ pub fn container_init_process(
         tracing::warn!("seccomp not available, unable to set seccomp privileges!")
     }
 
+    // add HOME into envs if not exists
+    if !envs.contains_key("HOME") {
+        if let Some(dir_home) = utils::get_user_home(proc.user().uid()) {
+            envs.insert("HOME".to_owned(), dir_home.to_string_lossy().to_string());
+        }
+    }
+
     args.executor.validate(spec)?;
+    args.executor.setup_envs(envs)?;
 
     // Notify main process that the init process is ready to execute the
     // payload.  Note, because we are already inside the pid namespace, the pid

diff --git a/crates/libcontainer/src/workload/default.rs b/crates/libcontainer/src/workload/default.rs
@@ -138,6 +138,11 @@ fn is_executable(path: &Path) -> std::result::Result<bool, std::io::Error> {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+    use std::env;
+
+    use serial_test::serial;
+
     use super::*;
 
     #[test]
@@ -177,4 +182,35 @@ mod tests {
         assert!(!is_executable(&non_executable_path).unwrap());
         assert!(!is_executable(directory_path).unwrap());
     }
+
+    #[test]
+    #[serial]
+    fn test_executor_set_envs() {
+        // Store original environment variables to restore later
+        let original_envs: HashMap<String, String> = env::vars().collect();
+
+        // Test setting environment variables
+        {
+            let executor = get_executor();
+            let envs = HashMap::from([
+                ("FOO".to_owned(), "hoge".to_owned()),
+                ("BAR".to_owned(), "fuga".to_owned()),
+                ("BAZ".to_owned(), "piyo".to_owned()),
+            ]);
+            assert!(executor.setup_envs(envs).is_ok());
+
+            // Check if the environment variables are set correctly
+            let current_envs = std::env::vars().collect::<HashMap<String, String>>();
+            assert_eq!(current_envs.get("FOO").unwrap(), "hoge");
+            assert_eq!(current_envs.get("BAR").unwrap(), "fuga");
+            assert_eq!(current_envs.get("BAZ").unwrap(), "piyo");
+            // No other environment variables should be set
+            assert_eq!(current_envs.len(), 3);
+        }
+
+        // Restore original environment variables
+        original_envs.iter().for_each(|(key, value)| {
+            env::set_var(key, value);
+        });
+    }
 }
diff --git a/crates/libcontainer/src/workload/mod.rs b/crates/libcontainer/src/workload/mod.rs
@@ -1,3 +1,6 @@
+use std::collections::HashMap;
+use std::env;
+
 use oci_spec::runtime::Spec;
 
 pub mod default;
@@ -24,6 +27,14 @@ pub enum ExecutorValidationError {
     ArgValidationError(String),
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum ExecutorSetEnvsError {
+    #[error("failed to set envs")]
+    SetEnvs(#[from] Box<dyn std::error::Error + Send + Sync>),
+    #[error("{0}")]
+    Other(String),
+}
+
 // Here is an explanation about the complexity below regarding to
 // CloneBoxExecutor and Executor traits. This is one of the places rust actually
 // makes our life harder. The usecase for the executor is to allow users of
@@ -60,6 +71,24 @@ pub trait Executor: CloneBoxExecutor {
     /// namespace and cgroups, and pivot_root into the rootfs. But this step
     /// runs before waiting for the container start signal.
     fn validate(&self, spec: &Spec) -> Result<(), ExecutorValidationError>;
+
+    /// Set environment variables for the container process to be executed.
+    /// This step runs after the container init process is created, entered
+    /// into the correct namespace and cgroups, and pivot_root into the rootfs.
+    /// But this step runs before waiting for the container start signal.
+    /// The host's environment variables are not cleared yet at this point.
+    /// They should be cleared explicitly if needed.
+    fn setup_envs(&self, envs: HashMap<String, String>) -> Result<(), ExecutorSetEnvsError> {
+        // The default implementation resets the process env based on the OCI spec.
+        // First, clear all host's envs.
+        env::vars().for_each(|(key, _value)| env::remove_var(key));
+
+        // Next, set envs based on the spec
+        envs.iter()
+            .for_each(|(key, value)| env::set_var(key, value));
+
+        Ok(())
+    }
 }
 
 impl<T> CloneBoxExecutor for T