Use memchr consistently for NUL checks

`memchr` is used throughout RustPython for searching through bytes
expediently. However, for NUL checks, it's scantily used. Instead, our
NUL checks either use `contains` or `memchr`.

I switched all of the `contains(b'\0')` I could find to using memchr
instead. I marked the failure paths as cold to hint to LLVM that
interior NULs are truly exceptional. This should help branch prediction
a bit which is nice for string/bytes functions since they are likely
called a lot.

Author: joshuamegnauth54

crates/host_env/src/multiprocessing.rs

@@ -32,12 +32,13 @@ pub enum SemError {
     AlreadyExists,
     NotFound,
     InvalidInput,
+    InteriorNul,
     Other(i32),
 }
 
 #[cfg(unix)]
 impl SemError {
-    fn from_errno(err: Errno) -> Self {
+    const fn from_errno(err: Errno) -> Self {
         match err {
             Errno::EAGAIN => Self::WouldBlock,
             Errno::ETIMEDOUT => Self::TimedOut,
@@ -49,14 +50,14 @@ impl SemError {
         }
     }
 
-    pub fn raw_os_error(self) -> i32 {
+    pub const fn raw_os_error(self) -> i32 {
         match self {
             Self::WouldBlock => Errno::EAGAIN as i32,
             Self::TimedOut => Errno::ETIMEDOUT as i32,
             Self::Interrupted => Errno::EINTR as i32,
             Self::AlreadyExists => Errno::EEXIST as i32,
             Self::NotFound => Errno::ENOENT as i32,
-            Self::InvalidInput => Errno::EINVAL as i32,
+            Self::InvalidInput | Self::InteriorNul => Errno::EINVAL as i32,
             Self::Other(code) => code,
         }
     }
@@ -119,7 +120,7 @@ impl SemHandle {
         value: u32,
         unlink: bool,
     ) -> Result<(Self, Option<String>), SemError> {
-        let cname = semaphore_name(name).map_err(|_| SemError::InvalidInput)?;
+        let cname = semaphore_name(name)?;
         let raw =
             unsafe { libc::sem_open(cname.as_ptr(), libc::O_CREAT | libc::O_EXCL, 0o600, value) };
         if raw == libc::SEM_FAILED {
@@ -141,7 +142,7 @@ impl SemHandle {
     }
 
     pub fn open_existing(name: &str) -> Result<Self, SemError> {
-        let cname = semaphore_name(name).map_err(|_| SemError::InvalidInput)?;
+        let cname = semaphore_name(name)?;
         let raw = unsafe { libc::sem_open(cname.as_ptr(), 0) };
         if raw == libc::SEM_FAILED {
             Err(SemError::from_errno(Errno::last()))
@@ -305,18 +306,18 @@ pub fn is_too_many_posts(err: u32) -> bool {
 }
 
 #[cfg(unix)]
-pub fn semaphore_name(name: &str) -> Result<CString, alloc::ffi::NulError> {
-    let mut full = String::with_capacity(name.len() + 1);
+pub fn semaphore_name(name: &str) -> Result<CString, SemError> {
+    let mut full = String::with_capacity(name.len() + 2);
     if !name.starts_with('/') {
         full.push('/');
     }
     full.push_str(name);
-    CString::new(full)
+    CString::new(full).map_err(|_| SemError::InteriorNul)
 }
 
 #[cfg(unix)]
 pub fn sem_unlink(name: &str) -> Result<(), SemError> {
-    let cname = semaphore_name(name).map_err(|_| SemError::InvalidInput)?;
+    let cname = semaphore_name(name)?;
     let res = unsafe { libc::sem_unlink(cname.as_ptr()) };
     if res < 0 {
         Err(SemError::from_errno(Errno::last()))

crates/host_env/src/time.rs

@@ -618,10 +618,9 @@ unsafe extern "C" {
 
 #[cfg(windows)]
 pub fn strftime_ascii(fmt: &str, tm: &libc::tm) -> Result<String, CheckedTmError> {
-    if fmt.contains('\0') {
-        return Err(CheckedTmError::EmbeddedNul);
-    }
-    let fmt_wide: Vec<u16> = fmt.encode_utf16().chain(core::iter::once(0)).collect();
+    let fmt_wide = widestring::WideCString::from_str(fmt)
+        .map_err(|_| CheckedTmError::EmbeddedNul)?
+        .into_vec_with_nul();
     let mut size = 1024usize;
     let max_scale = 256usize.saturating_mul(fmt.len().max(1));
     loop {

crates/host_env/src/winapi.rs

@@ -3,22 +3,20 @@
     reason = "This module mirrors Win32 APIs with raw handle and pointer parameters."
 )]
 
+use core::hint::cold_path;
 use std::{io, path::Path};
-use windows_sys::Win32::{
-    Foundation::{HANDLE, HMODULE, WAIT_FAILED},
-    System::Threading::PROCESS_INFORMATION,
-};
 
 use crate::windows::{CheckWin32Bool, CheckWin32Handle};
 
+use memchr::memchr;
 pub use windows_sys::Win32::{
     Foundation::{
         DUPLICATE_CLOSE_SOURCE, DUPLICATE_SAME_ACCESS, ERROR_ACCESS_DENIED, ERROR_ALREADY_EXISTS,
         ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_NETNAME_DELETED, ERROR_NO_DATA,
         ERROR_NO_SYSTEM_RESOURCES, ERROR_NOT_FOUND, ERROR_OPERATION_ABORTED, ERROR_PIPE_BUSY,
         ERROR_PIPE_CONNECTED, ERROR_PORT_UNREACHABLE, ERROR_PRIVILEGE_NOT_HELD, ERROR_SEM_TIMEOUT,
-        ERROR_SUCCESS, GENERIC_READ, GENERIC_WRITE, STILL_ACTIVE, WAIT_ABANDONED_0, WAIT_OBJECT_0,
-        WAIT_TIMEOUT,
+        ERROR_SUCCESS, GENERIC_READ, GENERIC_WRITE, HANDLE, HMODULE, STILL_ACTIVE,
+        WAIT_ABANDONED_0, WAIT_FAILED, WAIT_OBJECT_0, WAIT_TIMEOUT,
     },
     Globalization::{
         LCMAP_FULLWIDTH, LCMAP_HALFWIDTH, LCMAP_HIRAGANA, LCMAP_KATAKANA, LCMAP_LINGUISTIC_CASING,
@@ -57,12 +55,12 @@ pub use windows_sys::Win32::{
             ABOVE_NORMAL_PRIORITY_CLASS, BELOW_NORMAL_PRIORITY_CLASS, CREATE_BREAKAWAY_FROM_JOB,
             CREATE_DEFAULT_ERROR_MODE, CREATE_NEW_CONSOLE, CREATE_NEW_PROCESS_GROUP,
             CREATE_NO_WINDOW, DETACHED_PROCESS, HIGH_PRIORITY_CLASS, IDLE_PRIORITY_CLASS,
-            NORMAL_PRIORITY_CLASS, PROCESS_ALL_ACCESS, PROCESS_DUP_HANDLE, REALTIME_PRIORITY_CLASS,
-            STARTF_FORCEOFFFEEDBACK, STARTF_FORCEONFEEDBACK, STARTF_PREVENTPINNING,
-            STARTF_RUNFULLSCREEN, STARTF_TITLEISAPPID, STARTF_TITLEISLINKNAME,
-            STARTF_UNTRUSTEDSOURCE, STARTF_USECOUNTCHARS, STARTF_USEFILLATTRIBUTE,
-            STARTF_USEHOTKEY, STARTF_USEPOSITION, STARTF_USESHOWWINDOW, STARTF_USESIZE,
-            STARTF_USESTDHANDLES,
+            NORMAL_PRIORITY_CLASS, PROCESS_ALL_ACCESS, PROCESS_DUP_HANDLE, PROCESS_INFORMATION,
+            REALTIME_PRIORITY_CLASS, STARTF_FORCEOFFFEEDBACK, STARTF_FORCEONFEEDBACK,
+            STARTF_PREVENTPINNING, STARTF_RUNFULLSCREEN, STARTF_TITLEISAPPID,
+            STARTF_TITLEISLINKNAME, STARTF_UNTRUSTEDSOURCE, STARTF_USECOUNTCHARS,
+            STARTF_USEFILLATTRIBUTE, STARTF_USEHOTKEY, STARTF_USEPOSITION, STARTF_USESHOWWINDOW,
+            STARTF_USESIZE, STARTF_USESTDHANDLES,
         },
     },
     UI::WindowsAndMessaging::SW_HIDE,
@@ -312,7 +310,8 @@ pub fn build_environment_block(
 
     let mut last_entry: HashMap<String, Vec<u16>> = HashMap::new();
     for (key, value) in entries {
-        if key.contains('\0') || value.contains('\0') {
+        if memchr(b'\0', key.as_bytes()) || memchr(b'\0', value.as_bytes()) {
+            cold_path();
             return Err(BuildEnvironmentBlockError::ContainsNul);
         }
         if key.is_empty() || key[1..].contains('=') {

crates/stdlib/src/grp.rs

@@ -10,6 +10,7 @@ mod grp {
         exceptions,
         types::PyStructSequence,
     };
+    use core::hint::cold_path;
     use rustpython_host_env::grp as host_grp;
 
     #[pystruct_sequence_data]
@@ -61,10 +62,11 @@ mod grp {
 
     #[pyfunction]
     fn getgrnam(name: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult<GroupData> {
-        let gr_name = name.as_str();
-        if gr_name.contains('\0') {
+        if name.as_pystr().contains_nuls() {
+            cold_path();
             return Err(exceptions::nul_char_error(vm));
         }
+        let gr_name = name.as_str();
         let group = host_grp::getgrnam(gr_name).map_err(|err| err.into_pyexception(vm))?;
         let group = group.ok_or_else(|| {
             vm.new_key_error(

crates/stdlib/src/mmap.rs

@@ -24,9 +24,9 @@ mod mmap {
     use core::ops::{Deref, DerefMut};
     use crossbeam_utils::atomic::AtomicCell;
     use num_traits::Signed;
-    #[cfg(windows)]
-    use std::io;
     use std::io::Write;
+    #[cfg(windows)]
+    use {core::hint::cold_path, memchr::memchr, rustpython_vm::exceptions, std::io};
 
     #[cfg(unix)]
     use rustpython_host_env::crt_fd;
@@ -460,8 +460,9 @@ mod mmap {
                     let s = obj
                         .try_to_value::<String>(vm)
                         .map_err(|_| vm.new_type_error("tagname must be a string or None"))?;
-                    if s.contains('\0') {
-                        return Err(vm.new_value_error("tagname must not contain null characters"));
+                    if memchr(b'\0', s.as_bytes()).is_some() {
+                        cold_path();
+                        return Err(exceptions::nul_char_error(vm));
                     }
                     Some(s)
                 }

crates/stdlib/src/multiprocessing.rs

@@ -10,7 +10,10 @@ mod _multiprocessing {
         function::{ArgBytesLike, FuncArgs, KwArgs},
         types::Constructor,
     };
-    use core::sync::atomic::{AtomicI32, AtomicU32, Ordering};
+    use core::{
+        hint::cold_path,
+        sync::atomic::{AtomicI32, AtomicU32, Ordering},
+    };
     use rustpython_host_env::multiprocessing as host_multiprocessing;
 
     // These match the values in Lib/multiprocessing/synchronize.py
@@ -811,7 +814,7 @@ mod _multiprocessing {
             let value = args.value as u32;
             let (handle, name) =
                 SemHandle::create(&args.name, value, args.unlink).map_err(|err| {
-                    if err == SemError::InvalidInput && args.name.contains('\0') {
+                    if err == SemError::InteriorNul {
                         exceptions::nul_char_error(vm)
                     } else {
                         os_error(vm, err)
@@ -835,7 +838,7 @@ mod _multiprocessing {
     #[pyfunction]
     fn sem_unlink(name: String, vm: &VirtualMachine) -> PyResult<()> {
         host_multiprocessing::sem_unlink(&name).map_err(|err| {
-            if err == SemError::InvalidInput && name.contains('\0') {
+            if err == SemError::InteriorNul {
                 exceptions::nul_char_error(vm)
             } else {
                 os_error(vm, err)

crates/stdlib/src/openssl.rs

@@ -53,6 +53,8 @@ fn probe() -> &'static ProbeResult {
     #[cfg(ossl111)] ossl111,
     #[cfg(windows)] windows))]
 mod _ssl {
+    use core::hint::cold_path;
+
     use super::{bio, probe};
 
     // Import error types and helpers used in this module (others are exposed via pymodule(with(...)))
@@ -85,6 +87,7 @@ mod _ssl {
     };
     use crossbeam_utils::atomic::AtomicCell;
     use foreign_types_shared::{ForeignType, ForeignTypeRef};
+    use memchr::memchr;
     use openssl::{
         asn1::{Asn1Object, Asn1ObjectRef},
         error::ErrorStack,
@@ -1039,12 +1042,13 @@ mod _ssl {
 
         #[pymethod]
         fn set_ciphers(&self, cipherlist: PyStrRef, vm: &VirtualMachine) -> PyResult<()> {
-            let ciphers: &str = cipherlist.as_ref();
-            if ciphers.contains('\0') {
+            if cipherlist.contains_nuls() {
+                cold_path();
                 return Err(exceptions::nul_char_error(vm));
             }
+
             self.builder()
-                .set_cipher_list(ciphers)
+                .set_cipher_list(cipherlist.as_ref())
                 .map_err(|_| new_ssl_error(vm, "No cipher can be selected."))
         }
 
@@ -1096,9 +1100,6 @@ mod _ssl {
             let name_cstr = match name {
                 Either::A(s) => {
                     let s: &str = s.as_ref();
-                    if s.contains('\0') {
-                        return Err(exceptions::nul_char_error(vm));
-                    }
                     s.to_cstring(vm)?
                 }
                 Either::B(b) => std::ffi::CString::new(b.borrow_buf().to_vec())
@@ -2031,15 +2032,16 @@ mod _ssl {
 
             // Configure server hostname
             if let Some(hostname) = &server_hostname {
+                if hostname.contains_nuls() {
+                    cold_path();
+                    return Err(exceptions::nul_char_type_error(vm));
+                }
                 let hostname_str: &str = hostname.as_ref();
                 if hostname_str.is_empty() || hostname_str.starts_with('.') {
                     return Err(vm.new_value_error(
                         "server_hostname cannot be an empty string or start with a leading dot.",
                     ));
                 }
-                if hostname_str.contains('\0') {
-                    return Err(exceptions::nul_char_type_error(vm));
-                }
                 let ip = hostname_str.parse::<core::net::IpAddr>();
                 if ip.is_err() {
                     ssl.set_hostname(hostname_str)

crates/stdlib/src/ssl.rs

@@ -67,9 +67,11 @@ mod _ssl {
     use alloc::sync::Arc;
     use core::{
         hash::{Hash, Hasher},
+        hint::cold_path,
         sync::atomic::{AtomicUsize, Ordering},
         time::Duration,
     };
+    use memchr::memchr;
     use rustpython_vm::exceptions;
     use std::{
         collections::{HashMap, hash_map::DefaultHasher},
@@ -392,7 +394,8 @@ mod _ssl {
         // IP addresses are allowed as server_hostname
         // SNI will not be sent for IP addresses
 
-        if hostname.contains('\0') {
+        if memchr(b'\0', hostname.as_bytes()).is_some() {
+            cold_path();
             return Err(exceptions::nul_char_type_error(vm));
         }
 
@@ -1854,25 +1857,7 @@ mod _ssl {
             let hostname = match args.server_hostname.into_option().flatten() {
                 Some(hostname_str) => {
                     let hostname = hostname_str.as_str();
-
-                    // Validate hostname
-                    if hostname.is_empty() {
-                        return Err(vm.new_value_error("server_hostname cannot be an empty string"));
-                    }
-
-                    // Check if it starts with a dot
-                    if hostname.starts_with('.') {
-                        return Err(vm.new_value_error("server_hostname cannot start with a dot"));
-                    }
-
-                    // IP addresses are allowed
-                    // SNI will not be sent for IP addresses
-
-                    // Check for NULL bytes
-                    if hostname.contains('\0') {
-                        return Err(exceptions::nul_char_error(vm));
-                    }
-
+                    validate_hostname(hostname, vm)?;
                     Some(hostname.to_string())
                 }
                 None => None,

crates/vm/src/builtins/bytes.rs

@@ -31,6 +31,7 @@ use crate::{
 };
 use bstr::ByteSlice;
 use core::{mem::size_of, ops::Deref};
+use memchr::memchr;
 
 #[pyclass(module = false, name = "bytes")]
 #[derive(Clone, Debug)]
@@ -169,6 +170,13 @@ impl PyBytes {
                 .map(|x| vm.ctx.new_bytes(x).into()),
         }
     }
+
+    /// Check bytes for interior NULs.
+    #[inline]
+    #[must_use]
+    pub fn contains_nuls(&self) -> bool {
+        memchr(b'\0', self.as_bytes()).is_some()
+    }
 }
 
 impl PyRef<PyBytes> {
@@ -218,7 +226,7 @@ impl PyBytes {
 
     #[inline]
     #[must_use]
-    pub fn as_bytes(&self) -> &[u8] {
+    pub const fn as_bytes(&self) -> &[u8] {
         self.inner.as_bytes()
     }
 

crates/vm/src/builtins/str.rs

@@ -37,6 +37,7 @@ use bstr::ByteSlice;
 use core::ffi::CStr;
 use core::{char, mem, ops::Range};
 use itertools::Itertools;
+use memchr::memchr;
 use num_traits::ToPrimitive;
 use rustpython_common::{
     ascii,
@@ -545,6 +546,13 @@ impl PyStr {
         }
     }
 
+    /// Check string bytes for interior NULs.
+    #[inline]
+    #[must_use]
+    pub fn contains_nuls(&self) -> bool {
+        memchr(b'\0', self.as_bytes()).is_some()
+    }
+
     pub fn to_string_lossy(&self) -> Cow<'_, str> {
         self.to_str()
             .map_or_else(|| self.as_wtf8().to_string_lossy(), Cow::Borrowed)
@@ -2150,7 +2158,7 @@ impl PyUtf8Str {
 
 impl Py<PyUtf8Str> {
     /// Upcast to PyStr.
-    pub fn as_pystr(&self) -> &Py<PyStr> {
+    pub const fn as_pystr(&self) -> &Py<PyStr> {
         unsafe {
             // Safety: PyUtf8Str is a wrapper around PyStr, so this cast is safe.
             &*(self as *const Self as *const Py<PyStr>)