Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Guarantee address in split_off/split_to for empty slices #740

Merged
merged 2 commits into from
Oct 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 28 additions & 5 deletions src/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,22 @@ impl Bytes {
}
}

/// Creates a new `Bytes` with length zero and the given pointer as the address.
fn new_empty_with_ptr(ptr: *const u8) -> Self {
debug_assert!(!ptr.is_null());

// Detach this pointer's provenance from whichever allocation it came from, and reattach it
// to the provenance of the fake ZST [u8;0] at the same address.
let ptr = without_provenance(ptr as usize);

Bytes {
ptr,
len: 0,
data: AtomicPtr::new(ptr::null_mut()),
vtable: &STATIC_VTABLE,
}
}

/// Returns the number of bytes contained in this `Bytes`.
///
/// # Examples
Expand Down Expand Up @@ -366,7 +382,9 @@ impl Bytes {
/// Splits the bytes into two at the given index.
///
/// Afterwards `self` contains elements `[0, at)`, and the returned `Bytes`
/// contains elements `[at, len)`.
/// contains elements `[at, len)`. It's guaranteed that the memory does not
/// move, that is, the address of `self` does not change, and the address of
/// the returned slice is `at` bytes after that.
///
/// This is an `O(1)` operation that just increases the reference count and
/// sets a few indices.
Expand All @@ -389,11 +407,11 @@ impl Bytes {
#[must_use = "consider Bytes::truncate if you don't need the other half"]
pub fn split_off(&mut self, at: usize) -> Self {
if at == self.len() {
return Bytes::new();
return Bytes::new_empty_with_ptr(self.ptr.wrapping_add(at));
}

if at == 0 {
return mem::replace(self, Bytes::new());
return mem::replace(self, Bytes::new_empty_with_ptr(self.ptr));
}

assert!(
Expand Down Expand Up @@ -438,11 +456,12 @@ impl Bytes {
#[must_use = "consider Bytes::advance if you don't need the other half"]
pub fn split_to(&mut self, at: usize) -> Self {
if at == self.len() {
return mem::replace(self, Bytes::new());
let end_ptr = self.ptr.wrapping_add(at);
return mem::replace(self, Bytes::new_empty_with_ptr(end_ptr));
}

if at == 0 {
return Bytes::new();
return Bytes::new_empty_with_ptr(self.ptr);
}

assert!(
Expand Down Expand Up @@ -1426,6 +1445,10 @@ where
new_addr as *mut u8
}

fn without_provenance(ptr: usize) -> *const u8 {
core::ptr::null::<u8>().wrapping_add(ptr)
}

// compile-fails

/// ```compile_fail
Expand Down
4 changes: 3 additions & 1 deletion src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,9 @@ impl BytesMut {
/// Splits the bytes into two at the given index.
///
/// Afterwards `self` contains elements `[0, at)`, and the returned
/// `BytesMut` contains elements `[at, capacity)`.
/// `BytesMut` contains elements `[at, capacity)`. It's guaranteed that the
/// memory does not move, that is, the address of `self` does not change,
/// and the address of the returned slice is `at` bytes after that.
///
/// This is an `O(1)` operation that just increases the reference count
/// and sets a few indices.
Expand Down
80 changes: 80 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1399,3 +1399,83 @@ fn try_reclaim_arc() {
buf.advance(2);
assert_eq!(true, buf.try_reclaim(6));
}

#[test]
fn split_off_empty_addr() {
let mut buf = Bytes::from(vec![0; 1024]);

let ptr_start = buf.as_ptr();
let ptr_end = ptr_start.wrapping_add(1024);

let empty_end = buf.split_off(1024);
assert_eq!(empty_end.len(), 0);
assert_eq!(empty_end.as_ptr(), ptr_end);

let _ = buf.split_off(0);
assert_eq!(buf.len(), 0);
assert_eq!(buf.as_ptr(), ptr_start);

// Is miri happy about the provenance?
let _ = &empty_end[..];
let _ = &buf[..];
}

#[test]
fn split_to_empty_addr() {
let mut buf = Bytes::from(vec![0; 1024]);

let ptr_start = buf.as_ptr();
let ptr_end = ptr_start.wrapping_add(1024);

let empty_start = buf.split_to(0);
assert_eq!(empty_start.len(), 0);
assert_eq!(empty_start.as_ptr(), ptr_start);

let _ = buf.split_to(1024);
assert_eq!(buf.len(), 0);
assert_eq!(buf.as_ptr(), ptr_end);

// Is miri happy about the provenance?
let _ = &empty_start[..];
let _ = &buf[..];
}

#[test]
fn split_off_empty_addr_mut() {
let mut buf = BytesMut::from([0; 1024].as_slice());

let ptr_start = buf.as_ptr();
let ptr_end = ptr_start.wrapping_add(1024);

let empty_end = buf.split_off(1024);
assert_eq!(empty_end.len(), 0);
assert_eq!(empty_end.as_ptr(), ptr_end);

let _ = buf.split_off(0);
assert_eq!(buf.len(), 0);
assert_eq!(buf.as_ptr(), ptr_start);

// Is miri happy about the provenance?
let _ = &empty_end[..];
let _ = &buf[..];
}

#[test]
fn split_to_empty_addr_mut() {
let mut buf = BytesMut::from([0; 1024].as_slice());

let ptr_start = buf.as_ptr();
let ptr_end = ptr_start.wrapping_add(1024);

let empty_start = buf.split_to(0);
assert_eq!(empty_start.len(), 0);
assert_eq!(empty_start.as_ptr(), ptr_start);

let _ = buf.split_to(1024);
assert_eq!(buf.len(), 0);
assert_eq!(buf.as_ptr(), ptr_end);

// Is miri happy about the provenance?
let _ = &empty_start[..];
let _ = &buf[..];
}