Skip to content

Commit

Permalink
Use MaybeUninit in DecInt (#1201)
Browse files Browse the repository at this point in the history
* Use MaybeUninit in DecInt

* Fix DecInt::new panicking on integers greater than 64 bits and optimize out any panics in DecInt::new

* Add tests so people messing with DecInt can catch UB with miri

* Use new itoa

---------

Co-authored-by: Alex Saveau <[email protected]>
  • Loading branch information
SUPERCILEX and Alex Saveau authored Nov 21, 2024
1 parent 41b26db commit 61f17d6
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 47 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ rust-version = "1.63"

[dependencies]
bitflags = { version = "2.4.0", default-features = false }
itoa = { version = "1.0.1", default-features = false, optional = true }
itoa = { version = "1.0.13", default-features = false, optional = true }

# Special dependencies used in rustc-dep-of-std mode.
core = { version = "1.0.0", optional = true, package = "rustc-std-workspace-core" }
Expand Down
66 changes: 35 additions & 31 deletions src/path/dec_int.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

use crate::backend::fd::{AsFd, AsRawFd};
use crate::ffi::CStr;
use core::fmt::Write;
use core::mem::{self, MaybeUninit};
use itoa::{Buffer, Integer};
#[cfg(all(feature = "std", unix))]
use std::os::unix::ffi::OsStrExt;
Expand Down Expand Up @@ -36,23 +36,36 @@ use {core::fmt, std::ffi::OsStr, std::path::Path};
/// ```
#[derive(Clone)]
pub struct DecInt {
// 20 `u8`s is enough to hold the decimal ASCII representation of any
// `u64`, and we add one for a NUL terminator for `as_c_str`.
buf: [u8; 20 + 1],
// Enough to hold an {u,i}64 and NUL terminator.
buf: [MaybeUninit<u8>; u64::MAX_STR_LEN + 1],
len: usize,
}
const _: () = assert!(u64::MAX_STR_LEN == i64::MAX_STR_LEN);

impl DecInt {
/// Construct a new path component from an integer.
#[inline]
pub fn new<Int: Integer>(i: Int) -> Self {
let mut me = DecIntWriter(Self {
buf: [0; 20 + 1],
len: 0,
pub fn new<Int: Integer + 'static>(i: Int) -> Self {
let mut buf = [MaybeUninit::uninit(); 21];

let mut str_buf = Buffer::new();
let str_buf = str_buf.format(i);
assert!(
str_buf.len() < buf.len(),
"{str_buf}{} unsupported.",
core::any::type_name::<Int>()
);

buf[..str_buf.len()].copy_from_slice(unsafe {
// SAFETY: you can always go from init to uninit
mem::transmute::<&[u8], &[MaybeUninit<u8>]>(str_buf.as_bytes())
});
let mut buf = Buffer::new();
me.write_str(buf.format(i)).unwrap();
me.0
buf[str_buf.len()] = MaybeUninit::new(0);

Self {
buf,
len: str_buf.len(),
}
}

/// Construct a new path component from a file descriptor.
Expand All @@ -72,44 +85,35 @@ impl DecInt {
/// Return the raw byte buffer as a `&CStr`.
#[inline]
pub fn as_c_str(&self) -> &CStr {
let bytes_with_nul = &self.buf[..=self.len];
let bytes_with_nul = self.as_bytes_with_nul();
debug_assert!(CStr::from_bytes_with_nul(bytes_with_nul).is_ok());

// SAFETY: `self.buf` holds a single decimal ASCII representation and
// at least one extra NUL byte.
unsafe { CStr::from_bytes_with_nul_unchecked(bytes_with_nul) }
}

/// Return the raw byte buffer.
/// Return the raw byte buffer including the NUL byte.
#[inline]
pub fn as_bytes(&self) -> &[u8] {
&self.buf[..self.len]
pub fn as_bytes_with_nul(&self) -> &[u8] {
let init = &self.buf[..=self.len];
// SAFETY: we're guaranteed to have initialized len+1 bytes.
unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(init) }
}
}

/// A wrapper around `DecInt` that implements `Write` without exposing this
/// implementation to `DecInt`'s public API.
struct DecIntWriter(DecInt);

impl core::fmt::Write for DecIntWriter {
/// Return the raw byte buffer.
#[inline]
fn write_str(&mut self, s: &str) -> core::fmt::Result {
match self.0.buf.get_mut(self.0.len..self.0.len + s.len()) {
Some(slice) => {
slice.copy_from_slice(s.as_bytes());
self.0.len += s.len();
Ok(())
}
None => Err(core::fmt::Error),
}
pub fn as_bytes(&self) -> &[u8] {
let bytes = self.as_bytes_with_nul();
&bytes[..bytes.len() - 1]
}
}

#[cfg(feature = "std")]
impl AsRef<Path> for DecInt {
#[inline]
fn as_ref(&self) -> &Path {
let as_os_str: &OsStr = OsStrExt::from_bytes(&self.buf[..self.len]);
let as_os_str: &OsStr = OsStrExt::from_bytes(self.as_bytes());
Path::new(as_os_str)
}
}
Expand Down
55 changes: 40 additions & 15 deletions tests/path/dec_int.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,45 @@
use rustix::path::DecInt;

macro_rules! check {
($i:expr) => {
let i = $i;
assert_eq!(DecInt::new(i).as_ref().to_str().unwrap(), i.to_string());
};
}

#[test]
fn test_dec_int() {
assert_eq!(DecInt::new(0).as_ref().to_str().unwrap(), "0");
assert_eq!(DecInt::new(-1).as_ref().to_str().unwrap(), "-1");
assert_eq!(DecInt::new(789).as_ref().to_str().unwrap(), "789");
assert_eq!(
DecInt::new(i64::MIN).as_ref().to_str().unwrap(),
i64::MIN.to_string()
);
assert_eq!(
DecInt::new(i64::MAX).as_ref().to_str().unwrap(),
i64::MAX.to_string()
);
assert_eq!(
DecInt::new(u64::MAX).as_ref().to_str().unwrap(),
u64::MAX.to_string()
);
check!(0);
check!(-1);
check!(789);

check!(u8::MAX);
check!(i8::MIN);
check!(u16::MAX);
check!(i16::MIN);
check!(u32::MAX);
check!(i32::MIN);
check!(u64::MAX);
check!(i64::MIN);
#[cfg(any(
target_pointer_width = "16",
target_pointer_width = "32",
target_pointer_width = "64"
))]
{
check!(usize::MAX);
check!(isize::MIN);
}
}

#[test]
#[should_panic]
fn test_unsupported_max_u128_dec_int() {
check!(u128::MAX);
}

#[test]
#[should_panic]
fn test_unsupported_min_u128_dec_int() {
check!(i128::MIN);
}

0 comments on commit 61f17d6

Please sign in to comment.