env 参数完成

src/main.rs

@@ -1,6 +1,7 @@
 use std::arch::asm;
 use std::ffi::{CStr, CString};
 use std::fs::File;
+use std::io::Read;
 use std::os::fd::{AsFd, AsRawFd};
 use std::os::unix::fs::PermissionsExt;
 use nix::libc::{self, setgid, CLONE_NEWCGROUP, MS_NODEV, MS_NOSUID};
@@ -36,8 +37,10 @@ struct RockerArgs {
     #[arg(long)]
     image: Option<String>,
     #[arg(long)]
-    volume: Option<String>,
+    // --volume "/tmp/test1:tmp/test1,/tmp/test2:tmp/test2"
+    volume: Option<String>,  
     #[arg(long)]
+    // --env "a=1,b=2,c=3"
     env: Option<String>,
 
     // --run /bin/bash --exec container_id
@@ -60,11 +63,11 @@ struct RockerArgs {
     #[arg(long)]
     psa: bool,
     
-    // rm container_id
+    // rm "container_id_1, container_id_2, container_id_3"
     #[arg(long)]
     rm: Option<String>,
 
-    // stop container_id
+    // stop "container_id_1, container_id_2, container_id_3"
     #[arg(long)]
     stop: Option<String>
     
@@ -173,7 +176,7 @@ fn init_container_custom_volume<P: AsRef<Path>>(container_merged_path: P, custom
             .arg(host_path)
             .arg(container_path)
             .output()?;
-        let std_out = String::from_utf8_lossy(&out.stdout);
+        // let std_out = String::from_utf8_lossy(&out.stdout);
         let std_err = String::from_utf8_lossy(&out.stderr);
         if std_err.len() == 0 {
             println!("创建自定义 volume: {custom_volume:?}");
@@ -185,6 +188,39 @@ fn init_container_custom_volume<P: AsRef<Path>>(container_merged_path: P, custom
 }
 
 
+fn init_container_env(env: Option<&String>) -> Result<()>{
+    for (k, _) in std::env::vars(){
+        std::env::remove_var(k);
+    }
+
+    if let Some(env) = env {
+        let env_vec = if env.starts_with("./") || env.starts_with("/") {
+            // 读取出路径指定的文件作为env
+            let env_path = Path::new(env);
+            let mut env_file = File::open(env_path)?;
+            let text = {
+                let mut s = String::new();
+                env_file.read_to_string(&mut s)?;
+                s
+            };
+            text.lines().map(String::from).collect::<Vec<String>>()
+        } else {
+            env.split(",").map(String::from).collect::<Vec<String>>()
+        };
+
+        for item_env in env_vec.iter() {
+            let item_env_v = item_env.split("=").collect::<Vec<&str>>();
+            if item_env_v.len() == 2 {
+                std::env::set_var(item_env_v[0], item_env_v[1])
+            } else {
+                println!("env 格式不正确: {item_env}")
+            }
+        }
+    }
+    Ok(())
+}
+
+
 fn init_container_pivot<P: AsRef<Path>>(merged_path: P) -> Result<()> {
     // 在我们没有设置 chroot之前, 需要先把所有挂载点的传播类型改为 private, 避免进程中的系统调用污染全局
     mount(None::<&str>, "/", None::<&str>, MsFlags::MS_PRIVATE | MsFlags::MS_REC, None::<&str>)?;
@@ -201,8 +237,7 @@ fn init_container_pivot<P: AsRef<Path>>(merged_path: P) -> Result<()> {
     let pivot_root_dir = format!("{pwd_str}/.pivot_root");
 
     // 将系统rootfs切换到新的rootfs, 并设置权限
-    std::fs::create_dir(&pivot_root_dir)?;
-    std::fs::set_permissions(&pivot_root_dir, PermissionsExt::from_mode(0o777))?;
+    create_dir(&pivot_root_dir, true)?;
     pivot_root(pwd_str.as_str(), pivot_root_dir.as_str())?;
 
     // 修改当前进程工作目录(注意我们之前已经到rootfs内, 并且把根目录设置完毕了)
@@ -239,12 +274,7 @@ fn init_container_log(log: bool) -> Result<()> {
     Ok(())
 }
 
-fn init_container_env() -> Result<()> {
-    unsafe {
-        clearenv().map_err(|e|RockerError::OtherError(format!("清除env失败: {e:?}")))?;
-    }
-    Ok(())
-}
+
 fn init_container_user(uid: Uid, gid: Gid) -> Result<()>{
     unsafe {
         setgid(gid.as_raw());
@@ -352,14 +382,6 @@ fn run_container(_container_id: &String, cmd: &String, args: &RockerArgs, volume
 
     let clone_flags;
 
-    // 初始化容器工作目录
-    let container_work_path = Path::new(WORKSPACE).join("containers").join(&_container_id);
-    let container_upper_path = container_work_path.join("upper");
-    let container_merged_path = container_work_path.join("merged");
-    create_dir(&container_work_path, true)?;
-    create_dir(&container_upper_path, true)?;
-    create_dir(&container_merged_path, true)?;
-
     let rocker_user_info =  User::from_name(USER_NAME)?.ok_or(RockerError::OtherError(format!("没找到 用户: {USER_NAME}")))?;
     let rocker_uid = rocker_user_info.uid;
     let rocker_gid = rocker_user_info.gid;
@@ -368,7 +390,7 @@ fn run_container(_container_id: &String, cmd: &String, args: &RockerArgs, volume
         let _cb = move || {
             let container_info = get_container_info(_container_id).unwrap();
             init_exec_ns(container_info.pid).unwrap();
-            init_container_env().unwrap();
+            init_container_env(None).unwrap();
             init_container_user(rocker_uid, rocker_gid).unwrap();
 
             let cmd_vec = parse_cmd(cmd);
@@ -384,16 +406,24 @@ fn run_container(_container_id: &String, cmd: &String, args: &RockerArgs, volume
         clone_flags = CloneFlags::empty();
         Box::new(_cb) as CloneCb
     } else {
+        // 初始化容器工作目录
+        let container_work_path = Path::new(WORKSPACE).join("containers").join(&_container_id);
+        let container_upper_path = container_work_path.join("upper");
+        let container_merged_path = container_work_path.join("merged");
+        create_dir(&container_work_path, true)?;
+        create_dir(&container_upper_path, true)?;
+        create_dir(&container_merged_path, true)?;
+
         let _cb = move || {
             init_container_lock(&container_work_path).unwrap();
             init_container_overlay(volume_path, &container_upper_path, &container_merged_path).unwrap();
             if let Some(custom_volume) = &args.volume {
                 init_container_custom_volume(&container_merged_path, custom_volume).unwrap();
             }
+            init_container_env(args.env.as_ref()).unwrap();
             init_container_pivot(&container_merged_path).unwrap();
             init_container_mount().unwrap();
             init_container_log(args.log).unwrap();
-            init_container_env().unwrap();
             init_container_user(rocker_uid, rocker_gid).unwrap();
             
             let cmd_vec = parse_cmd(cmd);

test.env

@@ -0,0 +1,2 @@
+a=2
+b=3