Skip to content

Commit

Permalink
Ensure that HSTRING builder provides initialized memory (#3141)
Browse files Browse the repository at this point in the history
  • Loading branch information
kennykerr authored Jul 3, 2024
1 parent 12a60df commit f6c49f4
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 17 deletions.
8 changes: 7 additions & 1 deletion crates/libs/strings/src/hstring_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,13 @@ pub struct HStringBuilder(*mut HStringHeader);
impl HStringBuilder {
/// Creates a preallocated `HSTRING` value.
pub fn new(len: usize) -> Result<Self> {
Ok(Self(HStringHeader::alloc(len.try_into()?)?))
let header = HStringHeader::alloc(len.try_into()?)?;

if len > 0 {
unsafe { core::ptr::write_bytes((*header).data, 0, len) };
}

Ok(Self(header))
}

/// Shortens the string by removing any trailing 0 characters.
Expand Down
28 changes: 12 additions & 16 deletions crates/libs/strings/src/hstring_header.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,26 +14,26 @@ pub struct HStringHeader {
}

impl HStringHeader {
pub fn alloc(len: u32) -> Result<*mut HStringHeader> {
pub fn alloc(len: u32) -> Result<*mut Self> {
if len == 0 {
return Ok(core::ptr::null_mut());
}

// Allocate enough space for header and two bytes per character.
// The space for the terminating null character is already accounted for inside of `HStringHeader`.
let bytes = core::mem::size_of::<HStringHeader>() + 2 * len as usize;
let bytes = core::mem::size_of::<Self>() + 2 * len as usize;

#[cfg(windows)]
let header = unsafe { bindings::HeapAlloc(bindings::GetProcessHeap(), 0, bytes) }
as *mut HStringHeader;
let header =
unsafe { bindings::HeapAlloc(bindings::GetProcessHeap(), 0, bytes) } as *mut Self;

#[cfg(not(windows))]
let header = unsafe {
extern "C" {
fn malloc(bytes: usize) -> *mut core::ffi::c_void;
}

malloc(bytes) as *mut HStringHeader
malloc(bytes) as *mut Self
};

if header.is_null() {
Expand All @@ -42,7 +42,7 @@ impl HStringHeader {

unsafe {
// Use `ptr::write` (since `header` is unintialized). `HStringHeader` is safe to be all zeros.
header.write(core::mem::MaybeUninit::<HStringHeader>::zeroed().assume_init());
header.write(core::mem::MaybeUninit::<Self>::zeroed().assume_init());
(*header).len = len;
(*header).count = RefCount::new(1);
(*header).data = &mut (*header).buffer_start;
Expand All @@ -51,36 +51,32 @@ impl HStringHeader {
Ok(header)
}

pub unsafe fn free(header: *mut HStringHeader) {
pub unsafe fn free(header: *mut Self) {
if header.is_null() {
return;
}

let header = header as *mut _;

#[cfg(windows)]
{
bindings::HeapFree(bindings::GetProcessHeap(), 0, header);
}
bindings::HeapFree(bindings::GetProcessHeap(), 0, header as *mut _);

#[cfg(not(windows))]
{
extern "C" {
fn free(ptr: *mut core::ffi::c_void);
}

free(header);
free(header as *mut _);
}
}

pub fn duplicate(&self) -> Result<*mut HStringHeader> {
pub fn duplicate(&self) -> Result<*mut Self> {
if self.flags & HSTRING_REFERENCE_FLAG == 0 {
// If this is not a "fast pass" string then simply increment the reference count.
self.count.add_ref();
Ok(self as *const HStringHeader as *mut HStringHeader)
Ok(self as *const Self as *mut Self)
} else {
// Otherwise, allocate a new string and copy the value into the new string.
let copy = HStringHeader::alloc(self.len)?;
let copy = Self::alloc(self.len)?;
// SAFETY: since we are duplicating the string it is safe to copy all data from self to the initialized `copy`.
// We copy `len + 1` characters since `len` does not account for the terminating null character.
unsafe {
Expand Down
4 changes: 4 additions & 0 deletions crates/tests/strings/tests/hstring.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,5 +47,9 @@ fn hstring_builder() -> Result<()> {
assert_eq!(h.len(), 5);
assert_eq!(h.as_wide(), HELLO);

// HStringBuilder will initialize memory to zero.
let b = HStringBuilder::new(5)?;
assert_eq!(*b, [0, 0, 0, 0, 0]);

Ok(())
}

0 comments on commit f6c49f4

Please sign in to comment.