Merge #1302

1302: Add getrlimit and setrlimit r=asomers a=LMJW

This PR is based on the previous PR #1190, which has not been updated for a very long time.

I have fixed the segfault and compilation error of the original PR and rebased the changes to the latest master branch.

Please let me know if anything need to been changed.

Co-authored-by: LMJW <heysuperming@gmail.com>

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
 

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;
 

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<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)
+}

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;

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);
+}