tfhe_c_api_dynamic_buffer/

lib.rs

//! This crate provides a [`DynamicBuffer`] struct that allows to easily share a pointer to u8 with
//! C APIs and free that pointer properly by carrying a `destructor_callback`. In that regard it is
//! carrying a very barebone vtable so that freeing the memory pointed to by the [`DynamicBuffer`]
//! is easy either on the C or Rust side.
//!
//! A `From` implementation is provided to convert a `Vec` of `u8` into a [`DynamicBuffer`] easily,
//! the destructor being populated automatically.
//!
//! A [`DynamicBufferView`] is also provided to indicate that the struct does not own the data and
//! is merely used to share data in a read-only way.

#![cfg_attr(all(doc, not(doctest)), feature(doc_auto_cfg))]
#![cfg_attr(all(doc, not(doctest)), feature(doc_cfg))]
#![warn(rustdoc::broken_intra_doc_links)]

use std::ffi::c_int;

#[repr(C)]
pub struct DynamicBufferView {
    pub pointer: *const u8,
    pub length: usize,
}

impl DynamicBufferView {
    /// Returns a view to the memory borrowed by the [`DynamicBufferView`].
    ///
    /// # Safety
    ///
    /// This is safe to call as long as the pointer is valid and the length corresponds to the
    /// length of the underlying buffer.
    pub unsafe fn as_slice(&self) -> &[u8] {
        std::slice::from_raw_parts(self.pointer, self.length)
    }
}

impl From<&[u8]> for DynamicBufferView {
    fn from(a: &[u8]) -> Self {
        Self {
            pointer: a.as_ptr(),
            length: a.len(),
        }
    }
}

#[repr(C)]
pub struct DynamicBuffer {
    pub pointer: *mut u8,
    pub length: usize,
    pub destructor: Option<unsafe extern "C" fn(*mut u8, usize) -> c_int>,
}

impl DynamicBuffer {
    /// Destroy the [`DynamicBuffer`] freeing the underlying memory using the provided
    /// `destructor_callback` and clearing or zeroing all the members.
    ///
    /// If the `pointer` stored in [`DynamicBuffer`] is NULL, then `length` is zeroed out and the
    /// `destructor_callback` is set to `None`. It is similar to how free ignores NULL in C, we just
    /// do some additional housekeeping to signal the [`DynamicBuffer`] is an empty shell.
    ///
    /// # Safety
    ///
    /// Destroy is safe to call only if the `destructor_callback` is the method that needs to be
    /// called to free the stored pointer. For example in C++, memory allocated with `new` must be
    /// freed with `delete`, memory allocated with `new[]` must be freed with `delete[]`.
    ///
    /// Length must indicate how many `u8` are present in the allocation and can be used by the
    /// `destructor_callback` to free memory. For example in the case of a `Vec` being turned into a
    /// [`DynamicBuffer`] the length is obtained by first calling the `len` function on the `Vec`.
    pub unsafe fn destroy(&mut self) -> Result<(), &str> {
        if self.pointer.is_null() {
            // Finish emptying stuff
            self.length = 0;
            self.destructor = None;
            return Ok(());
        }

        match self.destructor {
            Some(destructor_callback) => {
                let res = destructor_callback(self.pointer, self.length);
                if res == 0 {
                    // If the deallocation is successful then we empty the buffer
                    self.pointer = std::ptr::null_mut();
                    self.length = 0;
                    self.destructor = None;
                    return Ok(());
                }
                Err("destructor returned a non 0 error code")
            }
            // We could not free because of a missing destructor, return an error
            None => Err("destructor is NULL, could not destroy DynamicBuffer"),
        }
    }
}

unsafe extern "C" fn free_u8_ptr_built_from_vec_u8(pointer: *mut u8, length: usize) -> c_int {
    if pointer.is_null() {
        return 0;
    }

    let slice = std::slice::from_raw_parts_mut(pointer, length);

    drop(Box::from_raw(slice));

    0
}

impl From<Vec<u8>> for DynamicBuffer {
    fn from(value: Vec<u8>) -> Self {
        let boxed_slice = value.into_boxed_slice();
        let length = boxed_slice.len();
        let pointer = Box::leak(boxed_slice);

        Self {
            pointer: pointer.as_mut_ptr(),
            length,
            destructor: Some(free_u8_ptr_built_from_vec_u8),
        }
    }
}

fn check_ptr_is_non_null_and_aligned<T>(ptr: *const T) -> Result<(), String> {
    if ptr.is_null() {
        return Err(format!("pointer is null, got: {ptr:p}"));
    }
    let expected_alignment = std::mem::align_of::<T>();
    if ptr as usize % expected_alignment != 0 {
        return Err(format!(
            "pointer is misaligned, expected {expected_alignment} bytes alignment, got pointer: \
            {ptr:p}. You May have mixed some pointers in your function call.",
        ));
    }
    Ok(())
}

/// Get a mutable reference from a pointer checking the pointer is well aligned for the given type.
///
/// # Safety
///
/// Caller of this function needs to make sure the pointer type corresponds to the data type being
/// pointed to.
///
/// Caller of this function needs to make sure the aliasing rules for mutable reference are
/// respected.
///
/// The basics are: at any time only a single mutable reference may exist to a given memory location
/// XOR any number of immutable reference may exist to a given memory location.
///
/// Failure to abide by the above rules will result in undefined behavior (UB).
unsafe fn get_mut_checked<'a, T>(ptr: *mut T) -> Result<&'a mut T, String> {
    match check_ptr_is_non_null_and_aligned(ptr) {
        Ok(()) => ptr
            .as_mut()
            .ok_or_else(|| "Error while converting to mut reference".into()),
        Err(e) => Err(e),
    }
}

/// C API to destroy a [`DynamicBuffer`].
///
/// # Safety
///
/// This function is safe to call if `dynamic_buffer` is not aliased to avoid double frees.
#[no_mangle]
pub unsafe extern "C" fn destroy_dynamic_buffer(dynamic_buffer: *mut DynamicBuffer) -> c_int {
    // Mimicks C for calls of free on NULL, nothing occurs
    if dynamic_buffer.is_null() {
        return 0;
    }

    let dynamic_buffer = match get_mut_checked(dynamic_buffer) {
        Ok(dynamic_buffer) => dynamic_buffer,
        Err(cause) => {
            if cfg!(feature = "c_api_print_error_source") {
                println!("{cause}");
            }

            return 1;
        }
    };

    match dynamic_buffer.destroy() {
        Ok(_) => 0,
        Err(cause) => {
            if cfg!(feature = "c_api_print_error_source") {
                println!("{cause}");
            }

            1
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_dynamic_buffer_vec_u8_custom_destructor() {
        let vec = vec![99u8; 1000];
        let len = vec.len();
        let ptr = vec.leak();

        unsafe extern "C" fn custom_destroy_vec_u8_buffer(
            pointer: *mut u8,
            length: usize,
        ) -> c_int {
            if pointer.is_null() {
                return 0;
            }

            let slice = std::slice::from_raw_parts_mut(pointer, length);

            drop(Box::from_raw(slice));

            0
        }

        let mut dynamic_buffer = DynamicBuffer {
            pointer: ptr.as_mut_ptr(),
            length: len,
            destructor: Some(custom_destroy_vec_u8_buffer),
        };

        let res = unsafe { destroy_dynamic_buffer(&mut dynamic_buffer as *mut DynamicBuffer) };

        assert_eq!(res, 0);
        assert!(dynamic_buffer.pointer.is_null());
        assert_eq!(dynamic_buffer.length, 0);
        assert!(dynamic_buffer.destructor.is_none());

        assert!(dynamic_buffer.pointer.is_null());
        let res = unsafe { destroy_dynamic_buffer(&mut dynamic_buffer as *mut DynamicBuffer) };
        // Same as free in C, destroy on a NULL pointer does nothing
        assert_eq!(res, 0);
        assert!(dynamic_buffer.pointer.is_null());
        assert_eq!(dynamic_buffer.length, 0);
        assert!(dynamic_buffer.destructor.is_none());

        let mut some_u8 = 0u8;

        dynamic_buffer.pointer = &mut some_u8 as *mut u8;

        assert!(dynamic_buffer.destructor.is_none());
        let res = unsafe { destroy_dynamic_buffer(&mut dynamic_buffer as *mut DynamicBuffer) };
        assert_eq!(res, 1);
    }

    #[test]
    fn test_dynamic_buffer_vec_u8_default_destructor() {
        let vec = vec![99u8; 1000];

        let mut dynamic_buffer: DynamicBuffer = vec.clone().into();

        let res = unsafe { destroy_dynamic_buffer(&mut dynamic_buffer as *mut DynamicBuffer) };

        assert_eq!(res, 0);
        assert!(dynamic_buffer.pointer.is_null());
        assert_eq!(dynamic_buffer.length, 0);
        assert!(dynamic_buffer.destructor.is_none());

        let mut dynamic_buffer: DynamicBuffer = vec.into();

        let res = unsafe { dynamic_buffer.destroy() };

        assert!(res.is_ok());
        assert!(dynamic_buffer.pointer.is_null());
        assert_eq!(dynamic_buffer.length, 0);
        assert!(dynamic_buffer.destructor.is_none());
    }
}