host_env: Fix wcslen, leverage NonNull<T> more (#8386)

`strlen` was both unused and unneeded. Rust's `CStr` handles the length
for us without needing to call into `libc` manually or implementing our
own version of `strlen`. Our `strlen` implemented a fallback that
incremented and dereferenced a pointer to check for NUL. This is slower
than what `libc` usually does. For example, `musl` operates on words
instead of individual bytes. This implementation also caused UB if the
caller passed in a null pointer.

`wcslen` is similar. The currently implementation can cause UB for null
pointers. It's also slow from the same reason mentioned above. Luckily,
the widestring crate can handle this for us. The crate panics on null
pointers, so using it was a good excuse to use NonNull in more places.

crates/host_env/Cargo.toml

@@ -17,6 +17,7 @@ libc = { workspace = true }
 num-traits = { workspace = true }
 parking_lot = { workspace = true }
 paste = { workspace = true }
+widestring = { workspace = true }
 
 [target.'cfg(unix)'.dependencies]
 nix = { workspace = true }
@@ -54,7 +55,6 @@ system-configuration = { workspace = true }
 memchr.workspace = true
 junction = { workspace = true }
 schannel = { workspace = true }
-widestring = { workspace = true }
 windows-sys = { workspace = true, features = [
     "Win32_Foundation",
     "Win32_Globalization",

crates/host_env/src/ctypes.rs

@@ -3,6 +3,7 @@ use core::ffi::{
     CStr, c_char, c_double, c_float, c_int, c_long, c_longlong, c_schar, c_short, c_uchar, c_uint,
     c_ulong, c_ulonglong, c_ushort, c_void,
 };
+use core::ptr::NonNull;
 #[cfg(all(
     any(
         target_os = "linux",
@@ -36,6 +37,7 @@ use rustpython_wtf8::Wtf8;
 use rustpython_wtf8::Wtf8Buf;
 #[cfg(any(unix, windows))]
 use std::{collections::HashMap, ffi::OsStr, sync::OnceLock};
+use widestring::WideCStr;
 
 #[cfg(all(
     any(
@@ -392,33 +394,11 @@ pub fn dyld_shared_cache_contains_path(path: &str) -> Result<bool, alloc::ffi::N
     Ok(unsafe { _dyld_shared_cache_contains_path(c_path.as_ptr()) })
 }
 
-/// # Safety
-///
-/// `ptr` must be valid to read until the first NUL byte.
-pub unsafe fn strlen(ptr: *const c_char) -> usize {
-    #[cfg(any(unix, windows, target_os = "wasi"))]
-    {
-        unsafe { libc::strlen(ptr) }
-    }
-    #[cfg(not(any(unix, windows, target_os = "wasi")))]
-    {
-        let mut len = 0;
-        while unsafe { *ptr.add(len) } != 0 {
-            len += 1;
-        }
-        len
-    }
-}
-
 /// # Safety
 ///
 /// `ptr` must be valid to read until the first NUL wide character.
-pub unsafe fn wcslen(ptr: *const WChar) -> usize {
-    let mut len = 0;
-    while unsafe { *ptr.add(len) } != 0 as WChar {
-        len += 1;
-    }
-    len
+pub unsafe fn wcslen(ptr: NonNull<WChar>) -> usize {
+    unsafe { WideCStr::from_ptr_str(ptr.as_ptr().cast()).len() }
 }
 
 /// # Safety
@@ -1309,10 +1289,10 @@ pub unsafe fn callback_arg_value(type_code: Option<&str>, ptr: *const c_void) ->
         }
         Some("Z") => {
             let wstr_ptr = unsafe { *(ptr as *const *const WChar) };
-            if wstr_ptr.is_null() {
-                DecodedValue::None
-            } else {
+            if let Some(wstr_ptr) = NonNull::new(wstr_ptr.cast_mut()) {
                 DecodedValue::String(unsafe { read_wide_string(wstr_ptr) }.to_string())
+            } else {
+                DecodedValue::None
             }
         }
         Some("P") => DecodedValue::Pointer(unsafe { *(ptr as *const usize) }),
@@ -2272,8 +2252,7 @@ pub unsafe fn borrowed_slice_as_mut(slice: &[u8]) -> &mut [u8] {
 pub fn wide_chars_to_wtf8(wchars: &[WChar]) -> Wtf8Buf {
     #[cfg(windows)]
     {
-        let wide: Vec<u16> = wchars.to_vec();
-        Wtf8Buf::from_wide(&wide)
+        Wtf8Buf::from_wide(wchars)
     }
     #[cfg(not(windows))]
     {
@@ -2292,10 +2271,10 @@ pub fn wide_chars_to_wtf8(wchars: &[WChar]) -> Wtf8Buf {
 /// # Safety
 ///
 /// `ptr` must be a valid NUL-terminated wide C string.
-pub unsafe fn read_wide_string(ptr: *const WChar) -> Wtf8Buf {
-    let len = unsafe { wcslen(ptr) };
-    let wchars = unsafe { core::slice::from_raw_parts(ptr, len) };
-    wide_chars_to_wtf8(wchars)
+pub unsafe fn read_wide_string(ptr: NonNull<WChar>) -> Wtf8Buf {
+    // SAFETY: WideCStr does not assume an encoding.
+    let wchars = unsafe { WideCStr::from_ptr_str(ptr.as_ptr().cast()) };
+    Wtf8Buf::from_string(wchars.to_string_lossy())
 }
 
 /// # Safety
@@ -2313,18 +2292,15 @@ pub unsafe fn read_c_string_from_address(addr: usize) -> Option<Vec<u8>> {
 ///
 /// `addr` must either be zero or a valid NUL-terminated wide C string pointer.
 pub unsafe fn read_wide_string_from_address(addr: usize) -> Option<Wtf8Buf> {
-    if addr == 0 {
-        None
-    } else {
-        Some(unsafe { read_wide_string(addr as *const WChar) })
-    }
+    let ptr = NonNull::new(addr as *mut WChar)?;
+    Some(unsafe { read_wide_string(ptr) })
 }
 
 /// # Safety
 ///
 /// `ptr` must point to `len` readable wide characters.
-pub unsafe fn read_wide_string_with_len(ptr: *const WChar, len: usize) -> Wtf8Buf {
-    let wchars = unsafe { core::slice::from_raw_parts(ptr, len) };
+pub unsafe fn read_wide_string_with_len(ptr: NonNull<WChar>, len: usize) -> Wtf8Buf {
+    let wchars = unsafe { core::slice::from_raw_parts(ptr.as_ptr(), len) };
     wide_chars_to_wtf8(wchars)
 }
 
@@ -2348,13 +2324,12 @@ pub fn string_at(ptr: usize, size: isize) -> Result<Vec<u8>, StringAtError> {
 }
 
 pub fn wstring_at(ptr: usize, size: isize) -> Result<Wtf8Buf, StringAtError> {
-    if ptr == 0 {
+    let Some(ptr) = NonNull::new(ptr as *mut WChar) else {
         return Err(StringAtError::NullPointer);
-    }
-    let w_ptr = ptr as *const WChar;
+    };
     if size < 0 {
         // SAFETY: caller passed a non-null NUL-terminated wide string pointer.
-        return Ok(unsafe { read_wide_string(w_ptr) });
+        return Ok(unsafe { read_wide_string(ptr) });
     }
     let len = {
         let size_usize = size as usize;
@@ -2364,7 +2339,7 @@ pub fn wstring_at(ptr: usize, size: isize) -> Result<Wtf8Buf, StringAtError> {
         size_usize
     };
     // SAFETY: caller requested exactly `len` readable wide characters from non-null pointer.
-    Ok(unsafe { read_wide_string_with_len(w_ptr, len) })
+    Ok(unsafe { read_wide_string_with_len(ptr, len) })
 }
 
 /// # Safety
@@ -2414,14 +2389,14 @@ pub unsafe fn read_pointer_char_slice(
 /// # Safety
 ///
 /// `start` must be valid to read `len` wide characters following `step`.
-pub unsafe fn read_wide_string_strided(start: *const WChar, len: usize, step: isize) -> Wtf8Buf {
+pub unsafe fn read_wide_string_strided(start: NonNull<WChar>, len: usize, step: isize) -> Wtf8Buf {
     if step == 1 {
         return unsafe { read_wide_string_with_len(start, len) };
     }
     let mut wchars = Vec::with_capacity(len);
     let mut cur = start;
     for _ in 0..len {
-        wchars.push(unsafe { *cur });
+        wchars.push(unsafe { cur.read() });
         cur = unsafe { cur.offset(step) };
     }
     wide_chars_to_wtf8(&wchars)
@@ -2436,10 +2411,9 @@ pub unsafe fn read_pointer_wchar_slice(
     start: isize,
     len: usize,
     step: isize,
-) -> Wtf8Buf {
-    let wchar_size = core::mem::size_of::<WChar>();
-    let start_addr = (ptr_value as isize + start * wchar_size as isize) as *const WChar;
-    unsafe { read_wide_string_strided(start_addr, len, step) }
+) -> Option<Wtf8Buf> {
+    let start_addr = unsafe { NonNull::new(ptr_value as *mut WChar)?.offset(start) };
+    Some(unsafe { read_wide_string_strided(start_addr, len, step) })
 }
 
 /// # Safety

crates/host_env/src/wmi.rs

@@ -6,7 +6,7 @@
 #![allow(unsafe_op_in_unsafe_fn)]
 
 use core::ffi::c_void;
-use core::ptr::{null, null_mut};
+use core::ptr::{NonNull, null, null_mut};
 use windows_sys::Win32::Foundation::{
     CloseHandle, ERROR_BROKEN_PIPE, ERROR_MORE_DATA, ERROR_NOT_ENOUGH_MEMORY, GetLastError, HANDLE,
     WAIT_OBJECT_0, WAIT_TIMEOUT,
@@ -17,6 +17,8 @@ use windows_sys::Win32::System::Threading::{
     CreateEventW, CreateThread, GetExitCodeThread, SetEvent, WaitForSingleObject,
 };
 
+use crate::ctypes::wcslen;
+
 pub const BUFFER_SIZE: usize = 8192;
 
 pub enum ExecQueryError {
@@ -238,7 +240,7 @@ unsafe fn object_end_enumeration(this: *mut c_void) -> HRESULT {
     method(this)
 }
 
-fn hresult_from_win32(err: u32) -> HRESULT {
+const fn hresult_from_win32(err: u32) -> HRESULT {
     if err == 0 {
         0
     } else {
@@ -246,11 +248,11 @@ fn hresult_from_win32(err: u32) -> HRESULT {
     }
 }
 
-fn succeeded(hr: HRESULT) -> bool {
+const fn succeeded(hr: HRESULT) -> bool {
     hr >= 0
 }
 
-fn failed(hr: HRESULT) -> bool {
+const fn failed(hr: HRESULT) -> bool {
     hr < 0
 }
 
@@ -258,14 +260,6 @@ fn wide_str(s: &str) -> Vec<u16> {
     s.encode_utf16().chain(core::iter::once(0)).collect()
 }
 
-unsafe fn wcslen(s: *const u16) -> usize {
-    let mut len = 0;
-    while unsafe { *s.add(len) } != 0 {
-        len += 1;
-    }
-    len
-}
-
 unsafe fn wait_event(event: HANDLE, timeout: u32) -> u32 {
     match unsafe { WaitForSingleObject(event, timeout) } {
         WAIT_OBJECT_0 => 0,
@@ -471,16 +465,25 @@ unsafe fn query_thread_impl(param: *mut c_void) -> u32 {
             }
 
             if succeeded(hr) && (flavor & WBEM_FLAVOR_MASK_ORIGIN) != WBEM_FLAVOR_ORIGIN_SYSTEM {
+                let Some(cb_str1) = NonNull::new(prop_name)
+                    .map(|prop_name| (unsafe { wcslen(prop_name) } * 2) as u32)
+                else {
+                    unsafe {
+                        SysFreeString(prop_name);
+                    }
+                    break;
+                };
+
                 let mut prop_str = [0u16; BUFFER_SIZE];
                 hr = unsafe {
                     VariantToString(&prop_value, prop_str.as_mut_ptr(), BUFFER_SIZE as u32)
                 };
+                let cb_str2 = NonNull::new(prop_str.as_ptr().cast_mut())
+                    .map(|prop_str| (unsafe { wcslen(prop_str) } * 2) as u32)
+                    .expect("prop_str is never null");
 
-                if succeeded(hr) {
-                    let cb_str1 = (unsafe { wcslen(prop_name) } * 2) as u32;
-                    let cb_str2 = (unsafe { wcslen(prop_str.as_ptr()) } * 2) as u32;
-
-                    if unsafe {
+                if succeeded(hr)
+                    && unsafe {
                         WriteFile(
                             write_pipe,
                             prop_name as *const _,
@@ -489,36 +492,35 @@ unsafe fn query_thread_impl(param: *mut c_void) -> u32 {
                             null_mut(),
                         )
                     } == 0
-                        || unsafe {
-                            WriteFile(
-                                write_pipe,
-                                &eq_sign as *const u16 as *const _,
-                                2,
-                                &mut written,
-                                null_mut(),
-                            )
-                        } == 0
-                        || unsafe {
-                            WriteFile(
-                                write_pipe,
-                                prop_str.as_ptr() as *const _,
-                                cb_str2,
-                                &mut written,
-                                null_mut(),
-                            )
-                        } == 0
-                        || unsafe {
-                            WriteFile(
-                                write_pipe,
-                                &null_sep as *const u16 as *const _,
-                                2,
-                                &mut written,
-                                null_mut(),
-                            )
-                        } == 0
-                    {
-                        hr = hresult_from_win32(unsafe { GetLastError() });
-                    }
+                    || unsafe {
+                        WriteFile(
+                            write_pipe,
+                            &eq_sign as *const u16 as *const _,
+                            2,
+                            &mut written,
+                            null_mut(),
+                        )
+                    } == 0
+                    || unsafe {
+                        WriteFile(
+                            write_pipe,
+                            prop_str.as_ptr() as *const _,
+                            cb_str2,
+                            &mut written,
+                            null_mut(),
+                        )
+                    } == 0
+                    || unsafe {
+                        WriteFile(
+                            write_pipe,
+                            &null_sep as *const u16 as *const _,
+                            2,
+                            &mut written,
+                            null_mut(),
+                        )
+                    } == 0
+                {
+                    hr = hresult_from_win32(unsafe { GetLastError() });
                 }
 
                 unsafe {

crates/vm/src/stdlib/_ctypes/pointer.rs

@@ -524,13 +524,12 @@ impl PyCPointer {
 
         // c_wchar → str
         if type_code.as_deref() == Some("u") {
-            if len == 0 {
-                return Ok(vm.ctx.new_str("").into());
+            if len > 0
+                && let Some(s) = unsafe { read_pointer_wchar_slice(ptr_value, start, len, step) }
+            {
+                return Ok(vm.ctx.new_str(s).into());
             }
-            return Ok(vm
-                .ctx
-                .new_str(unsafe { read_pointer_wchar_slice(ptr_value, start, len, step) })
-                .into());
+            return Ok(vm.ctx.new_str("").into());
         }
 
         // other types → list with Pointer_item for each