lolicron/src/sys/fork.rs

//! Forking utils

use super::*;
use errno::Errno;

use libc::{
    fork,
    setgid,
    setuid,
};
use std::{
    fmt,
};
use crate::util::PhantomDrop;
use cfg_if::cfg_if;

use super::pipe::{
    self,
    unix_pipe,
    pipe_read_value,
    pipe_write_value,
};

/// Forking error
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
#[repr(u32)]
pub enum Error {
    Fork = 1,
    SetUid = 2,
    SetGid = 3,
    
    Pipe = 4,
    PipeRead = 5,
    PipeWrite = 6,
    PipeBroken = 7,
    WaitPid = 8,
    
    Unknown = 0,
}

impl From<pipe::Error> for Error
{
    fn from(from: pipe::Error) -> Self
    {
	use pipe::Error::*;
	match from {
	    Create => Self::Pipe,
	    Read => Self::PipeRead,
	    Write => Self::PipeWrite,
	    Broken => Self::PipeBroken,
	    _ => Self::Unknown,
	}
    }
}


impl Error
{
    #[inline] fn from_u32(from: u32) -> Self
    {
	use Error::*;
	match from {
	    x if x == Fork as u32 => Fork,
	    x if x == SetUid as u32 => SetUid,
	    x if x == SetGid as u32 => SetGid,
	    x if x == Pipe as u32 => Pipe,
	    x if x == PipeRead as u32 => PipeRead,
	    x if x == PipeWrite as u32 => PipeWrite,
	    x if x == PipeBroken as u32 => PipeBroken,
	    x if x == WaitPid as u32 => WaitPid,
	    _ => Self::Unknown,
	}
    }
}

impl From<u32> for Error
{
    #[inline] fn from(from: u32) -> Self
    {
	Self::from_u32(from)
    }
}

impl From<Error> for u32
{
    #[inline] fn from(from: Error) -> Self
    {
	from as Self
    }
}

impl std::error::Error for Error{}
impl std::fmt::Display for Error
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
    {
	match self {
	    Self::Fork => write!(f, "fork() failed"),
	    Self::Pipe => write!(f, "pipe() failed"),
	    Self::PipeBroken => write!(f, "broken pipe with child"),
	    Self::PipeRead => write!(f, "pipe read failed"),
	    Self::PipeWrite => write!(f, "pipe write failed"),
	    Self::SetGid => write!(f, "child reported setgid() failed"),
	    Self::SetUid => write!(f, "child reported setuid() failed"),
	    Self::WaitPid => write!(f, "waitpid() failed unexpectedly"),
	    _ => write!(f, "child reported unknown error"),
	}
    }
}

/// Represents the detached child
#[derive(Debug)]
pub struct Child
{
    pid: i32,
}

#[derive(Debug)]
pub struct Parent {
    pid: i32,
}


/// Run a closure as a fork and then exit, optionally trying to set `uid` and `gid`.
///
/// # Notes
/// This fork does not set up the async runtime. Using Tokio from the child process will require manually setting up the runtime
pub async fn detach_closure<F: FnOnce(Parent)>(as_uid: Option<u32>, as_gid: Option<u32>, into: F) -> Result<Child, Errno<Error>> {

    let (rx, tx) = unix_pipe().map_inner(|x| Error::from(x))?;
    
    let child = unsafe{fork()};
    if child == 0 {
	// Is child
	
	let complete = move || {
	    unsafe {
		if let Ok(_) = pipe_write_value(tx, &!0u32)
		{
		    into(Parent{pid: libc::getppid()});
		}
	    }
	};

	let complete_err = move |err: Error| {
	    unsafe{let _ = pipe_write_value(tx, &u32::from(err));}
	};
	
	{
	    let _guard = PhantomDrop::new((), move |_| {
		unsafe {
		    libc::close(tx);
		}
	    });// Close pipe sender when we're done here

	    unsafe {
		loop {
		    if let Some(as_uid) = as_uid {
			// Set UID
			if setuid(as_uid) != 0 {
			    complete_err(Error::SetUid);
			    break;
			}
		    }
		    if let Some(as_gid) = as_gid {
			// Set GID
			if setgid(as_gid) != 0 {
			    complete_err(Error::SetGid);
			    break;
			}
		    }

		    complete();
		    break;
		}
	    }
	}
	std::process::exit(0);
    } else if child > 0 {
	// Fork succeeded
	let _guard = PhantomDrop::new((), move |_| {
	    unsafe {
		libc::close(rx);
	    }
	});

	
	let err: u32 = unsafe{
	    cfg_if! {
		if #[cfg(feature="threaded")] {
		    use tokio::{
			task,
		    };
		    let waiter = task::spawn_blocking(move || {
			pipe_read_value(rx).map_inner(|e| Error::from(e))
		    });

		    waiter.await.expect("Panic while waiting for child status")?
		} else {
		    unsafe {
			pipe_read_value(rx)?
		    }
		}
	    }
	};
	if err == !0u32 {   
	    Ok(Child{pid:child})
	} else {
	    Err(Error::from(err).into())
	}
    } else {
	let rval = Error::Fork.into();
	unsafe {
	    libc::close(tx);
	    libc::close(rx);
	}
	Err(rval)
    }
}

impl Child
{
    /// Call `waitpid` on this child. Returns the status code if possible, or `Error` if error.
    ///
    /// # Notes
    ///  - When `threaded` feature is disabled, this function will not block at all, instead returning error if there is not status available, this is to prevent deadlock.
    ///  - This function (at present) will return `Error` when the child process exits, too. This is because we don't have `errno` access yet, I'm working on it.
    pub async fn waitpid(&self) -> Result<i32, Errno<Error>>
    {
	cfg_if! {
	    if #[cfg(feature="threaded")] {
		let pid = self.pid;
		let waiter = tokio::task::spawn_blocking(move || {
		    let mut status: i32 = 0;
		    let (out, status) = (unsafe{libc::waitpid(pid, &mut status as *mut i32, 0)}, status);
		    if out != pid {
			Err(Errno::from(Error::WaitPid))
		    } else {
			Ok(status)
		    }
		});
		waiter.await.expect("Waiter panicked")   
	    } else {
		let mut status: i32 = 0;
		if unsafe{libc::waitpid(self.pid, &mut status as *mut i32, 1)} == pid { // We can't afford to block here
		    Ok(status)
		} else {
		    Err(Error::WaitPid.into())
		}
	    }
	}
    }

    /// Wait for the child process to end, ignoring any other signals.
    ///
    /// # Notes
    ///  - When `threaded` feature is disabled, this function will return immediately, this is an error-handling bug because we don't have `errno` access as of yet.
    ///  - This function will call `waitpid` untill it returns status `ECHILD` (child has exited). Other status values are ignored, if any.
    pub async fn wait(&self) -> Result<(), Errno<Error>>
    {
	loop {
	    match self.waitpid().await {
		Ok(v) => {
		    // keep going until we get `no child process'
		    debug!("Got status {} from child {}, ignoring. ", v, self.pid);
		},
		Err(e) if e.internal == Error::WaitPid && e.error() == 10 /* No child processes `ECHILD` */ => {
		    //debug!("Error: {}", e);
		    break Ok(());
		},
		Err(e) => {
		    break Err(e);
		},
	    }
	}
    }
}