Linux: DXC BSTR now resembles its Windows counterpart with explicit size (#11)

We [recently] fixed memory leaks when calling into the intellisense part
of DXC and made the decision shortly after merging to [implement
null-invariant `BSTR`s properly] in DXC's shimmed type for non-Windows,
which was promptly accepted.

Currently our implementation of `SysFreeString` together with that PR
segfaults because it frees the pointer at offset `+4` of the allocation.
That is corrected together with `SysStringLen` now reading the size of
the string from this prefix instead of looking for a (wrong!) null
termination character.  At least our Windows and Linux implementation
`utils.rs` is the same again.

[recently]: #10
[implement null-invariant `BSTR`s properly]: microsoft/DirectXShaderCompiler#3250

Author: MarijnS95

src/os.rs

@@ -9,13 +9,14 @@ mod os_defs {
     };
 
     pub use winapi::um::combaseapi::CoTaskMemFree;
-    pub use winapi::um::oleauto::SysFreeString;
+    pub use winapi::um::oleauto::{SysFreeString, SysStringLen};
 }
 
 #[cfg(not(windows))]
 mod os_defs {
     pub type CHAR = i8;
-    pub type WCHAR = u32;
+    pub type UINT = u32;
+    pub type WCHAR = widestring::WideChar;
     pub type OLECHAR = WCHAR;
     pub type LPSTR = *mut CHAR;
     pub type LPWSTR = *mut WCHAR;
@@ -25,20 +26,57 @@ mod os_defs {
     pub type LPBSTR = *mut BSTR;
     pub type HRESULT = i32;
 
+    /// Returns a mutable pointer to the length prefix of the string
+    fn len_ptr(p: BSTR) -> *mut UINT {
+        // The first four bytes before the pointer contain the length prefix:
+        // https://docs.microsoft.com/en-us/previous-versions/windows/desktop/automat/bstr#remarks
+        unsafe { p.cast::<UINT>().offset(-1) }
+    }
+
     #[allow(non_snake_case)]
     /// # Safety
-    /// `p` must be a valid pointer to an allocation made with `malloc`
+    /// `p` must be a valid pointer to an allocation made with `malloc`, or null.
     pub unsafe fn CoTaskMemFree(p: *mut libc::c_void) {
-        // https://github.com/microsoft/DirectXShaderCompiler/blob/a8d9780046cb64a1cea842fa6fc28a250e3e2c09/include/dxc/Support/WinAdapter.h#L46
-        libc::free(p)
+        // https://github.com/microsoft/DirectXShaderCompiler/blob/56e22b30c5e43632f56a1f97865f37108ec35463/include/dxc/Support/WinAdapter.h#L46
+        if !p.is_null() {
+            libc::free(p)
+        }
     }
 
     #[allow(non_snake_case)]
     /// # Safety
-    /// `p` must be a valid pointer to an allocation made with `malloc`
+    /// `p` must be a valid pointer to an allocation made with `malloc`, or null.
     pub unsafe fn SysFreeString(p: BSTR) {
-        // https://github.com/microsoft/DirectXShaderCompiler/blob/a8d9780046cb64a1cea842fa6fc28a250e3e2c09/include/dxc/Support/WinAdapter.h#L48-L50
-        libc::free(p as _)
+        // https://github.com/microsoft/DirectXShaderCompiler/blob/56e22b30c5e43632f56a1f97865f37108ec35463/lib/DxcSupport/WinAdapter.cpp#L50-L53
+        if !p.is_null() {
+            libc::free(len_ptr(p).cast::<_>())
+        }
+    }
+
+    #[allow(non_snake_case)]
+    /// Returns the size of `p` in bytes, excluding terminating NULL character
+    ///
+    /// # Safety
+    /// `p` must be a valid pointer to a [`BSTR`] with size-prefix in the `4` leading bytes, or null.
+    ///
+    /// https://docs.microsoft.com/en-us/previous-versions/windows/desktop/automat/bstr#remarks
+    pub unsafe fn SysStringByteLen(p: BSTR) -> UINT {
+        if p.is_null() {
+            0
+        } else {
+            *len_ptr(p)
+        }
+    }
+
+    #[allow(non_snake_case)]
+    /// Returns the size of `p` in characters, excluding terminating NULL character
+    ///
+    /// # Safety
+    /// `p` must be a valid pointer to a [`BSTR`] with size-prefix in the `4` leading bytes, or null.
+    ///
+    /// https://docs.microsoft.com/en-us/previous-versions/windows/desktop/automat/bstr#remarks
+    pub unsafe fn SysStringLen(p: BSTR) -> UINT {
+        SysStringByteLen(p) / std::mem::size_of::<OLECHAR>() as UINT
     }
 }
 

src/utils.rs

@@ -1,12 +1,9 @@
 use std::path::PathBuf;
 
-use crate::os::{SysFreeString, BSTR, HRESULT, LPSTR, LPWSTR, WCHAR};
+use crate::os::{SysFreeString, SysStringLen, BSTR, HRESULT, LPSTR, LPWSTR, WCHAR};
 use crate::wrapper::*;
 use thiserror::Error;
 
-#[cfg(windows)]
-use winapi::um::oleauto::SysStringLen;
-
 pub(crate) fn to_wide(msg: &str) -> Vec<WCHAR> {
     widestring::WideCString::from_str(msg)
         .unwrap()
@@ -21,7 +18,6 @@ pub(crate) fn from_wide(wide: LPWSTR) -> String {
     }
 }
 
-#[cfg(windows)]
 pub(crate) fn from_bstr(string: BSTR) -> String {
     unsafe {
         let len = SysStringLen(string) as usize;
@@ -35,20 +31,6 @@ pub(crate) fn from_bstr(string: BSTR) -> String {
     }
 }
 
-#[cfg(not(windows))]
-pub(crate) fn from_bstr(string: BSTR) -> String {
-    // TODO (Marijn): This does NOT cover embedded NULLs
-
-    // BSTR contains its size in the four bytes preceding the pointer, in order to contain NULL bytes:
-    // https://docs.microsoft.com/en-us/previous-versions/windows/desktop/automat/bstr
-    // DXC on non-Windows does not adhere to that and simply allocates a buffer without prepending the size:
-    // https://github.com/microsoft/DirectXShaderCompiler/blob/a8d9780046cb64a1cea842fa6fc28a250e3e2c09/include/dxc/Support/WinAdapter.h#L49-L50
-    let result = from_wide(string as LPWSTR);
-
-    unsafe { SysFreeString(string) };
-    result
-}
-
 pub(crate) fn from_lpstr(string: LPSTR) -> String {
     unsafe {
         let len = (0..).take_while(|&i| *string.offset(i) != 0).count();