Workaround for mem safety in third party dlls

1 files changed, 73 insertions, 0 deletions

diff --git a/library/std/src/sys/pal/windows/api.rs b/library/std/src/sys/pal/windows/api.rs
index 773455c572f..25a6c2d7d8e 100644
--- a/library/std/src/sys/pal/windows/api.rs
+++ b/library/std/src/sys/pal/windows/api.rs
@@ -30,6 +30,7 @@
 //! should go in sys/pal/windows/mod.rs rather than here. See `IoResult` as an example.
 
 use core::ffi::c_void;
+use core::marker::PhantomData;
 
 use super::c;
 
@@ -291,3 +292,75 @@ impl WinError {
     pub const TIMEOUT: Self = Self::new(c::ERROR_TIMEOUT);
     // tidy-alphabetical-end
 }
+
+/// A wrapper around a UNICODE_STRING that is equivalent to `&[u16]`.
+///
+/// It is preferable to use the `unicode_str!` macro as that contains mitigations for  #143078.
+///
+/// If the MaximumLength field of the underlying UNICODE_STRING is greater than
+/// the Length field then you can test if the string is null terminated by inspecting
+/// the u16 directly after the string. You cannot otherwise depend on nul termination.
+#[derive(Copy, Clone)]
+pub struct UnicodeStrRef<'a> {
+    s: c::UNICODE_STRING,
+    lifetime: PhantomData<&'a [u16]>,
+}
+
+static EMPTY_STRING_NULL_TERMINATED: &[u16] = &[0];
+
+impl UnicodeStrRef<'_> {
+    const fn new(slice: &[u16], is_null_terminated: bool) -> Self {
+        let (len, max_len, ptr) = if slice.is_empty() {
+            (0, 2, EMPTY_STRING_NULL_TERMINATED.as_ptr().cast_mut())
+        } else {
+            let len = slice.len() - (is_null_terminated as usize);
+            (len * 2, size_of_val(slice), slice.as_ptr().cast_mut())
+        };
+        Self {
+            s: c::UNICODE_STRING { Length: len as _, MaximumLength: max_len as _, Buffer: ptr },
+            lifetime: PhantomData,
+        }
+    }
+
+    pub const fn from_slice_with_nul(slice: &[u16]) -> Self {
+        if !slice.is_empty() {
+            debug_assert!(slice[slice.len() - 1] == 0);
+        }
+        Self::new(slice, true)
+    }
+
+    pub const fn from_slice(slice: &[u16]) -> Self {
+        Self::new(slice, false)
+    }
+
+    /// Returns a pointer to the underlying UNICODE_STRING
+    pub const fn as_ptr(&self) -> *const c::UNICODE_STRING {
+        &self.s
+    }
+}
+
+/// Create a UnicodeStringRef from a literal str or a u16 array.
+///
+/// To mitigate #143078, when using a literal str the created UNICODE_STRING
+/// will be nul terminated. The MaximumLength field of the UNICODE_STRING will
+/// be set greater than the Length field to indicate that a nul may be present.
+///
+/// If using a u16 array, the array is used exactly as provided and you cannot
+/// count on the string being nul terminated.
+/// This should generally be used for strings that come from the OS.
+///
+/// **NOTE:** we lack a UNICODE_STRING builder type as we don't currently have
+/// a use for it. If needing to dynamically build a UNICODE_STRING, the builder
+/// should try to ensure there's a nul one past the end of the string.
+pub macro unicode_str {
+    ($str:literal) => {const {
+        crate::sys::pal::windows::api::UnicodeStrRef::from_slice_with_nul(
+            crate::sys::pal::windows::api::wide_str!($str),
+        )
+    }},
+    ($array:expr) => {
+        crate::sys::pal::windows::api::UnicodeStrRef::from_slice(
+            $array,
+        )
+    }
+}