From ca50c595b6ea1b355aa1d15aeaeac60823d8ecbd Mon Sep 17 00:00:00 2001
From: Marijn Suijten <marijn@traverseresearch.nl>
Date: Tue, 10 Nov 2020 15:54:08 +0100
Subject: [PATCH] Linux: DXC BSTR now resembles its Windows counterpart with
 explicit size

Since https://github.com/microsoft/DirectXShaderCompiler/pull/3250
---
 src/os.rs    | 33 ++++++++++++++++++++++++++++-----
 src/utils.rs | 20 +-------------------
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/src/os.rs b/src/os.rs
index d46c25e..90d9e99 100644
--- a/src/os.rs
+++ b/src/os.rs
@@ -6,13 +6,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,13 +26,35 @@ mod os_defs {
     #[allow(non_snake_case)]
     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)
+        if !p.is_null() {
+            libc::free(p)
+        }
     }
 
     #[allow(non_snake_case)]
     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/2ade6f84d6b95bfd96eec1d6d15e3aa3b519d180/lib/DxcSupport/WinAdapter.cpp#L73-L76
+        if !p.is_null() {
+            libc::free((p as *mut UINT).offset(-1) as *mut _)
+        }
+    }
+
+    /// Returns the size of `p` in bytes, without terminating NULL character
+    #[allow(non_snake_case)]
+    pub unsafe fn SysStringByteLen(p: BSTR) -> UINT {
+        // The first four bytes before the pointer contain the length prefix:
+        // https://docs.microsoft.com/en-us/previous-versions/windows/desktop/automat/bstr
+        if p.is_null() {
+            0
+        } else {
+            *(p as *const UINT).offset(-1)
+        }
+    }
+
+    /// Returns the size of `p` in characters, without terminating NULL character
+    #[allow(non_snake_case)]
+    pub unsafe fn SysStringLen(p: BSTR) -> UINT {
+        SysStringByteLen(p) / std::mem::size_of::<OLECHAR>() as UINT
     }
 }
 
diff --git a/src/utils.rs b/src/utils.rs
index a24f7f4..6fa2932 100644
--- a/src/utils.rs
+++ b/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().into_vec()
 }
@@ -19,7 +16,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;
@@ -33,20 +29,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();