From 57259b00ac94cf460fe9a95346c6e8ff39c6df3f Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Thu, 8 Apr 2021 10:51:08 -0500 Subject: [PATCH] Add the resource module --- vm/src/stdlib/mod.rs | 3 + vm/src/stdlib/os.rs | 2 +- vm/src/stdlib/resource.rs | 168 +++++++++++++++++++++++++++++++++ vm/src/stdlib/sysconfigdata.rs | 6 +- 4 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 vm/src/stdlib/resource.rs diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs index ffa90764f..b2c9ab208 100644 --- a/vm/src/stdlib/mod.rs +++ b/vm/src/stdlib/mod.rs @@ -56,6 +56,8 @@ mod multiprocessing; mod posixsubprocess; #[cfg(all(unix, not(any(target_os = "android", target_os = "redox"))))] mod pwd; +#[cfg(unix)] +mod resource; #[cfg(not(target_arch = "wasm32"))] mod select; #[cfg(not(target_arch = "wasm32"))] @@ -164,6 +166,7 @@ pub fn get_module_inits() -> HashMap "_posixsubprocess".to_owned(), Box::new(posixsubprocess::make_module), ); + modules.insert("resource".to_owned(), Box::new(resource::make_module)); } // Windows-only diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index 7f4bc313d..0e1b9397a 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -1270,7 +1270,7 @@ mod _os { Ok(()) } - #[cfg(unix)] + #[cfg(all(unix, not(any(target_os = "redox", target_os = "android"))))] #[pyfunction] fn getloadavg(vm: &VirtualMachine) -> PyResult<(f64, f64, f64)> { let mut loadavg = [0f64; 3]; diff --git a/vm/src/stdlib/resource.rs b/vm/src/stdlib/resource.rs new file mode 100644 index 000000000..76212dbc8 --- /dev/null +++ b/vm/src/stdlib/resource.rs @@ -0,0 +1,168 @@ +pub(crate) use resource::make_module; + +#[pymodule] +mod resource { + use super::super::os; + use crate::exceptions::IntoPyException; + use crate::pyobject::{IntoPyObject, PyObjectRef, PyResult, PyStructSequence, TryFromObject}; + use crate::VirtualMachine; + use std::{io, mem}; + + cfg_if::cfg_if! { + if #[cfg(target_os = "linux")] { + use libc::RLIMIT_NLIMITS as RLIM_NLIMITS; + } else if #[cfg(target_os = "android")] { + pub const RLIM_NLIMITS: i32 = 16; + } else { + use libc::RLIM_NLIMITS; + } + } + + // TODO: RLIMIT_OFILE, + #[pyattr] + use libc::{ + RLIMIT_AS, RLIMIT_CORE, RLIMIT_CPU, RLIMIT_DATA, RLIMIT_FSIZE, RLIMIT_MEMLOCK, + RLIMIT_NOFILE, RLIMIT_NPROC, RLIMIT_RSS, RLIMIT_STACK, RLIM_INFINITY, + }; + + #[cfg(any(target_os = "linux", target_os = "android", target_os = "emscripten"))] + #[pyattr] + use libc::{RLIMIT_MSGQUEUE, RLIMIT_NICE, RLIMIT_RTPRIO, RLIMIT_SIGPENDING}; + // TODO: I think this is supposed to be defined for all linux_like? + #[cfg(target_os = "linux")] + #[pyattr] + use libc::RLIMIT_RTTIME; + + #[cfg(any(target_os = "freebsd", target_os = "netbsd"))] + #[pyattr] + use libc::{RLIMIT_NPTS, RLIMIT_SBSIZE, RLIMIT_SWAP}; + + #[cfg(any(target_os = "freebsd", target_os = "solaris", target_os = "illumos"))] + #[pyattr] + use libc::{RLIMIT_NPTS, RLIMIT_SBSIZE, RLIMIT_SWAP, RLIMIT_VMEM}; + + #[pyattr] + #[pyclass(name = "struct_rusage")] + #[derive(PyStructSequence)] + struct Rusage { + ru_utime: f64, + ru_stime: f64, + ru_maxrss: libc::c_long, + ru_ixrss: libc::c_long, + ru_idrss: libc::c_long, + ru_isrss: libc::c_long, + ru_minflt: libc::c_long, + ru_majflt: libc::c_long, + ru_nswap: libc::c_long, + ru_inblock: libc::c_long, + ru_oublock: libc::c_long, + ru_msgsnd: libc::c_long, + ru_msgrcv: libc::c_long, + ru_nsignals: libc::c_long, + ru_nvcsw: libc::c_long, + ru_nivcsw: libc::c_long, + } + + #[pyimpl(with(PyStructSequence))] + impl Rusage {} + + impl From for Rusage { + fn from(rusage: libc::rusage) -> Self { + let tv = |tv: libc::timeval| tv.tv_sec as f64 + (tv.tv_usec as f64 / 1_000_000.0); + Rusage { + ru_utime: tv(rusage.ru_utime), + ru_stime: tv(rusage.ru_utime), + ru_maxrss: rusage.ru_maxrss, + ru_ixrss: rusage.ru_ixrss, + ru_idrss: rusage.ru_idrss, + ru_isrss: rusage.ru_isrss, + ru_minflt: rusage.ru_minflt, + ru_majflt: rusage.ru_majflt, + ru_nswap: rusage.ru_nswap, + ru_inblock: rusage.ru_inblock, + ru_oublock: rusage.ru_oublock, + ru_msgsnd: rusage.ru_msgsnd, + ru_msgrcv: rusage.ru_msgrcv, + ru_nsignals: rusage.ru_nsignals, + ru_nvcsw: rusage.ru_nvcsw, + ru_nivcsw: rusage.ru_nivcsw, + } + } + } + + #[pyfunction] + fn getrusage(who: i32, vm: &VirtualMachine) -> PyResult { + let res = unsafe { + let mut rusage = mem::MaybeUninit::::uninit(); + if libc::getrusage(who, rusage.as_mut_ptr()) == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(rusage.assume_init()) + } + }; + res.map(Rusage::from).map_err(|e| { + if e.kind() == io::ErrorKind::InvalidInput { + vm.new_value_error("invalid who parameter".to_owned()) + } else { + e.into_pyexception(vm) + } + }) + } + + struct Limits(libc::rlimit); + impl TryFromObject for Limits { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + let seq = vm.extract_elements::(&obj)?; + match *seq { + [cur, max] => Ok(Self(libc::rlimit { + rlim_cur: cur & RLIM_INFINITY, + rlim_max: max & RLIM_INFINITY, + })), + _ => Err(vm.new_value_error("expected a tuple of 2 integers".to_owned())), + } + } + } + impl IntoPyObject for Limits { + fn into_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { + (self.0.rlim_cur, self.0.rlim_max).into_pyobject(vm) + } + } + + #[pyfunction] + fn getrlimit(resource: i32, vm: &VirtualMachine) -> PyResult { + if resource < 0 || resource >= RLIM_NLIMITS as _ { + return Err(vm.new_value_error("invalid resource specified".to_owned())); + } + let rlimit = unsafe { + let mut rlimit = mem::MaybeUninit::::uninit(); + if libc::getrlimit(resource as _, rlimit.as_mut_ptr()) == -1 { + return Err(os::errno_err(vm)); + } + rlimit.assume_init() + }; + Ok(Limits(rlimit)) + } + + #[pyfunction] + fn setrlimit(resource: i32, limits: Limits, vm: &VirtualMachine) -> PyResult<()> { + if resource < 0 || resource >= RLIM_NLIMITS as _ { + return Err(vm.new_value_error("invalid resource specified".to_owned())); + } + let res = unsafe { + if libc::setrlimit(resource as _, &limits.0) == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } + }; + res.map_err(|e| match e.kind() { + io::ErrorKind::InvalidInput => { + vm.new_value_error("current limit exceeds maximum limit".to_owned()) + } + io::ErrorKind::PermissionDenied => { + vm.new_value_error("not allowed to raise maximum limit".to_owned()) + } + _ => e.into_pyexception(vm), + }) + } +} diff --git a/vm/src/stdlib/sysconfigdata.rs b/vm/src/stdlib/sysconfigdata.rs index 5b15e9b1c..d785bd4d8 100644 --- a/vm/src/stdlib/sysconfigdata.rs +++ b/vm/src/stdlib/sysconfigdata.rs @@ -1,4 +1,4 @@ -use crate::pyobject::{ItemProtocol, PyObjectRef}; +use crate::pyobject::{IntoPyObject, ItemProtocol, PyObjectRef}; use crate::VirtualMachine; use crate::sysmodule::MULTIARCH; @@ -7,13 +7,15 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let vars = vm.ctx.new_dict(); macro_rules! hashmap { ($($key:literal => $value:expr),*$(,)?) => {{ - $(vars.set_item($key, vm.ctx.new_str($value), vm).unwrap();)* + $(vars.set_item($key, $value.into_pyobject(vm), vm).unwrap();)* }}; } hashmap! { // fake shared module extension "EXT_SUFFIX" => format!(".rustpython-{}", MULTIARCH), "MULTIARCH" => MULTIARCH, + // enough for tests to stop expecting urandom() to fail after restricting file resources + "HAVE_GETRANDOM" => 1, } include!(concat!(env!("OUT_DIR"), "/env_vars.rs"));