Merge pull request #85 from tsturzl/main

Clean up use of unsafe
youki-dev · Jun 12, 2021 · a6a8435 · a6a8435
2 parents 42d877a + c34c14b
commit a6a8435
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 82 deletions.
diff --git a/src/process/fork.rs b/src/process/fork.rs
@@ -34,51 +34,49 @@ pub fn fork_first<P: AsRef<Path>>(
     // create a new child process structure with sending end of parent process
     let child = child::ChildProcess::new(sender_for_parent)?;
 
-    unsafe {
-        // fork the process
-        match unistd::fork()? {
-            // in the child process
-            unistd::ForkResult::Child => {
-                // if Out-of-memory score adjustment is set in specification.
-                // set the score value for the current process
-                // check https://dev.to/rrampage/surviving-the-linux-oom-killer-2ki9 for some more information
-                if let Some(ref r) = linux.resources {
-                    if let Some(adj) = r.oom_score_adj {
-                        let mut f = fs::File::create("/proc/self/oom_score_adj")?;
-                        f.write_all(adj.to_string().as_bytes())?;
-                    }
-                }
-
-                // if new user is specified in specification, this will be true
-                // and new namespace will be created, check https://man7.org/linux/man-pages/man7/user_namespaces.7.html
-                // for more information
-                if is_userns {
-                    sched::unshare(sched::CloneFlags::CLONE_NEWUSER)?;
+    // fork the process
+    match unsafe { unistd::fork()? } {
+        // in the child process
+        unistd::ForkResult::Child => {
+            // if Out-of-memory score adjustment is set in specification.
+            // set the score value for the current process
+            // check https://dev.to/rrampage/surviving-the-linux-oom-killer-2ki9 for some more information
+            if let Some(ref r) = linux.resources {
+                if let Some(adj) = r.oom_score_adj {
+                    let mut f = fs::File::create("/proc/self/oom_score_adj")?;
+                    f.write_all(adj.to_string().as_bytes())?;
                 }
+            }
 
-                ccond.notify()?;
-                Ok(Process::Child(child))
+            // if new user is specified in specification, this will be true
+            // and new namespace will be created, check https://man7.org/linux/man-pages/man7/user_namespaces.7.html
+            // for more information
+            if is_userns {
+                sched::unshare(sched::CloneFlags::CLONE_NEWUSER)?;
             }
-            // in the parent process
-            unistd::ForkResult::Parent { child } => {
-                ccond.wait()?;
 
-                // wait for child to fork init process and report back its pid
-                let init_pid = parent.wait_for_child_ready()?;
-                log::debug!("init pid is {:?}", init_pid);
-                cmanager.apply(&linux.resources.as_ref().unwrap(), Pid::from_raw(init_pid))?;
+            ccond.notify()?;
+            Ok(Process::Child(child))
+        }
+        // in the parent process
+        unistd::ForkResult::Parent { child } => {
+            ccond.wait()?;
 
-                // update status and pid of the container process
-                container
-                    .update_status(ContainerStatus::Created)?
-                    .set_pid(init_pid)
-                    .save()?;
-                // if file to write the pid to is specified, write pid of the child
-                if let Some(pid_file) = pid_file {
-                    fs::write(&pid_file, format!("{}", child))?;
-                }
-                Ok(Process::Parent(parent))
+            // wait for child to fork init process and report back its pid
+            let init_pid = parent.wait_for_child_ready()?;
+            log::debug!("init pid is {:?}", init_pid);
+            cmanager.apply(&linux.resources.as_ref().unwrap(), Pid::from_raw(init_pid))?;
+
+            // update status and pid of the container process
+            container
+                .update_status(ContainerStatus::Created)?
+                .set_pid(init_pid)
+                .save()?;
+            // if file to write the pid to is specified, write pid of the child
+            if let Some(pid_file) = pid_file {
+                fs::write(&pid_file, format!("{}", child))?;
             }
+            Ok(Process::Parent(parent))
         }
     }
 }
@@ -87,33 +85,31 @@ pub fn fork_first<P: AsRef<Path>>(
 pub fn fork_init(mut child_process: ChildProcess) -> Result<Process> {
     // setup sockets for init process
     let sender_for_child = child_process.setup_pipe()?;
-    unsafe {
-        // for the process into current process (C1) (which is child of first_fork) and init process
-        match unistd::fork()? {
-            // if it is child process, create new InitProcess structure and return
-            unistd::ForkResult::Child => Ok(Process::Init(InitProcess::new(sender_for_child))),
-            // in the forking process C1
-            unistd::ForkResult::Parent { child } => {
-                // wait for init process to be ready
-                child_process.wait_for_init_ready()?;
-                // notify the parent process (original youki process) that init process is forked and ready
-                child_process.notify_parent(child)?;
+    // for the process into current process (C1) (which is child of first_fork) and init process
+    match unsafe { unistd::fork()? } {
+        // if it is child process, create new InitProcess structure and return
+        unistd::ForkResult::Child => Ok(Process::Init(InitProcess::new(sender_for_child))),
+        // in the forking process C1
+        unistd::ForkResult::Parent { child } => {
+            // wait for init process to be ready
+            child_process.wait_for_init_ready()?;
+            // notify the parent process (original youki process) that init process is forked and ready
+            child_process.notify_parent(child)?;
 
-                // wait for the init process, which is container process, to change state
-                // check https://man7.org/linux/man-pages/man3/wait.3p.html for more information
-                match waitpid(child, None)? {
-                    // if normally exited
-                    WaitStatus::Exited(pid, status) => {
-                        log::debug!("exited pid: {:?}, status: {:?}", pid, status);
-                        exit(status);
-                    }
-                    // if terminated by a signal
-                    WaitStatus::Signaled(pid, status, _) => {
-                        log::debug!("signaled pid: {:?}, status: {:?}", pid, status);
-                        exit(0);
-                    }
-                    _ => bail!("abnormal exited!"),
+            // wait for the init process, which is container process, to change state
+            // check https://man7.org/linux/man-pages/man3/wait.3p.html for more information
+            match waitpid(child, None)? {
+                // if normally exited
+                WaitStatus::Exited(pid, status) => {
+                    log::debug!("exited pid: {:?}, status: {:?}", pid, status);
+                    exit(status);
+                }
+                // if terminated by a signal
+                WaitStatus::Signaled(pid, status, _) => {
+                    log::debug!("signaled pid: {:?}, status: {:?}", pid, status);
+                    exit(0);
                 }
+                _ => bail!("abnormal exited!"),
             }
         }
     }

diff --git a/src/utils.rs b/src/utils.rs
@@ -1,12 +1,13 @@
 //! Utility functionality
 
+use std::env;
 use std::ffi::CString;
 use std::fs;
 use std::path::{Path, PathBuf};
 use std::time::Duration;
 
 use anyhow::{bail, Result};
-use nix::{env::clearenv, errno::Errno, unistd};
+use nix::unistd;
 
 pub trait PathBufExt {
     fn as_in_container(&self) -> Result<PathBuf>;
@@ -40,29 +41,25 @@ pub fn do_exec(path: impl AsRef<Path>, args: &[String], envs: &[String]) -> Resu
         .iter()
         .map(|s| CString::new(s.to_string()).unwrap_or_default())
         .collect();
-    let envs: Vec<CString> = envs
-        .iter()
-        .map(|s| CString::new(s.to_string()).unwrap_or_default())
-        .collect();
 
-    unsafe {
-        clearenv()?;
-    }
-    for e in envs {
-        putenv(&e)?
-    }
+    // clear env vars
+    env::vars().for_each(|(key, _value)| std::env::remove_var(key));
+    // set env vars
+    envs.iter().for_each(|e| {
+        let mut split = e.split("=");
+        match split.next() {
+            Some(key) => {
+                let value: String = split.collect::<Vec<&str>>().join("=");
+                env::set_var(key, value)
+            }
+            None => {}
+        };
+    });
 
     unistd::execvp(&p, &a)?;
     Ok(())
 }
 
-#[inline]
-fn putenv(string: &CString) -> nix::Result<()> {
-    let ptr = string.clone().into_raw();
-    let res = unsafe { libc::putenv(ptr as *mut libc::c_char) };
-    Errno::result(res).map(drop)
-}
-
 // TODO implement
 pub fn set_name(_name: &str) -> Result<()> {
     Ok(())