use rustpython_vm::{AsObject, TryFromBorrowedObject};

use super::*;

use crate::vm::{PyResult, VirtualMachine, convert::IntoPyException};

#[pyfunction]

fn fork_exec(args: ForkExecArgs<'_>, vm: &VirtualMachine) -> PyResult<libc::pid_t> {

// Check for interpreter shutdown when preexec_fn is used

if args.preexec_fn.is_some()

&& vm

.state

.finalizing

.load(core::sync::atomic::Ordering::Acquire)

{

return Err(vm.new_python_finalization_error(

"preexec_fn not supported at interpreter shutdown".to_owned(),

));

}

let extra_groups = args

.groups_list

.as_ref()

.map(|l| Vec::<RawGid>::try_from_borrowed_object(vm, l.as_object()))

.map(|res| res.map(|groups| groups.into_iter().map(|gid| gid.0).collect::<Vec<_>>()))

.transpose()?;

let argv = args.args.iter().collect::<CharPtrVec<'_>>();

let envp = args.env_list.as_ref().map(CharPtrVec::from_iter);

let procargs = ProcArgs {

argv: &argv,

envp: envp.as_deref(),

extra_groups: extra_groups.as_deref(),

};

match host_posix::fork().map_err(|err| err.into_pyexception(vm))? {

0 => exec(&args, procargs, vm),

child => Ok(child),

}

}

($($field:ident: $t:ty),*$(,)?) => {

#[derive(FromArgs)]

struct ForkExecArgs<'fd> {

$(#[pyarg(positional)] $field: $t,)*

}

};

fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

let s = OsPath::try_from_object(vm, obj)?.into_cstring(vm)?;

Ok(Self { s })

}

fn as_ref(&self) -> &CStr {

&self.s

}

fn from_iter<I: IntoIterator<Item = &'a T>>(iter: I) -> Self {

let vec = iter

.into_iter()

.map(|x| x.as_ref().as_ptr())

.chain(core::iter::once(core::ptr::null()))

.collect();

Self {

vec,

marker: PhantomData,

}

}

type Target = CharPtrSlice<'a>;

fn deref(&self) -> &Self::Target {

unsafe {

&*(self.vec.as_slice() as *const [*const libc::c_char] as *const CharPtrSlice<'a>)

}

}

const fn as_ptr(&self) -> *const *const libc::c_char {

self.slice.as_ptr()

}

fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

match MaybeFd::try_from_object(vm, obj)? {

MaybeFd::Valid(fd) => Ok(fd),

MaybeFd::Invalid => Err(vm.new_value_error("invalid fd")),

}

}

fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {

host_posix::write_fd(self.as_fd(), buf)

}

fn flush(&mut self) -> std::io::Result<()> {

Ok(())

}

fn as_raw_fd(&self) -> RawFd {

self.0.as_raw_fd()

}

fn into_raw_fd(self) -> RawFd {

self.0.as_raw_fd()

}

fn as_fd(&self) -> BorrowedFd<'_> {

self.0.as_fd()

}

fn from(fd: OwnedFd) -> Self {

Self(unsafe { BorrowedFd::borrow_raw(fd.into_raw_fd()) })

}

fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

let fd = i32::try_from_object(vm, obj)?;

Ok(if fd == -1 {

Self::Invalid

} else {

Self::Valid(Fd(unsafe { BorrowedFd::borrow_raw(fd) }))

})

}

fn as_raw_fd(&self) -> RawFd {

match self {

Self::Valid(fd) => fd.as_raw_fd(),

Self::Invalid => -1,

}

}

use core::cmp::Ordering;

let i = obj

.try_to_ref::<crate::vm::builtins::PyInt>(vm)

.map_err(|_| {

vm.new_type_error(format!(

"an integer is required (got type {})",

obj.class().name()

))

})?

.try_to_primitive::<i64>(vm)?;

match i.cmp(&-1) {

Ordering::Greater => Ok(i

.try_into()

.map_err(|_| vm.new_overflow_error(format!("{typ_name} is larger than maximum")))?),

Ordering::Less => Err(vm.new_overflow_error(format!("{typ_name} is less than minimum"))),

Ordering::Equal => Ok(-1i32 as u32),

}

fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

try_from_id(vm, obj, "uid").map(Self)

}

fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

try_from_id(vm, obj, "gid").map(Self)

}

argv: &'a CharPtrSlice<'a>,

envp: Option<&'a CharPtrSlice<'a>>,

extra_groups: Option<&'a [u32]>,

let mut ctx = ExecErrorContext::NoExec;

match exec_inner(args, procargs, &mut ctx, vm) {

Ok(x) => match x {},

Err(e) => {

let mut pipe = args.errpipe_write;

if matches!(ctx, ExecErrorContext::PreExec) {

// For preexec_fn errors, use SubprocessError format (errno=0)

let _ = write!(pipe, "SubprocessError:0:{}", ctx.as_msg());

} else {

// errno is written in hex format

let errno = e.raw_os_error().unwrap_or(0);

let _ = write!(pipe, "OSError:{errno:x}:{}", ctx.as_msg());

}

rustpython_host_env::os::exit(255)

}

}

NoExec,

ChDir,

PreExec,

Exec,

const fn as_msg(&self) -> &'static str {

match self {

Self::NoExec => "noexec",

Self::ChDir => "noexec:chdir",

Self::PreExec => "Exception occurred in preexec_fn.",

Self::Exec => "",

}

}

args: &ForkExecArgs<'_>,

procargs: ProcArgs<'_>,

ctx: &mut ExecErrorContext,

vm: &VirtualMachine,

) -> std::io::Result<Never> {

host_posix::setup_child_fds(

args.fds_to_keep.as_slice(),

args.errpipe_write.as_fd(),

args.p2cread.as_raw_fd(),

args.p2cwrite.as_raw_fd(),

args.c2pread.as_raw_fd(),

args.c2pwrite.as_raw_fd(),

args.errread.as_raw_fd(),

args.errwrite.as_raw_fd(),

args.errpipe_read.as_raw_fd(),

)?;

if let Some(ref cwd) = args.cwd {

host_posix::chdir(cwd.s.as_c_str()).inspect_err(|_| *ctx = ExecErrorContext::ChDir)?

}

host_posix::set_umask(args.child_umask);

if args.restore_signals {

host_posix::restore_signals();

}

host_posix::setsid_if_needed(args.call_setsid)?;

host_posix::setpgid_if_needed(args.pgid_to_set)?;

host_posix::setgroups_if_needed(procargs.extra_groups)?;

host_posix::setregid_if_needed(args.gid.map(|gid| gid.0))?;

host_posix::setreuid_if_needed(args.uid.map(|uid| uid.0))?;

// Call preexec_fn after all process setup but before closing FDs

if let Some(ref preexec_fn) = args.preexec_fn {

match preexec_fn.call((), vm) {

Ok(_) => {}

Err(_e) => {

// Cannot safely stringify exception after fork

*ctx = ExecErrorContext::PreExec;

return Err(std::io::Error::from_raw_os_error(0));

}

}

}

*ctx = ExecErrorContext::Exec;

if args.close_fds {

host_posix::close_fds(2, args.fds_to_keep.as_slice());

}

let err = host_posix::exec_replace(

args.exec_list.as_slice(),

procargs.argv.as_ptr(),

procargs.envp.map(CharPtrSlice::as_ptr),

);

Err(std::io::Error::from_raw_os_error(err as i32))