Skip to content

Commit

Permalink
Add getrlimit and setrlimit
Browse files Browse the repository at this point in the history
This work is a continutation on previou contribution by @kpcyrd and @j1ah0ng.
  • Loading branch information
LMJW committed Jul 31, 2021
1 parent 8d36e9a commit 0fb6859
Show file tree
Hide file tree
Showing 4 changed files with 222 additions and 0 deletions.
3 changes: 3 additions & 0 deletions src/sys/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
195 changes: 195 additions & 0 deletions src/sys/resource.rs
Original file line number Diff line number Diff line change
@@ -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<rlim_t>, Option<rlim_t>)> {
let mut old_rlim = mem::MaybeUninit::<rlimit>::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<rlim_t>,
hard_limit: Option<rlim_t>,
) -> 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)
}
1 change: 1 addition & 0 deletions test/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
23 changes: 23 additions & 0 deletions test/test_resource.rs
Original file line number Diff line number Diff line change
@@ -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);
}

0 comments on commit 0fb6859

Please sign in to comment.