view src/libpam/memory.rs @ 141:a508a69c068a

Remove a lot of Results from functions. Many functions are documented to only return failing Results when given improper inputs or when there is a memory allocation failure (which can be verified by looking at the source). In cases where we know our input is correct, we don't need to check for memory allocation errors for the same reason that Rust doesn't do so when you, e.g., create a new Vec.
author Paul Fisher <paul@pfish.zone>
date Sat, 05 Jul 2025 17:16:56 -0400
parents 33b9622ed6d2
children ebb71a412b58
line wrap: on
line source

//! Things for dealing with memory.

use crate::ErrorCode;
use crate::Result;
use libpam_sys_helpers::memory::{Buffer, OwnedBinaryPayload};
use std::ffi::{c_char, CStr, CString};
use std::marker::{PhantomData, PhantomPinned};
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;
use std::{mem, ptr, slice};

/// Allocates `count` elements to hold `T`.
#[inline]
pub fn calloc<T>(count: usize) -> NonNull<T> {
    // SAFETY: it's always safe to allocate! Leaking memory is fun!
    unsafe { NonNull::new_unchecked(libc::calloc(count, mem::size_of::<T>()).cast()) }
}

/// Wrapper for [`libc::free`] to make debugging calls/frees easier.
///
/// # Safety
///
/// If you double-free, it's all your fault.
#[inline]
pub unsafe fn free<T>(p: *mut T) {
    libc::free(p.cast())
}

/// Makes whatever it's in not [`Send`], [`Sync`], or [`Unpin`].
#[repr(C)]
#[derive(Debug, Default)]
pub struct Immovable(pub PhantomData<(*mut u8, PhantomPinned)>);

/// Safely converts a `&str` option to a `CString` option.
pub fn option_cstr(prompt: Option<&str>) -> Result<Option<CString>> {
    prompt
        .map(CString::new)
        .transpose()
        .map_err(|_| ErrorCode::ConversationError)
}

/// Gets the pointer to the given CString, or a null pointer if absent.
pub fn prompt_ptr(prompt: Option<&CString>) -> *const c_char {
    match prompt {
        Some(c_str) => c_str.as_ptr(),
        None => ptr::null(),
    }
}

/// It's like a [`Box`], but C heap managed.
#[derive(Debug)]
#[repr(transparent)]
pub struct CHeapBox<T>(NonNull<T>);

// Lots of "as" and "into" associated functions.
#[allow(clippy::wrong_self_convention)]
impl<T> CHeapBox<T> {
    /// Creates a new CHeapBox holding the given data.
    pub fn new(value: T) -> Self {
        let memory = calloc(1);
        unsafe { ptr::write(memory.as_ptr(), value) }
        // SAFETY: We literally just allocated this.
        Self(memory)
    }

    /// Takes ownership of the given pointer.
    ///
    /// # Safety
    ///
    /// You have to provide a valid pointer to the start of an allocation
    /// that was made with `malloc`.
    pub unsafe fn from_ptr(ptr: NonNull<T>) -> Self {
        Self(ptr)
    }

    /// Converts this CBox into a raw pointer.
    pub fn into_ptr(this: Self) -> NonNull<T> {
        ManuallyDrop::new(this).0
    }

    /// Gets a pointer from this but doesn't convert this into a raw pointer.
    ///
    /// You are responsible for ensuring the CHeapBox lives long enough.
    pub fn as_ptr(this: &Self) -> NonNull<T> {
        this.0
    }

    /// Because it's annoying to type `CHeapBox.as_ptr(...).as_ptr()`.
    pub fn as_raw_ptr(this: &Self) -> *mut T {
        this.0.as_ptr()
    }

    /// Converts this into a Box of a different type.
    ///
    /// # Safety
    ///
    /// The other type has to have the same size and alignment and
    /// have compatible drop behavior with respect to other resources.
    pub unsafe fn cast<R>(this: Self) -> CHeapBox<R> {
        mem::transmute(this)
    }
}

impl<T: Default> Default for CHeapBox<T> {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

impl Buffer for CHeapBox<u8> {
    fn allocate(len: usize) -> Self {
        // SAFETY: This is all freshly-allocated memory!
        unsafe { Self::from_ptr(calloc(len)) }
    }

    fn as_ptr(this: &Self) -> *const u8 {
        this.0.as_ptr()
    }

    unsafe fn as_mut_slice(this: &mut Self, len: usize) -> &mut [u8] {
        slice::from_raw_parts_mut(this.0.as_ptr(), len)
    }

    fn into_ptr(this: Self) -> NonNull<u8> {
        CHeapBox::into_ptr(this)
    }

    unsafe fn from_ptr(ptr: NonNull<u8>, _: usize) -> Self {
        CHeapBox::from_ptr(ptr)
    }
}

pub type CHeapPayload = OwnedBinaryPayload<CHeapBox<u8>>;

impl<T> Deref for CHeapBox<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        // SAFETY: We own this pointer and it is guaranteed valid.
        unsafe { Self::as_ptr(self).as_ref() }
    }
}

impl<T> DerefMut for CHeapBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: We own this pointer and it is guaranteed valid.
        unsafe { Self::as_ptr(self).as_mut() }
    }
}

impl<T> Drop for CHeapBox<T> {
    fn drop(&mut self) {
        // SAFETY: We own a valid pointer, and will never use it after this.
        unsafe {
            let ptr = self.0.as_ptr();
            ptr::drop_in_place(ptr);
            free(ptr)
        }
    }
}

/// A null-terminated string allocated on the C heap.
///
/// Basically [`CString`], but managed by malloc.
#[derive(Debug)]
#[repr(transparent)]
pub struct CHeapString(CHeapBox<c_char>);

impl CHeapString {
    /// Creates a new C heap string with the given contents.
    pub fn new(text: &str) -> Result<Self> {
        let data = text.as_bytes();
        if data.contains(&0) {
            return Err(ErrorCode::ConversationError);
        }
        // +1 for the null terminator
        let data_alloc: NonNull<c_char> = calloc(data.len() + 1);
        // SAFETY: we just allocated this and we have enough room.
        unsafe {
            let dest = slice::from_raw_parts_mut(data_alloc.as_ptr().cast(), data.len());
            dest.copy_from_slice(data);
            Ok(Self(CHeapBox::from_ptr(data_alloc)))
        }
    }

    /// Converts this C heap string into a raw pointer.
    ///
    /// You are responsible for freeing it later.
    pub fn into_ptr(self) -> NonNull<c_char> {
        let this = ManuallyDrop::new(self);
        CHeapBox::as_ptr(&this.0)
    }

    /// Converts this into a dumb box. It will no longer be zeroed upon drop.
    pub fn into_box(self) -> CHeapBox<c_char> {
        unsafe { mem::transmute(self) }
    }

    /// Takes ownership of a C heap string.
    ///
    /// # Safety
    ///
    /// You have to provide a pointer to the start of an allocation that is
    /// a valid 0-terminated C string.
    unsafe fn from_ptr(ptr: *mut c_char) -> Option<Self> {
        NonNull::new(ptr).map(|p| unsafe { Self(CHeapBox::from_ptr(p)) })
    }

    unsafe fn from_box<T>(bx: CHeapBox<T>) -> Self {
        Self(CHeapBox::cast(bx))
    }

    /// Zeroes the contents of a C string.
    ///
    /// # Safety
    ///
    /// You have to provide a valid pointer to a null-terminated C string.
    pub unsafe fn zero(ptr: NonNull<c_char>) {
        let cstr = ptr.as_ptr();
        let len = libc::strlen(cstr.cast());
        for x in 0..len {
            ptr::write_volatile(cstr.byte_offset(x as isize), mem::zeroed())
        }
    }
}

impl Drop for CHeapString {
    fn drop(&mut self) {
        // SAFETY: We own a valid C String
        unsafe { Self::zero(CHeapBox::as_ptr(&self.0)) }
    }
}

impl Deref for CHeapString {
    type Target = CStr;

    fn deref(&self) -> &Self::Target {
        // SAFETY: We know we own a valid C string pointer.
        let ptr = CHeapBox::as_ptr(&self.0).as_ptr();
        unsafe { CStr::from_ptr(ptr) }
    }
}

/// Creates an owned copy of a string that is returned from a
/// <code>pam_get_<var>whatever</var></code> function.
///
/// # Safety
///
/// It's on you to provide a valid string.
pub unsafe fn copy_pam_string(result_ptr: *const c_char) -> Result<Option<String>> {
    let borrowed = match NonNull::new(result_ptr.cast_mut()) {
        Some(data) => Some(
            CStr::from_ptr(data.as_ptr())
                .to_str()
                .map_err(|_| ErrorCode::ConversationError)?,
        ),
        None => return Ok(None),
    };
    Ok(borrowed.map(String::from))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cell::Cell;
    #[test]
    fn test_box() {
        let drop_count: Cell<u32> = Cell::new(0);

        struct Dropper<'a>(&'a Cell<u32>);

        impl Drop for Dropper<'_> {
            fn drop(&mut self) {
                self.0.set(self.0.get() + 1)
            }
        }

        let mut dropbox = CHeapBox::new(Dropper(&drop_count));
        _ = dropbox;
        // ensure the old value is dropped when the new one is assigned.
        dropbox = CHeapBox::new(Dropper(&drop_count));
        assert_eq!(1, drop_count.get());
        *dropbox = Dropper(&drop_count);
        assert_eq!(2, drop_count.get());
        drop(dropbox);
        assert_eq!(3, drop_count.get());
    }
    #[test]
    fn test_strings() {
        let str = CHeapString::new("hello there").unwrap();
        let str_ptr = str.into_ptr().as_ptr();
        CHeapString::new("hell\0 there").unwrap_err();
        unsafe {
            let copied = copy_pam_string(str_ptr).unwrap();
            assert_eq!("hello there", copied.unwrap());
            CHeapString::zero(NonNull::new(str_ptr).unwrap());
            let idx_three = str_ptr.add(3).as_mut().unwrap();
            *idx_three = 0x80u8 as i8;
            let zeroed = copy_pam_string(str_ptr).unwrap().unwrap();
            assert!(zeroed.is_empty());
            let _ = CHeapString::from_ptr(str_ptr);
        }
    }

    #[test]
    fn test_option_str() {
        let good = option_cstr(Some("whatever")).unwrap();
        assert_eq!("whatever", good.unwrap().to_str().unwrap());
        let no_str = option_cstr(None).unwrap();
        assert!(no_str.is_none());
        let bad_str = option_cstr(Some("what\0ever")).unwrap_err();
        assert_eq!(ErrorCode::ConversationError, bad_str);
    }

    #[test]
    fn test_prompt() {
        let prompt_cstr = CString::new("good").ok();
        let prompt = prompt_ptr(prompt_cstr.as_ref());
        assert!(!prompt.is_null());
        let no_prompt = prompt_ptr(None);
        assert!(no_prompt.is_null());
    }
}