Ensure that `HSTRING` builder provides initialized memory (#3141)

This commit is contained in:
Kenny Kerr 2024-07-03 10:46:10 -05:00 коммит произвёл GitHub
Родитель 12a60df721
Коммит f6c49f4f5e
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: B5690EEEBB952194
3 изменённых файлов: 23 добавлений и 17 удалений

Просмотреть файл

@ -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.

Просмотреть файл

@ -14,18 +14,18 @@ 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 {
@ -33,7 +33,7 @@ impl HStringHeader {
fn malloc(bytes: usize) -> *mut core::ffi::c_void;
}
malloc(bytes) as *mut HStringHeader
malloc(bytes) as *mut Self
};
if header.is_null() {
@ -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;
@ -51,17 +51,13 @@ 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))]
{
@ -69,18 +65,18 @@ impl HStringHeader {
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 {

Просмотреть файл

@ -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(())
}