Skip to content

Commit bae642c

Browse files
joshuamegnauth54joshuamegnauth54
authored
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.
1 parent a49e8f9 commit bae642c

4 files changed

Lines changed: 80 additions & 105 deletions

File tree

‎crates/host_env/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ libc = { workspace = true }
1717
num-traits = { workspace = true }
1818
parking_lot = { workspace = true }
1919
paste = { workspace = true }
20+
widestring = { workspace = true }
2021

2122
[target.'cfg(unix)'.dependencies]
2223
nix = { workspace = true }
@@ -54,7 +55,6 @@ system-configuration = { workspace = true }
5455
memchr.workspace = true
5556
junction = { workspace = true }
5657
schannel = { workspace = true }
57-
widestring = { workspace = true }
5858
windows-sys = { workspace = true, features = [
5959
"Win32_Foundation",
6060
"Win32_Globalization",

‎crates/host_env/src/ctypes.rs‎

Lines changed: 25 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use core::ffi::{
33
CStr, c_char, c_double, c_float, c_int, c_long, c_longlong, c_schar, c_short, c_uchar, c_uint,
44
c_ulong, c_ulonglong, c_ushort, c_void,
55
};
6+
use core::ptr::NonNull;
67
#[cfg(all(
78
any(
89
target_os = "linux",
@@ -36,6 +37,7 @@ use rustpython_wtf8::Wtf8;
3637
use rustpython_wtf8::Wtf8Buf;
3738
#[cfg(any(unix, windows))]
3839
use std::{collections::HashMap, ffi::OsStr, sync::OnceLock};
40+
use widestring::WideCStr;
3941

4042
#[cfg(all(
4143
any(
@@ -392,33 +394,11 @@ pub fn dyld_shared_cache_contains_path(path: &str) -> Result<bool, alloc::ffi::N
392394
Ok(unsafe { _dyld_shared_cache_contains_path(c_path.as_ptr()) })
393395
}
394396

395-
/// # Safety
396-
///
397-
/// `ptr` must be valid to read until the first NUL byte.
398-
pub unsafe fn strlen(ptr: *const c_char) -> usize {
399-
#[cfg(any(unix, windows, target_os = "wasi"))]
400-
{
401-
unsafe { libc::strlen(ptr) }
402-
}
403-
#[cfg(not(any(unix, windows, target_os = "wasi")))]
404-
{
405-
let mut len = 0;
406-
while unsafe { *ptr.add(len) } != 0 {
407-
len += 1;
408-
}
409-
len
410-
}
411-
}
412-
413397
/// # Safety
414398
///
415399
/// `ptr` must be valid to read until the first NUL wide character.
416-
pub unsafe fn wcslen(ptr: *const WChar) -> usize {
417-
let mut len = 0;
418-
while unsafe { *ptr.add(len) } != 0 as WChar {
419-
len += 1;
420-
}
421-
len
400+
pub unsafe fn wcslen(ptr: NonNull<WChar>) -> usize {
401+
unsafe { WideCStr::from_ptr_str(ptr.as_ptr().cast()).len() }
422402
}
423403

424404
/// # Safety
@@ -1309,10 +1289,10 @@ pub unsafe fn callback_arg_value(type_code: Option<&str>, ptr: *const c_void) ->
13091289
}
13101290
Some("Z") => {
13111291
let wstr_ptr = unsafe { *(ptr as *const *const WChar) };
1312-
if wstr_ptr.is_null() {
1313-
DecodedValue::None
1314-
} else {
1292+
if let Some(wstr_ptr) = NonNull::new(wstr_ptr.cast_mut()) {
13151293
DecodedValue::String(unsafe { read_wide_string(wstr_ptr) }.to_string())
1294+
} else {
1295+
DecodedValue::None
13161296
}
13171297
}
13181298
Some("P") => DecodedValue::Pointer(unsafe { *(ptr as *const usize) }),
@@ -2272,8 +2252,7 @@ pub unsafe fn borrowed_slice_as_mut(slice: &[u8]) -> &mut [u8] {
22722252
pub fn wide_chars_to_wtf8(wchars: &[WChar]) -> Wtf8Buf {
22732253
#[cfg(windows)]
22742254
{
2275-
let wide: Vec<u16> = wchars.to_vec();
2276-
Wtf8Buf::from_wide(&wide)
2255+
Wtf8Buf::from_wide(wchars)
22772256
}
22782257
#[cfg(not(windows))]
22792258
{
@@ -2292,10 +2271,10 @@ pub fn wide_chars_to_wtf8(wchars: &[WChar]) -> Wtf8Buf {
22922271
/// # Safety
22932272
///
22942273
/// `ptr` must be a valid NUL-terminated wide C string.
2295-
pub unsafe fn read_wide_string(ptr: *const WChar) -> Wtf8Buf {
2296-
let len = unsafe { wcslen(ptr) };
2297-
let wchars = unsafe { core::slice::from_raw_parts(ptr, len) };
2298-
wide_chars_to_wtf8(wchars)
2274+
pub unsafe fn read_wide_string(ptr: NonNull<WChar>) -> Wtf8Buf {
2275+
// SAFETY: WideCStr does not assume an encoding.
2276+
let wchars = unsafe { WideCStr::from_ptr_str(ptr.as_ptr().cast()) };
2277+
Wtf8Buf::from_string(wchars.to_string_lossy())
22992278
}
23002279

23012280
/// # Safety
@@ -2313,18 +2292,15 @@ pub unsafe fn read_c_string_from_address(addr: usize) -> Option<Vec<u8>> {
23132292
///
23142293
/// `addr` must either be zero or a valid NUL-terminated wide C string pointer.
23152294
pub unsafe fn read_wide_string_from_address(addr: usize) -> Option<Wtf8Buf> {
2316-
if addr == 0 {
2317-
None
2318-
} else {
2319-
Some(unsafe { read_wide_string(addr as *const WChar) })
2320-
}
2295+
let ptr = NonNull::new(addr as *mut WChar)?;
2296+
Some(unsafe { read_wide_string(ptr) })
23212297
}
23222298

23232299
/// # Safety
23242300
///
23252301
/// `ptr` must point to `len` readable wide characters.
2326-
pub unsafe fn read_wide_string_with_len(ptr: *const WChar, len: usize) -> Wtf8Buf {
2327-
let wchars = unsafe { core::slice::from_raw_parts(ptr, len) };
2302+
pub unsafe fn read_wide_string_with_len(ptr: NonNull<WChar>, len: usize) -> Wtf8Buf {
2303+
let wchars = unsafe { core::slice::from_raw_parts(ptr.as_ptr(), len) };
23282304
wide_chars_to_wtf8(wchars)
23292305
}
23302306

@@ -2348,13 +2324,12 @@ pub fn string_at(ptr: usize, size: isize) -> Result<Vec<u8>, StringAtError> {
23482324
}
23492325

23502326
pub fn wstring_at(ptr: usize, size: isize) -> Result<Wtf8Buf, StringAtError> {
2351-
if ptr == 0 {
2327+
let Some(ptr) = NonNull::new(ptr as *mut WChar) else {
23522328
return Err(StringAtError::NullPointer);
2353-
}
2354-
let w_ptr = ptr as *const WChar;
2329+
};
23552330
if size < 0 {
23562331
// SAFETY: caller passed a non-null NUL-terminated wide string pointer.
2357-
return Ok(unsafe { read_wide_string(w_ptr) });
2332+
return Ok(unsafe { read_wide_string(ptr) });
23582333
}
23592334
let len = {
23602335
let size_usize = size as usize;
@@ -2364,7 +2339,7 @@ pub fn wstring_at(ptr: usize, size: isize) -> Result<Wtf8Buf, StringAtError> {
23642339
size_usize
23652340
};
23662341
// SAFETY: caller requested exactly `len` readable wide characters from non-null pointer.
2367-
Ok(unsafe { read_wide_string_with_len(w_ptr, len) })
2342+
Ok(unsafe { read_wide_string_with_len(ptr, len) })
23682343
}
23692344

23702345
/// # Safety
@@ -2414,14 +2389,14 @@ pub unsafe fn read_pointer_char_slice(
24142389
/// # Safety
24152390
///
24162391
/// `start` must be valid to read `len` wide characters following `step`.
2417-
pub unsafe fn read_wide_string_strided(start: *const WChar, len: usize, step: isize) -> Wtf8Buf {
2392+
pub unsafe fn read_wide_string_strided(start: NonNull<WChar>, len: usize, step: isize) -> Wtf8Buf {
24182393
if step == 1 {
24192394
return unsafe { read_wide_string_with_len(start, len) };
24202395
}
24212396
let mut wchars = Vec::with_capacity(len);
24222397
let mut cur = start;
24232398
for _ in 0..len {
2424-
wchars.push(unsafe { *cur });
2399+
wchars.push(unsafe { cur.read() });
24252400
cur = unsafe { cur.offset(step) };
24262401
}
24272402
wide_chars_to_wtf8(&wchars)
@@ -2436,10 +2411,9 @@ pub unsafe fn read_pointer_wchar_slice(
24362411
start: isize,
24372412
len: usize,
24382413
step: isize,
2439-
) -> Wtf8Buf {
2440-
let wchar_size = core::mem::size_of::<WChar>();
2441-
let start_addr = (ptr_value as isize + start * wchar_size as isize) as *const WChar;
2442-
unsafe { read_wide_string_strided(start_addr, len, step) }
2414+
) -> Option<Wtf8Buf> {
2415+
let start_addr = unsafe { NonNull::new(ptr_value as *mut WChar)?.offset(start) };
2416+
Some(unsafe { read_wide_string_strided(start_addr, len, step) })
24432417
}
24442418

24452419
/// # Safety

‎crates/host_env/src/wmi.rs‎

Lines changed: 49 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
#![allow(unsafe_op_in_unsafe_fn)]
77

88
use core::ffi::c_void;
9-
use core::ptr::{null, null_mut};
9+
use core::ptr::{NonNull, null, null_mut};
1010
use windows_sys::Win32::Foundation::{
1111
CloseHandle, ERROR_BROKEN_PIPE, ERROR_MORE_DATA, ERROR_NOT_ENOUGH_MEMORY, GetLastError, HANDLE,
1212
WAIT_OBJECT_0, WAIT_TIMEOUT,
@@ -17,6 +17,8 @@ use windows_sys::Win32::System::Threading::{
1717
CreateEventW, CreateThread, GetExitCodeThread, SetEvent, WaitForSingleObject,
1818
};
1919

20+
use crate::ctypes::wcslen;
21+
2022
pub const BUFFER_SIZE: usize = 8192;
2123

2224
pub enum ExecQueryError {
@@ -238,34 +240,26 @@ unsafe fn object_end_enumeration(this: *mut c_void) -> HRESULT {
238240
method(this)
239241
}
240242

241-
fn hresult_from_win32(err: u32) -> HRESULT {
243+
const fn hresult_from_win32(err: u32) -> HRESULT {
242244
if err == 0 {
243245
0
244246
} else {
245247
((err & 0xFFFF) | 0x80070000) as HRESULT
246248
}
247249
}
248250

249-
fn succeeded(hr: HRESULT) -> bool {
251+
const fn succeeded(hr: HRESULT) -> bool {
250252
hr >= 0
251253
}
252254

253-
fn failed(hr: HRESULT) -> bool {
255+
const fn failed(hr: HRESULT) -> bool {
254256
hr < 0
255257
}
256258

257259
fn wide_str(s: &str) -> Vec<u16> {
258260
s.encode_utf16().chain(core::iter::once(0)).collect()
259261
}
260262

261-
unsafe fn wcslen(s: *const u16) -> usize {
262-
let mut len = 0;
263-
while unsafe { *s.add(len) } != 0 {
264-
len += 1;
265-
}
266-
len
267-
}
268-
269263
unsafe fn wait_event(event: HANDLE, timeout: u32) -> u32 {
270264
match unsafe { WaitForSingleObject(event, timeout) } {
271265
WAIT_OBJECT_0 => 0,
@@ -471,16 +465,25 @@ unsafe fn query_thread_impl(param: *mut c_void) -> u32 {
471465
}
472466

473467
if succeeded(hr) && (flavor & WBEM_FLAVOR_MASK_ORIGIN) != WBEM_FLAVOR_ORIGIN_SYSTEM {
468+
let Some(cb_str1) = NonNull::new(prop_name)
469+
.map(|prop_name| (unsafe { wcslen(prop_name) } * 2) as u32)
470+
else {
471+
unsafe {
472+
SysFreeString(prop_name);
473+
}
474+
break;
475+
};
476+
474477
let mut prop_str = [0u16; BUFFER_SIZE];
475478
hr = unsafe {
476479
VariantToString(&prop_value, prop_str.as_mut_ptr(), BUFFER_SIZE as u32)
477480
};
481+
let cb_str2 = NonNull::new(prop_str.as_ptr().cast_mut())
482+
.map(|prop_str| (unsafe { wcslen(prop_str) } * 2) as u32)
483+
.expect("prop_str is never null");
478484

479-
if succeeded(hr) {
480-
let cb_str1 = (unsafe { wcslen(prop_name) } * 2) as u32;
481-
let cb_str2 = (unsafe { wcslen(prop_str.as_ptr()) } * 2) as u32;
482-
483-
if unsafe {
485+
if succeeded(hr)
486+
&& unsafe {
484487
WriteFile(
485488
write_pipe,
486489
prop_name as *const _,
@@ -489,36 +492,35 @@ unsafe fn query_thread_impl(param: *mut c_void) -> u32 {
489492
null_mut(),
490493
)
491494
} == 0
492-
|| unsafe {
493-
WriteFile(
494-
write_pipe,
495-
&eq_sign as *const u16 as *const _,
496-
2,
497-
&mut written,
498-
null_mut(),
499-
)
500-
} == 0
501-
|| unsafe {
502-
WriteFile(
503-
write_pipe,
504-
prop_str.as_ptr() as *const _,
505-
cb_str2,
506-
&mut written,
507-
null_mut(),
508-
)
509-
} == 0
510-
|| unsafe {
511-
WriteFile(
512-
write_pipe,
513-
&null_sep as *const u16 as *const _,
514-
2,
515-
&mut written,
516-
null_mut(),
517-
)
518-
} == 0
519-
{
520-
hr = hresult_from_win32(unsafe { GetLastError() });
521-
}
495+
|| unsafe {
496+
WriteFile(
497+
write_pipe,
498+
&eq_sign as *const u16 as *const _,
499+
2,
500+
&mut written,
501+
null_mut(),
502+
)
503+
} == 0
504+
|| unsafe {
505+
WriteFile(
506+
write_pipe,
507+
prop_str.as_ptr() as *const _,
508+
cb_str2,
509+
&mut written,
510+
null_mut(),
511+
)
512+
} == 0
513+
|| unsafe {
514+
WriteFile(
515+
write_pipe,
516+
&null_sep as *const u16 as *const _,
517+
2,
518+
&mut written,
519+
null_mut(),
520+
)
521+
} == 0
522+
{
523+
hr = hresult_from_win32(unsafe { GetLastError() });
522524
}
523525

524526
unsafe {

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -524,13 +524,12 @@ impl PyCPointer {
524524

525525
// c_wchar → str
526526
if type_code.as_deref() == Some("u") {
527-
if len == 0 {
528-
return Ok(vm.ctx.new_str("").into());
527+
if len > 0
528+
&& let Some(s) = unsafe { read_pointer_wchar_slice(ptr_value, start, len, step) }
529+
{
530+
return Ok(vm.ctx.new_str(s).into());
529531
}
530-
return Ok(vm
531-
.ctx
532-
.new_str(unsafe { read_pointer_wchar_slice(ptr_value, start, len, step) })
533-
.into());
532+
return Ok(vm.ctx.new_str("").into());
534533
}
535534

536535
// other types → list with Pointer_item for each

0 commit comments

Comments
 (0)