src/fork.rs

   1 //! Fork helper.
   2 //!
   3 //! Note that forking in rust can be dangerous. A fork must consider all mutexes to be in a broken
   4 //! state, and cannot rely on any of its reference life times, so we be careful what kind of data
   5 //! we continue to work with.
   6
   7 use std::io;
   8 use std::os::raw::c_int;
   9 use std::os::unix::io::{FromRawFd, IntoRawFd};
  10 use std::panic::UnwindSafe;
  11 use std::pin::Pin;
  12 use std::task::{Context, Poll};
  13
  14 use futures::future::poll_fn;
  15 use futures::io::AsyncRead;
  16
  17 use crate::syscall::SyscallStatus;
  18 use crate::tools::Fd;
  19 use crate::{libc_try, libc_wrap};
  20
  21 pub async fn forking_syscall<F>(func: F) -> io::Result<SyscallStatus>
  22 where
  23     F: FnOnce() -> io::Result<SyscallStatus> + UnwindSafe,
  24 {
  25     let mut fork = Fork::new(func)?;
  26     let mut buf = [0u8; 10];
  27
  28     use futures::io::AsyncReadExt;
  29     fork.read_exact(&mut buf).await?;
  30     fork.wait()?;
  31
  32     Ok(SyscallStatus::Err(libc::ENOENT))
  33 }
  34
  35 pub struct Fork {
  36     pid: Option<libc::pid_t>,
  37     // FIXME: abuse! tokio-fs is not updated to futures@0.3 yet, but a TcpStream does the same
  38     // thing as a file when it's already open anyway...
  39     out: crate::tools::GenericStream,
  40 }
  41
  42 impl Drop for Fork {
  43     fn drop(&mut self) {
  44         if self.pid.is_some() {
  45             let _ = self.wait();
  46         }
  47     }
  48 }
  49
  50 impl Fork {
  51     pub fn new<F>(func: F) -> io::Result<Self>
  52     where
  53         F: FnOnce() -> io::Result<SyscallStatus> + UnwindSafe,
  54     {
  55         let mut pipe: [c_int; 2] = [0, 0];
  56         libc_try!(unsafe { libc::pipe2(pipe.as_mut_ptr(), libc::O_CLOEXEC | libc::O_NONBLOCK) });
  57         let (pipe_r, pipe_w) = (Fd(pipe[0]), Fd(pipe[1]));
  58
  59         let pipe_r = match crate::tools::GenericStream::from_fd(pipe_r) {
  60             Ok(o) => o,
  61             Err(err) => return Err(io::Error::new(io::ErrorKind::Other, err.to_string())),
  62         };
  63
  64         let pid = libc_try!(unsafe { libc::fork() });
  65         if pid == 0 {
  66             std::mem::drop(pipe_r);
  67             let mut pipe_w = unsafe { std::fs::File::from_raw_fd(pipe_w.into_raw_fd()) };
  68
  69             let _ = std::panic::catch_unwind(move || {
  70                 let mut buf = [0u8; 10];
  71
  72                 match func() {
  73                     Ok(SyscallStatus::Ok(value)) => unsafe {
  74                         std::ptr::write(buf.as_mut_ptr().add(1) as *mut i64, value);
  75                     },
  76                     Ok(SyscallStatus::Err(value)) => {
  77                         buf[0] = 1;
  78                         unsafe {
  79                             std::ptr::write(buf.as_mut_ptr().add(1) as *mut i32, value);
  80                         }
  81                     }
  82                     Err(err) => match err.raw_os_error() {
  83                         Some(err) => {
  84                             buf[0] = 2;
  85                             unsafe {
  86                                 std::ptr::write(buf.as_mut_ptr().add(1) as *mut i32, err);
  87                             }
  88                         }
  89                         None => {
  90                             buf[0] = 3;
  91                         }
  92                     },
  93                 }
  94
  95                 use std::io::Write;
  96                 match pipe_w.write_all(&buf) {
  97                     Ok(()) => unsafe { libc::_exit(0) },
  98                     Err(_) => unsafe { libc::_exit(1) },
  99                 }
 100             });
 101             unsafe {
 102                 libc::_exit(-1);
 103             }
 104         }
 105
 106         Ok(Self {
 107             pid: Some(pid),
 108             out: pipe_r,
 109         })
 110     }
 111
 112     pub fn wait(&mut self) -> io::Result<()> {
 113         let my_pid = self.pid.take().unwrap();
 114
 115         loop {
 116             let mut status: c_int = -1;
 117             match libc_wrap!(unsafe { libc::waitpid(my_pid, &mut status, 0) }) {
 118                 Ok(pid) if pid == my_pid => break,
 119                 Ok(_other) => continue,
 120                 Err(ref err) if err.kind() == io::ErrorKind::Interrupted => continue,
 121                 Err(other) => return Err(other),
 122             }
 123         }
 124
 125         Ok(())
 126     }
 127
 128     pub async fn async_read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
 129         poll_fn(|cx| Pin::new(&mut *self).poll_read(cx, buf)).await
 130     }
 131 }
 132
 133 // default impl will work
 134 impl AsyncRead for Fork {
 135     fn poll_read(
 136         self: Pin<&mut Self>,
 137         cx: &mut Context,
 138         buf: &mut [u8],
 139     ) -> Poll<io::Result<usize>> {
 140         unsafe { self.map_unchecked_mut(|this| &mut this.out) }.poll_read(cx, buf)
 141     }
 142
 143     unsafe fn initializer(&self) -> futures::io::Initializer {
 144         self.out.initializer()
 145     }
 146
 147     fn poll_read_vectored(
 148         self: Pin<&mut Self>,
 149         cx: &mut Context,
 150         bufs: &mut [futures::io::IoSliceMut],
 151     ) -> Poll<io::Result<usize>> {
 152         unsafe { self.map_unchecked_mut(|this| &mut this.out) }.poll_read_vectored(cx, bufs)
 153     }
 154 }