use std::{
ops::RangeInclusive,
os::fd::{AsFd, AsRawFd, BorrowedFd, OwnedFd, RawFd},
};
use libc::{c_int, dup3, GRND_RANDOM};
use nix::{
errno::Errno,
fcntl::{OFlag, ResolveFlag},
sys::resource::{getrlimit, Resource},
unistd::{close, UnlinkatFlags},
NixPath,
};
use crate::{
cookie::safe_unlinkat,
fs::{is_active_fd, retry_on_eintr, safe_open},
log::{now, Tm},
path::{XPath, XPathBuf, PATH_MAX},
};
#[must_use = "hold the guard to keep cancellation disabled"]
pub struct CancelGuard(c_int);
const _PTHREAD_CANCEL_ENABLE: c_int = 0;
const PTHREAD_CANCEL_DISABLE: c_int = 1;
extern "C" {
fn pthread_setcancelstate(state: c_int, oldstate: *mut c_int) -> c_int;
}
impl CancelGuard {
pub fn acquire() -> Result<Self, Errno> {
let mut old: c_int = 0;
let err = unsafe { pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, &raw mut old) };
if err == 0 {
Ok(Self(old))
} else {
Err(Errno::from_raw(err))
}
}
}
impl Drop for CancelGuard {
fn drop(&mut self) {
unsafe {
pthread_setcancelstate(self.0, std::ptr::null_mut());
}
}
}
pub trait RandUint: Copy + Ord {
const ZERO: Self;
const ONE: Self;
const MAX: Self;
fn rand_from_os() -> Result<Self, Errno>;
fn checked_add(self, rhs: Self) -> Option<Self>;
fn checked_sub(self, rhs: Self) -> Option<Self>;
fn checked_mul(self, rhs: Self) -> Option<Self>;
fn div_euclid_opt(self, rhs: Self) -> Option<Self>;
}
macro_rules! impl_rand_uint {
($($t:ty),* $(,)?) => {$(
impl RandUint for $t {
const ZERO: Self = 0;
const ONE: Self = 1;
const MAX: Self = <$t>::MAX;
#[inline]
fn rand_from_os() -> Result<Self, Errno> {
let mut buf = [0u8; { std::mem::size_of::<$t>() }];
fillrandom(&mut buf)?;
Ok(<$t>::from_ne_bytes(buf))
}
#[inline] fn checked_add(self, rhs: Self) -> Option<Self> { self.checked_add(rhs) }
#[inline] fn checked_sub(self, rhs: Self) -> Option<Self> { self.checked_sub(rhs) }
#[inline] fn checked_mul(self, rhs: Self) -> Option<Self> { self.checked_mul(rhs) }
#[inline]
fn div_euclid_opt(self, rhs: Self) -> Option<Self> {
if rhs == 0 { None } else { Some(self.div_euclid(rhs)) }
}
}
)*};
}
impl_rand_uint!(u8, u16, u32, u64, u128, usize);
pub fn randint<T>(range: RangeInclusive<T>) -> Result<T, Errno>
where
T: RandUint,
{
let (lo, hi) = range.into_inner();
if lo >= hi {
return Err(Errno::EINVAL);
}
if lo == T::ZERO && hi == T::MAX {
return T::rand_from_os();
}
let span = hi
.checked_sub(lo)
.ok_or(Errno::EOVERFLOW)?
.checked_add(T::ONE)
.ok_or(Errno::EOVERFLOW)?;
let k = T::MAX.div_euclid_opt(span).ok_or(Errno::EOVERFLOW)?;
let accept_top = k.checked_mul(span).ok_or(Errno::EOVERFLOW)?;
loop {
let r = T::rand_from_os()?;
if r < accept_top {
let q = r.div_euclid_opt(span).ok_or(Errno::EOVERFLOW)?;
let qspan = q.checked_mul(span).ok_or(Errno::EOVERFLOW)?;
let off = r.checked_sub(qspan).ok_or(Errno::EOVERFLOW)?;
let v = lo.checked_add(off).ok_or(Errno::EOVERFLOW)?;
return Ok(v);
}
}
}
#[inline]
pub fn randport() -> Result<u16, Errno> {
randint(1025u16..=0xFFFF)
}
pub fn getrandom(size: usize) -> Result<Vec<u8>, Errno> {
if size == 0 {
return Err(Errno::EINVAL);
}
let mut buf = Vec::new();
if buf.try_reserve(size).is_err() {
return Err(Errno::ENOMEM);
}
buf.resize(size, 0);
fillrandom(&mut buf)?;
Ok(buf)
}
pub fn fillrandom(buf: &mut [u8]) -> Result<(), Errno> {
let siz = buf.len();
if siz == 0 {
return Err(Errno::EINVAL);
}
let guard = CancelGuard::acquire()?;
let mut n = 0;
while n < siz {
let ptr = &mut buf[n..];
let ptr = ptr.as_mut_ptr().cast();
let siz = siz.checked_sub(n).ok_or(Errno::EOVERFLOW)?;
n = n
.checked_add(
retry_on_eintr(|| {
Errno::result(unsafe { libc::getrandom(ptr, siz, GRND_RANDOM) })
})?
.try_into()
.or(Err(Errno::EINVAL))?,
)
.ok_or(Errno::EOVERFLOW)?;
}
drop(guard);
Ok(())
}
pub fn duprand(oldfd: RawFd, mut flags: OFlag) -> Result<RawFd, Errno> {
let range_start = 7u64;
let (range_end, _) = getrlimit(Resource::RLIMIT_NOFILE)?;
#[expect(clippy::unnecessary_cast)]
let range_end = range_end.saturating_sub(1) as u64;
let range_end = range_end.min(0x10000);
if range_end <= range_start {
return Err(Errno::EMFILE);
}
let range = range_start..=range_end;
let close_old = flags.contains(OFlag::O_EXCL);
flags.remove(OFlag::O_EXCL);
for _ in range.clone() {
#[expect(clippy::cast_possible_truncation)]
let fd_rand = randint(range.clone())? as RawFd;
let fd_rand = unsafe { BorrowedFd::borrow_raw(fd_rand) };
if is_active_fd(fd_rand) {
continue;
}
match retry_on_eintr(|| {
Errno::result(unsafe { dup3(oldfd, fd_rand.as_raw_fd(), flags.bits()) })
}) {
Ok(_) => {
if close_old {
let _ = close(oldfd);
}
return Ok(fd_rand.as_raw_fd());
}
Err(Errno::EMFILE) => return Err(Errno::EMFILE),
Err(_) => {}
}
}
Err(Errno::EBADF)
}
pub fn mkstempat<Fd: AsFd>(dirfd: Fd, prefix: &XPath) -> Result<OwnedFd, Errno> {
const MAX_TCOUNT: usize = 8;
const SUFFIX_LEN: usize = 128;
const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
let mut flags = OFlag::O_TMPFILE | OFlag::O_EXCL | OFlag::O_RDWR;
match safe_open(&dirfd, XPath::from_bytes(b"."), flags, ResolveFlag::empty()) {
Ok(fd) => return Ok(fd),
Err(Errno::EISDIR | Errno::ENOENT | Errno::EOPNOTSUPP) => {}
Err(errno) => return Err(errno),
}
flags.remove(OFlag::O_TMPFILE);
flags.insert(OFlag::O_CREAT);
if prefix.is_absolute() {
return Err(Errno::EINVAL);
} else if prefix.len().saturating_sub(SUFFIX_LEN) > PATH_MAX {
return Err(Errno::ENAMETOOLONG);
}
let mut attempts = 0;
let mut rng_data = [0u8; SUFFIX_LEN];
#[expect(clippy::arithmetic_side_effects)]
loop {
attempts += 1;
if attempts > MAX_TCOUNT {
return Err(Errno::EEXIST);
}
fillrandom(&mut rng_data)?;
let mut base = XPathBuf::with_capacity(prefix.len() + SUFFIX_LEN);
base.append_bytes(prefix.as_bytes());
for &b in &rng_data {
base.append_byte(CHARSET[(b as usize) % CHARSET.len()]);
}
match safe_open(&dirfd, &base, flags, ResolveFlag::empty()) {
Ok(fd) => {
safe_unlinkat(dirfd, &base, UnlinkatFlags::NoRemoveDir)?;
return Ok(fd);
}
Err(Errno::EEXIST) => {
continue;
}
Err(errno) => return Err(errno),
}
}
}
pub fn rand_version() -> Result<String, Errno> {
const VERMAGICS: &[&str] = &[
"SMP",
"SMP PREEMPT",
"SMP PREEMPT_DYNAMIC",
"SMP PREEMPT_RT",
];
const MONTHS: &[&str] = &[
"Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
];
const WKDAYS: &[&str] = &["Sun", "Mon", "Tue", "Wed", "Thu", "Fri", "Sat"];
const TS_WINDOW: u64 = 366 * 86_400;
let now = now();
let offset: u64 = randint(0..=TS_WINDOW)?;
let target: i64 = now
.saturating_sub(offset)
.try_into()
.or(Err(Errno::EOVERFLOW))?;
let tm = Tm::try_from(target)?;
let build_no = randint(1u8..=64)?;
#[expect(clippy::arithmetic_side_effects)]
let vermagic = VERMAGICS[randint(0usize..=(VERMAGICS.len() - 1))?];
#[expect(clippy::arithmetic_side_effects)]
#[expect(clippy::cast_sign_loss)]
let mon = MONTHS[(tm.month() - 1) as usize];
#[expect(clippy::cast_sign_loss)]
let wday = WKDAYS[tm.weekday() as usize];
let mday = tm.day();
let hh = tm.hour();
let mm = tm.minute();
let ss = tm.second();
let year = tm.year();
Ok(format!(
"#{build_no} {vermagic} {wday} {mon} {mday:>2} {hh:02}:{mm:02}:{ss:02} UTC {year}",
))
}
#[cfg(test)]
mod tests {
use std::fmt::Debug;
use super::*;
#[test]
fn test_fillrandom() {
assert_eq!(fillrandom(&mut []), Err(Errno::EINVAL));
assert_eq!(fillrandom(&mut [0u8; 257]), Ok(()));
}
#[test]
fn test_getrandom() {
assert_eq!(getrandom(0), Err(Errno::EINVAL));
let result = getrandom(257);
assert!(result.is_ok(), "result:{result:?}");
}
fn draw<T: RandUint + Debug>(lo: T, hi: T) -> T {
match randint::<T>(lo..=hi) {
Ok(v) => v,
Err(e) => panic!("randint failed for [{:?},{:?}] -> {:?}", lo, hi, e),
}
}
fn sample<T: RandUint + Debug>(lo: T, hi: T, n: usize) -> Vec<T> {
(0..n).map(|_| draw::<T>(lo, hi)).collect()
}
fn all_in_range<T: RandUint + Debug>(xs: &[T], lo: T, hi: T) -> bool {
xs.iter().all(|&v| v >= lo && v <= hi)
}
#[test]
fn test_randint_invalid_u8() {
assert!(matches!(randint::<u8>(200..=100), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_invalid_u16() {
assert!(matches!(randint::<u16>(5000..=4999), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_invalid_u32() {
assert!(matches!(randint::<u32>(42..=41), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_invalid_u64() {
assert!(matches!(randint::<u64>(999..=998), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_invalid_u128() {
assert!(matches!(randint::<u128>(500..=499), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_invalid_usize() {
assert!(matches!(randint::<usize>(100..=99), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_onepoint_u8() {
assert!(matches!(randint::<u8>(77..=77), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_onepoint_u16() {
assert!(matches!(randint::<u16>(31337..=31337), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_onepoint_u32() {
assert!(matches!(
randint::<u32>(1_000_000..=1_000_000),
Err(Errno::EINVAL)
));
}
#[test]
fn test_randint_onepoint_u64() {
assert!(matches!(
randint::<u64>(123456789..=123456789),
Err(Errno::EINVAL)
));
}
#[test]
fn test_randint_onepoint_u128() {
assert!(matches!(randint::<u128>(999..=999), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_onepoint_usize() {
assert!(matches!(randint::<usize>(4242..=4242), Err(Errno::EINVAL)));
}
#[test]
fn test_randint_fulldomain_u8_inbounds() {
let xs = sample::<u8>(u8::MIN, u8::MAX, 4096);
assert!(all_in_range(&xs, u8::MIN, u8::MAX));
}
#[test]
fn test_randint_fulldomain_u16_inbounds() {
let xs = sample::<u16>(u16::MIN, u16::MAX, 2048);
assert!(all_in_range(&xs, u16::MIN, u16::MAX));
}
#[test]
fn test_randint_fulldomain_u32_inbounds() {
let xs = sample::<u32>(u32::MIN, u32::MAX, 2048);
assert!(all_in_range(&xs, u32::MIN, u32::MAX));
}
#[test]
fn test_randint_fulldomain_u64_inbounds() {
let xs = sample::<u64>(u64::MIN, u64::MAX, 1024);
assert!(all_in_range(&xs, u64::MIN, u64::MAX));
}
#[test]
fn test_randint_fulldomain_u128_inbounds() {
let xs = sample::<u128>(u128::MIN, u128::MAX, 256);
assert!(all_in_range(&xs, u128::MIN, u128::MAX));
}
#[test]
fn test_randint_fulldomain_usize_inbounds() {
let xs = sample::<usize>(usize::MIN, usize::MAX, 1024);
assert!(all_in_range(&xs, usize::MIN, usize::MAX));
}
#[test]
fn test_randint_u8_nearmax_inbounds() {
let lo = u8::MAX.saturating_sub(15);
let xs = sample::<u8>(lo, u8::MAX, 2000);
assert!(all_in_range(&xs, lo, u8::MAX));
}
#[test]
fn test_randint_u16_nearmax_inbounds() {
let lo = u16::MAX.saturating_sub(1023);
let xs = sample::<u16>(lo, u16::MAX, 4000);
assert!(all_in_range(&xs, lo, u16::MAX));
}
#[test]
fn test_randint_u32_nearmax_inbounds() {
let lo = u32::MAX.saturating_sub(1000);
let xs = sample::<u32>(lo, u32::MAX, 3000);
assert!(all_in_range(&xs, lo, u32::MAX));
}
#[test]
fn test_randint_u64_nearmax_inbounds() {
let lo = u64::MAX.saturating_sub(1000);
let xs = sample::<u64>(lo, u64::MAX, 3000);
assert!(all_in_range(&xs, lo, u64::MAX));
}
#[test]
fn test_randint_u128_nearmax_inbounds() {
let lo = u128::MAX.saturating_sub(1000);
let xs = sample::<u128>(lo, u128::MAX, 2000);
assert!(all_in_range(&xs, lo, u128::MAX));
}
#[test]
fn test_randint_usize_nearmax_inbounds() {
let lo = usize::MAX.saturating_sub(1000);
let xs = sample::<usize>(lo, usize::MAX, 3000);
assert!(all_in_range(&xs, lo, usize::MAX));
}
}