From cd004aa8615e4eb83937993ed05abd559c40c49c Mon Sep 17 00:00:00 2001 From: LMJW Date: Fri, 30 Jul 2021 09:33:21 +1000 Subject: [PATCH] Add getrlimit and setrlimit This work is a continutation on previou contribution by @kpcyrd and @j1ah0ng. --- CHANGELOG.md | 2 + src/sys/mod.rs | 3 + src/sys/resource.rs | 195 ++++++++++++++++++++++++++++++++++++++++++ test/test.rs | 1 + test/test_resource.rs | 23 +++++ 5 files changed, 224 insertions(+) create mode 100644 src/sys/resource.rs create mode 100644 test/test_resource.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a7dcfafb9..fdf796ce3a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,8 @@ This project adheres to [Semantic Versioning](https://semver.org/). (#[1472](https://github.com/nix-rust/nix/pull/1472)) - Added `mknodat`. (#[1473](https://github.com/nix-rust/nix/pull/1473)) +- Added `setrlimit` and `getrlimit`. + (#[1302](https://github.com/nix-rust/nix/pull/1302)) ### Changed diff --git a/src/sys/mod.rs b/src/sys/mod.rs index b43587b8e1..cffefdc3af 100644 --- a/src/sys/mod.rs +++ b/src/sys/mod.rs @@ -59,6 +59,9 @@ pub mod quota; #[cfg(any(target_os = "linux"))] pub mod reboot; +#[cfg(not(any(target_os = "redox", target_os = "fuchsia", target_os = "illumos")))] +pub mod resource; + #[cfg(not(target_os = "redox"))] pub mod select; diff --git a/src/sys/resource.rs b/src/sys/resource.rs new file mode 100644 index 0000000000..8f71375b25 --- /dev/null +++ b/src/sys/resource.rs @@ -0,0 +1,195 @@ +//! Configure the process resource limits. +use cfg_if::cfg_if; + +use crate::errno::Errno; +use crate::Result; +pub use libc::rlim_t; +use std::mem; + +cfg_if! { + if #[cfg(all(target_os = "linux", target_env = "gnu"))]{ + use libc::{__rlimit_resource_t, rlimit, RLIM_INFINITY}; + }else if #[cfg(any( + target_os = "freebsd", + target_os = "openbsd", + target_os = "netbsd", + target_os = "macos", + target_os = "ios", + target_os = "android", + target_os = "dragonfly", + all(target_os = "linux", not(target_env = "gnu")) + ))]{ + use libc::{c_int, rlimit, RLIM_INFINITY}; + } +} + +libc_enum! { + /// The Resource enum is platform dependent. Check different platform + /// manuals for more details. Some platform links has been provided for + /// earier reference (non-exhaustive). + /// + /// * [Linux](https://man7.org/linux/man-pages/man2/getrlimit.2.html) + /// * [FreeBSD](https://www.freebsd.org/cgi/man.cgi?query=setrlimit) + + // linux-gnu uses u_int as resource enum, which is implemented in libc as + // well. + // + // https://gcc.gnu.org/legacy-ml/gcc/2015-08/msg00441.html + // https://github.com/rust-lang/libc/blob/master/src/unix/linux_like/linux/gnu/mod.rs + #[cfg_attr(all(target_os = "linux", target_env = "gnu"), repr(u32))] + #[cfg_attr(any( + target_os = "freebsd", + target_os = "openbsd", + target_os = "netbsd", + target_os = "macos", + target_os = "ios", + target_os = "android", + target_os = "dragonfly", + all(target_os = "linux", not(target_env = "gnu")) + ), repr(i32))] + #[non_exhaustive] + pub enum Resource { + #[cfg(not(any(target_os = "netbsd", target_os = "freebsd")))] + RLIMIT_AS, + RLIMIT_CORE, + RLIMIT_CPU, + RLIMIT_DATA, + RLIMIT_FSIZE, + RLIMIT_NOFILE, + RLIMIT_STACK, + + #[cfg(target_os = "freebsd")] + RLIMIT_KQUEUES, + + #[cfg(any(target_os = "android", target_os = "linux"))] + RLIMIT_LOCKS, + + #[cfg(any(target_os = "android", target_os = "freebsd", target_os = "openbsd", target_os = "linux"))] + RLIMIT_MEMLOCK, + + #[cfg(any(target_os = "android", target_os = "linux"))] + RLIMIT_MSGQUEUE, + + #[cfg(any(target_os = "android", target_os = "linux"))] + RLIMIT_NICE, + + #[cfg(any(target_os = "android", target_os = "freebsd", target_os = "openbsd", target_os = "linux"))] + RLIMIT_NPROC, + + #[cfg(target_os = "freebsd")] + RLIMIT_NPTS, + + #[cfg(any(target_os = "android", target_os = "freebsd", target_os = "openbsd", target_os = "linux"))] + RLIMIT_RSS, + + #[cfg(any(target_os = "android", target_os = "linux"))] + RLIMIT_RTPRIO, + + #[cfg(any(target_os = "linux"))] + RLIMIT_RTTIME, + + #[cfg(any(target_os = "android", target_os = "linux"))] + RLIMIT_SIGPENDING, + + #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))] + RLIMIT_SBSIZE, + + #[cfg(target_os = "freebsd")] + RLIMIT_SWAP, + + #[cfg(target_os = "freebsd")] + RLIMIT_VMEM, + } +} + +/// Get the current processes resource limits +/// +/// A value of `None` indicates the value equals to `RLIM_INFINITY` which means +/// there is no limit. +/// +/// # Parameters +/// +/// * `resource`: The [`Resource`] that we want to get the limits of. +/// +/// # Examples +/// +/// ``` +/// # use nix::sys::resource::{getrlimit, Resource}; +/// +/// let (soft_limit, hard_limit) = getrlimit(Resource::RLIMIT_NOFILE).unwrap(); +/// println!("current soft_limit: {:?}", soft_limit); +/// println!("current hard_limit: {:?}", hard_limit); +/// ``` +/// +/// # References +/// +/// [getrlimit(2)](https://pubs.opengroup.org/onlinepubs/9699919799/functions/getrlimit.html#tag_16_215) +/// +/// [`Resource`]: enum.Resource.html +pub fn getrlimit(resource: Resource) -> Result<(Option, Option)> { + let mut old_rlim = mem::MaybeUninit::::uninit(); + + cfg_if! { + if #[cfg(all(target_os = "linux", target_env = "gnu"))]{ + let res = unsafe { libc::getrlimit(resource as __rlimit_resource_t, old_rlim.as_mut_ptr()) }; + }else{ + let res = unsafe { libc::getrlimit(resource as c_int, old_rlim.as_mut_ptr()) }; + } + } + + Errno::result(res).map(|_| { + let rlimit { rlim_cur, rlim_max } = unsafe { old_rlim.assume_init() }; + (Some(rlim_cur), Some(rlim_max)) + }) +} + +/// Set the current processes resource limits +/// +/// # Parameters +/// +/// * `resource`: The [`Resource`] that we want to set the limits of. +/// * `soft_limit`: The value that the kernel enforces for the corresponding +/// resource. Note: `None` input will be replaced by constant `RLIM_INFINITY`. +/// * `hard_limit`: The ceiling for the soft limit. Must be lower or equal to +/// the current hard limit for non-root users. Note: `None` input will be +/// replaced by constant `RLIM_INFINITY`. +/// +/// > Note: for some os (linux_gnu), setting hard_limit to `RLIM_INFINITY` can +/// > results `EPERM` Error. So you will need to set the number explicitly. +/// +/// # Examples +/// +/// ``` +/// # use nix::sys::resource::{setrlimit, Resource}; +/// +/// let soft_limit = Some(1024); +/// let hard_limit = Some(1048576); +/// setrlimit(Resource::RLIMIT_NOFILE, soft_limit, hard_limit).unwrap(); +/// ``` +/// +/// # References +/// +/// [setrlimit(2)](https://pubs.opengroup.org/onlinepubs/9699919799/functions/getrlimit.html#tag_16_215) +/// +/// [`Resource`]: enum.Resource.html +/// +/// Note: `setrlimit` provides a safe wrapper to libc's `setrlimit`. +pub fn setrlimit( + resource: Resource, + soft_limit: Option, + hard_limit: Option, +) -> Result<()> { + let new_rlim = rlimit { + rlim_cur: soft_limit.unwrap_or(RLIM_INFINITY), + rlim_max: hard_limit.unwrap_or(RLIM_INFINITY), + }; + cfg_if! { + if #[cfg(all(target_os = "linux", target_env = "gnu"))]{ + let res = unsafe { libc::setrlimit(resource as __rlimit_resource_t, &new_rlim as *const rlimit) }; + }else{ + let res = unsafe { libc::setrlimit(resource as c_int, &new_rlim as *const rlimit) }; + } + } + + Errno::result(res).map(drop) +} diff --git a/test/test.rs b/test/test.rs index 94f8e220ea..b882d17485 100644 --- a/test/test.rs +++ b/test/test.rs @@ -24,6 +24,7 @@ mod test_mq; #[cfg(not(target_os = "redox"))] mod test_net; mod test_nix_path; +mod test_resource; mod test_poll; #[cfg(not(any(target_os = "redox", target_os = "fuchsia")))] mod test_pty; diff --git a/test/test_resource.rs b/test/test_resource.rs new file mode 100644 index 0000000000..5969750091 --- /dev/null +++ b/test/test_resource.rs @@ -0,0 +1,23 @@ +#[cfg(not(any(target_os = "redox", target_os = "fuchsia", target_os = "illumos")))] +use nix::sys::resource::{getrlimit, setrlimit, Resource}; + +/// Tests the RLIMIT_NOFILE functionality of getrlimit(), where the resource RLIMIT_NOFILE refers +/// to the maximum file descriptor number that can be opened by the process (aka the maximum number +/// of file descriptors that the process can open, since Linux 4.5). +/// +/// We first fetch the existing file descriptor maximum values using getrlimit(), then edit the +/// soft limit to make sure it has a new and distinct value to the hard limit. We then setrlimit() +/// to put the new soft limit in effect, and then getrlimit() once more to ensure the limits have +/// been updated. +#[test] +#[cfg(not(any(target_os = "redox", target_os = "fuchsia", target_os = "illumos")))] +pub fn test_resource_limits_nofile() { + let (soft_limit, hard_limit) = getrlimit(Resource::RLIMIT_NOFILE).unwrap(); + + let soft_limit = Some(soft_limit.map_or(1024, |v| v - 1)); + assert_ne!(soft_limit, hard_limit); + setrlimit(Resource::RLIMIT_NOFILE, soft_limit, hard_limit).unwrap(); + + let (new_soft_limit, _) = getrlimit(Resource::RLIMIT_NOFILE).unwrap(); + assert_eq!(new_soft_limit, soft_limit); +}