Noeda / rllama

Rust+OpenCL+AVX2 implementation of LLaMA inference code
GNU Affero General Public License v3.0
539 stars 29 forks source link

Potential Undefined Behavior in data_u16_to_gpu Due to Unsafe Pointer Usage #19

Open lwz23 opened 3 days ago

lwz23 commented 3 days ago

Description The data_u16_to_gpu function contains unsafe code that can lead to undefined behavior (UB) due to unchecked pointer usage and assumptions about input validity. These issues undermine the safety guarantees of Rust and can lead to runtime crashes or memory corruption. https://github.com/Noeda/rllama/blob/1e1131faaaf7013ed19639ad96f252458efdb45b/src/tensor_opencl_support.rs#L149 Code in Question:

pub fn data_u16_to_gpu(
    &self,
    data: *const u16,
    data_layout: Layout,
    nitems: usize,
    rows: i64,
    cols: i64,
    cols_capacity: i64,
) -> Result<OpenCLTensor, OpenCLError> {
    unsafe {
        let buf = Buffer::builder()
            .queue(self.queue.clone())
            .len(nitems)
            .build()?;
        let mut event = Event::empty();
        let data_slice: &[u16] = std::slice::from_raw_parts(data, nitems);
        buf.cmd()
            .write(data_slice)
            .block(false)
            .enew(&mut event)
            .enq()?;
        Ok(OpenCLTensor {
            buf,
            initial_write_event: Some(event),
            last_event: None,
            data,
            data_layout,
            nitems,
            rows,
            cols,
            cols_capacity,
            queue: self.queue.clone(),
            cl: self.clone(),
        })
    }
}

Problems: this function is 'pub' so the user can control the 'data' field, if the 'data' is eg. a null pointer, there will be a UB

Steps to Reproduce: Here is an example that might lead to UB:

let invalid_data = std::ptr::null(); // Null pointer
let nitems = 100;
let layout = Layout::from_size_align(100 * std::mem::size_of::<u16>(), 2).unwrap();

let result = gpu_context.data_u16_to_gpu(
    invalid_data,
    layout,
    nitems,
    10, // rows
    10, // cols
    10, // cols_capacity
);

assert!(result.is_err()); // Undefined behavior may occur before reaching this point

Suggestion

  1. mark this function as 'unsafe'.
  2. add 'assert' in the function body.

Expected Behavior: Validate that data is non-null, properly aligned, and points to valid memory.

Additional Context: This function operates in an unsafe block, and mistakes in pointer handling can lead to serious consequences, such as crashes, data corruption, or security vulnerabilities. Input validation and robust error handling are essential to ensure soundness in Rust.

lwz23 commented 1 day ago

Unmatained?