diff --git a/cuda-core/src/cudarc_shim.rs b/cuda-core/src/cudarc_shim.rs index 1a0e63c5b..17b757ea0 100644 --- a/cuda-core/src/cudarc_shim.rs +++ b/cuda-core/src/cudarc_shim.rs @@ -362,13 +362,12 @@ pub(crate) mod stream { pub(crate) mod module { use super::{DriverError, IntoResult}; use core::ffi::c_void; - use std::ffi::CString; + use std::ffi::CStr; use std::mem::MaybeUninit; /// Loads a CUDA module from a file path. - pub fn load(filename: &str) -> Result { - let c_str = CString::new(filename).unwrap(); - let fname_ptr = c_str.as_c_str().as_ptr(); + pub fn load(filename: &CStr) -> Result { + let fname_ptr = filename.as_ptr(); let mut module = MaybeUninit::uninit(); unsafe { cuda_bindings::cuModuleLoad(module.as_mut_ptr(), fname_ptr).result()?; @@ -382,7 +381,7 @@ pub(crate) mod module { /// The PTX source must be valid. pub unsafe fn load_ptx_str(src_str: &str) -> Result { let mut module = MaybeUninit::uninit(); - let c_str = CString::new(src_str).unwrap(); + let c_str = crate::error::c_string(src_str)?; let module_res = cuda_bindings::cuModuleLoadData(module.as_mut_ptr(), c_str.as_ptr() as *const _); (module_res, module).result() @@ -406,7 +405,7 @@ pub(crate) mod module { module: cuda_bindings::CUmodule, name: &str, ) -> Result { - let name = CString::new(name).unwrap(); + let name = crate::error::c_string(name)?; let name_ptr = name.as_c_str().as_ptr(); let mut func = MaybeUninit::uninit(); let res = cuda_bindings::cuModuleGetFunction(func.as_mut_ptr(), module, name_ptr); @@ -809,3 +808,26 @@ pub(crate) mod memory { Ok((free, total)) } } + +#[cfg(test)] +mod tests { + use super::module; + + const INVALID_VALUE: cuda_bindings::CUresult = + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE; + + /// The name is rejected before any driver call. The test needs no device + /// or context. `error::c_string` has the full table of inputs. + #[test] + fn module_helpers_reject_interior_nul_without_panicking() { + let input = "a\0b"; + + let ptx = unsafe { module::load_ptx_str(input) } + .expect_err("load_ptx_str accepted an interior NUL"); + assert_eq!(ptx.0, INVALID_VALUE, "load_ptx_str({input:?})"); + + let func = unsafe { module::get_function(std::ptr::null_mut(), input) } + .expect_err("get_function accepted an interior NUL"); + assert_eq!(func.0, INVALID_VALUE, "get_function({input:?})"); + } +} diff --git a/cuda-core/src/error.rs b/cuda-core/src/error.rs index 8dd21f3fd..6610bb528 100644 --- a/cuda-core/src/error.rs +++ b/cuda-core/src/error.rs @@ -5,7 +5,7 @@ //! CUDA driver error types and result conversion utilities. -use std::ffi::CStr; +use std::ffi::{CStr, CString}; use std::mem::MaybeUninit; use std::{ error, @@ -146,9 +146,46 @@ impl DriverError { } } +/// Converts a Rust string into a NUL-terminated C string for a driver call. +/// +/// A single NUL in the final byte is a terminator and is dropped. PTX read +/// back from NVRTC ends in one. Any other NUL byte would truncate the string +/// at the driver and maps to `CUDA_ERROR_INVALID_VALUE`. The driver returns +/// that code for a bad argument. +pub(crate) fn c_string(s: &str) -> Result { + let s = s.strip_suffix('\0').unwrap_or(s); + CString::new(s).map_err(|_| DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE)) +} + #[cfg(test)] mod tests { - use super::DriverError; + use super::{c_string, DriverError}; + + #[test] + fn c_string_drops_a_final_nul_and_rejects_any_other() { + let cases: [(&str, Option<&[u8]>); 9] = [ + ("", Some(b"")), + ("kernel", Some(b"kernel")), + ("/tmp/m.ptx", Some(b"/tmp/m.ptx")), + ("\0", Some(b"")), + ("trailing\0", Some(b"trailing")), + ("a\0b", None), + ("\0a", None), + ("a\0\0", None), + ("a\0b\0", None), + ]; + for (input, expected) in cases { + match (c_string(input), expected) { + (Ok(c), Some(bytes)) => assert_eq!(c.as_bytes(), bytes, "input {input:?}"), + (Err(e), None) => assert_eq!( + e.0, + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + "input {input:?}" + ), + (got, want) => panic!("input {input:?}: got {got:?}, expected {want:?}"), + } + } + } #[test] fn identifies_unsupported_ptx_version() { diff --git a/cuda-core/src/runtime.rs b/cuda-core/src/runtime.rs index b1893e92b..312786a16 100644 --- a/cuda-core/src/runtime.rs +++ b/cuda-core/src/runtime.rs @@ -9,7 +9,7 @@ //! handles with RAII lifetimes and provide `borrow_raw` constructors //! for interop with external frameworks (cudarc, etc.). -use std::ffi::{c_int, c_void, CString}; +use std::ffi::{c_int, c_void}; use std::sync::Arc; use crate::cudarc_shim::{ctx, device, module, pool, primary_ctx, stream}; @@ -299,11 +299,9 @@ impl Device { self: &Arc, ptx_src: &str, ) -> Result, DriverError> { + let c_src = crate::error::c_string(ptx_src)?; self.bind_to_thread()?; - let cu_module = { - let c_src = CString::new(ptx_src).unwrap(); - unsafe { module::load_data(c_src.as_ptr() as *const _) } - }?; + let cu_module = unsafe { module::load_data(c_src.as_ptr() as *const _) }?; Ok(Arc::new(Module { cu_module, device: self.clone(), @@ -316,8 +314,9 @@ impl Device { self: &Arc, filename: &str, ) -> Result, DriverError> { + let c_filename = crate::error::c_string(filename)?; self.bind_to_thread()?; - let cu_module = { module::load(filename) }?; + let cu_module = module::load(&c_filename)?; Ok(Arc::new(Module { cu_module, device: self.clone(), diff --git a/cuda-core/src/simt/module.rs b/cuda-core/src/simt/module.rs index a2b4cb646..a20f6be3c 100644 --- a/cuda-core/src/simt/module.rs +++ b/cuda-core/src/simt/module.rs @@ -30,7 +30,7 @@ use crate::error::{DriverError, IntoResult}; use crate::simt::context::CudaContext; use std::borrow::Cow; -use std::ffi::{c_void, CString}; +use std::ffi::c_void; use std::mem::MaybeUninit; use std::sync::Arc; @@ -71,18 +71,18 @@ impl CudaContext { /// JIT-compiles PTX source and loads the resulting module into this /// context. /// - /// `ptx_src` must be a valid, null-terminator-free PTX string. The driver + /// `ptx_src` is PTX source that may end in one terminating NUL. The driver /// performs JIT compilation targeting the current device architecture. /// - /// # Panics + /// # Errors /// - /// Panics if `ptx_src` contains interior null bytes. + /// Returns `CUDA_ERROR_INVALID_VALUE` if `ptx_src` contains an interior null byte. pub fn load_module_from_ptx_src( self: &Arc, ptx_src: &str, ) -> Result, DriverError> { + let c_src = crate::error::c_string(ptx_src)?; self.bind_to_thread()?; - let c_src = CString::new(ptx_src).unwrap(); let cu_module = unsafe { let mut cu_module = MaybeUninit::uninit(); cuda_bindings::cuModuleLoadData(cu_module.as_mut_ptr(), c_src.as_ptr() as *const _) @@ -124,15 +124,15 @@ impl CudaContext { /// `filename` is the filesystem path. The driver selects the loader based /// on file contents (PTX text or cubin ELF). /// - /// # Panics + /// # Errors /// - /// Panics if `filename` contains interior null bytes. + /// Returns `CUDA_ERROR_INVALID_VALUE` if `filename` contains an interior null byte. pub fn load_module_from_file( self: &Arc, filename: &str, ) -> Result, DriverError> { + let c_str = crate::error::c_string(filename)?; self.bind_to_thread()?; - let c_str = CString::new(filename).unwrap(); let mut cu_module = MaybeUninit::uninit(); let cu_module = unsafe { cuda_bindings::cuModuleLoad(cu_module.as_mut_ptr(), c_str.as_ptr()).result()?; @@ -235,12 +235,12 @@ impl CudaModule { /// Returns an error if binding the module's context fails or if /// `cuModuleGetFunction` cannot resolve `fn_name` in this module. /// - /// # Panics + /// # Errors /// - /// Panics if `fn_name` contains interior null bytes. + /// Returns `CUDA_ERROR_INVALID_VALUE` if `fn_name` contains an interior null byte. pub fn load_function(self: &Arc, fn_name: &str) -> Result { + let c_name = crate::error::c_string(fn_name)?; self.ctx.bind_to_thread()?; - let c_name = CString::new(fn_name).unwrap(); let cu_function = unsafe { let mut cu_function = MaybeUninit::uninit(); cuda_bindings::cuModuleGetFunction( @@ -395,15 +395,15 @@ impl CudaModule { /// Returns an error if binding fails or if `name` cannot be resolved /// in this module. /// - /// # Panics + /// # Errors /// - /// Panics if `name` contains interior null bytes. + /// Returns `CUDA_ERROR_INVALID_VALUE` if `name` contains an interior null byte. pub fn get_global( self: &Arc, name: &str, ) -> Result<(cuda_bindings::CUdeviceptr, usize), DriverError> { + let c_name = crate::error::c_string(name)?; self.ctx.bind_to_thread()?; - let c_name = CString::new(name).unwrap(); let mut dptr = MaybeUninit::::uninit(); let mut size = MaybeUninit::::uninit(); unsafe { diff --git a/cuda-core/tests/simt_module_interior_nul.rs b/cuda-core/tests/simt_module_interior_nul.rs new file mode 100644 index 000000000..f6de53d7d --- /dev/null +++ b/cuda-core/tests/simt_module_interior_nul.rs @@ -0,0 +1,86 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Module-loading APIs return `CUDA_ERROR_INVALID_VALUE` for a name or source +//! with an interior NUL byte. A single NUL in the final byte is a terminator +//! and loads. + +use cuda_core::CudaContext; + +const INVALID_VALUE: cuda_bindings::CUresult = + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE; + +const VALID_PTX: &str = "\ +.version 7.0 +.target sm_50 +.address_size 64 +.visible .global .align 4 .u32 g; +.visible .entry k() { ret; } +"; + +const BAD_STRINGS: [&str; 2] = ["a\0b", "a\0\0"]; + +#[test] +fn load_module_from_ptx_src_rejects_interior_nul() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + for input in BAD_STRINGS { + let err = ctx + .load_module_from_ptx_src(input) + .err() + .unwrap_or_else(|| panic!("accepted {input:?}")); + assert_eq!(err.0, INVALID_VALUE, "ptx_src {input:?}"); + } +} + +#[test] +fn load_module_from_ptx_src_accepts_a_terminated_source() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let terminated = format!("{VALID_PTX}\0"); + ctx.load_module_from_ptx_src(&terminated) + .expect("PTX with its terminator failed to load"); +} + +#[test] +fn load_module_from_file_rejects_interior_nul() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + for input in BAD_STRINGS { + let err = ctx + .load_module_from_file(input) + .err() + .unwrap_or_else(|| panic!("accepted {input:?}")); + assert_eq!(err.0, INVALID_VALUE, "filename {input:?}"); + } +} + +#[test] +fn function_and_global_lookup_reject_interior_nul() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let module = ctx + .load_module_from_ptx_src(VALID_PTX) + .expect("valid PTX failed to load"); + + // Both names resolve with and without a terminator. Each rejection below + // therefore comes from the NUL check. + for name in ["k", "k\0"] { + module.load_function(name).expect("entry `k` not found"); + } + for name in ["g", "g\0"] { + module.get_global(name).expect("global `g` not found"); + } + + for input in BAD_STRINGS { + let err = module + .load_function(input) + .err() + .unwrap_or_else(|| panic!("load_function accepted {input:?}")); + assert_eq!(err.0, INVALID_VALUE, "load_function {input:?}"); + + let err = module + .get_global(input) + .err() + .unwrap_or_else(|| panic!("get_global accepted {input:?}")); + assert_eq!(err.0, INVALID_VALUE, "get_global {input:?}"); + } +}