musli_zerocopy/buf/

padder.rs

use core::marker::PhantomData;
use core::mem::{MaybeUninit, align_of, size_of, size_of_val, transmute};
use core::ptr::NonNull;

use crate::buf;
use crate::traits::{ZeroCopy, ZeroSized};

/// A struct padder as provided to the [`ZeroCopy::pad`] method.
///
/// This knows how to find and initialize padding regions in `repr(C)` types,
/// and provides a builder-like API to doing so.
#[must_use = "For the writer to have an effect on `OwnedBuf` you must call `Padder::remaining` / `Padder::remaining_unsized`"]
pub struct Padder<'a, T: ?Sized> {
    data: NonNull<MaybeUninit<u8>>,
    offset: usize,
    _marker: PhantomData<&'a mut T>,
}

impl<'a, T: ?Sized> Padder<'a, T> {
    #[inline]
    pub(crate) fn new(data: NonNull<MaybeUninit<u8>>) -> Self {
        Self {
            data,
            offset: 0,
            _marker: PhantomData,
        }
    }

    /// Indicate that this validate is transparent over `U`.
    //
    /// # Safety
    ///
    /// This is only allowed if `T` is `#[repr(transparent)]` over `U`.
    #[inline]
    pub unsafe fn transparent<U>(&mut self) -> &mut Padder<'a, U> {
        unsafe { transmute(self) }
    }

    /// Pad around the given field with zeros.
    ///
    /// Note that this is necessary to do correctly in order to satisfy the
    /// safety requirements by [`remaining()`].
    ///
    /// This is typically not called directly, but rather is implemented by the
    /// [`ZeroCopy`] derive.
    ///
    /// [`remaining()`]: Self::remaining
    /// [`ZeroCopy`]: derive@crate::ZeroCopy
    ///
    /// # Safety
    ///
    /// The caller must ensure that the field type `F` is an actual field in
    /// order in the struct being padded.
    #[inline]
    pub unsafe fn pad<F>(&mut self)
    where
        F: ZeroCopy,
    {
        unsafe {
            self.pad_with::<F>(align_of::<F>());
        }
    }

    /// Pad around the given field with zeros for a struct marked
    /// `#[repr(packed(align))]`.
    ///
    /// The field is aligned to `min(align, align_of::<F>())`, which is the
    /// alignment the compiler uses for fields in a packed struct.
    ///
    /// Note that this is necessary to do correctly in order to satisfy the
    /// safety requirements by [`remaining()`].
    ///
    /// This is typically not called directly, but rather is implemented by the
    /// [`ZeroCopy`] derive.
    ///
    /// [`remaining()`]: Self::remaining
    /// [`ZeroCopy`]: derive@crate::ZeroCopy
    ///
    /// # Safety
    ///
    /// The caller must ensure that the field type `F` is an actual field in
    /// order in the struct being padded and that `align` matches the argument
    /// provided to `#[repr(packed)]` (note that empty means 1). `align` must
    /// be a power of two.
    #[inline]
    pub unsafe fn pad_with<F>(&mut self, align: usize)
    where
        F: ZeroCopy,
    {
        unsafe {
            let align = align.min(align_of::<F>());
            let count = buf::padding_to(self.offset, align);
            // zero out padding.
            self.data.as_ptr().add(self.offset).write_bytes(0, count);
            self.offset += count;

            if F::PADDED {
                let ptr = NonNull::new_unchecked(self.data.as_ptr().add(self.offset));
                let mut padder = Padder::new(ptr);
                F::pad(&mut padder);
                padder.remaining();
            }

            self.offset += size_of::<F>();
        }
    }

    /// Pad up to an ignored zero-sized field `F`.
    ///
    /// Zero-sized fields can have an alignment larger than 1, in which case
    /// they affect the offset of the fields that follow them.
    ///
    /// This is typically not called directly, but rather is implemented by the
    /// [`ZeroCopy`] derive.
    ///
    /// [`ZeroCopy`]: derive@crate::ZeroCopy
    ///
    /// # Safety
    ///
    /// The caller must ensure that the field type `F` is an actual field in
    /// order in the struct being padded.
    #[inline]
    pub unsafe fn pad_zero_sized<F>(&mut self)
    where
        F: ZeroSized,
    {
        unsafe {
            self.pad_zero_sized_with::<F>(align_of::<F>());
        }
    }

    /// Pad up to an ignored zero-sized field `F` in a struct marked
    /// `#[repr(packed(align))]`.
    ///
    /// The field is aligned to `min(align, align_of::<F>())`.
    ///
    /// This is typically not called directly, but rather is implemented by the
    /// [`ZeroCopy`] derive.
    ///
    /// [`ZeroCopy`]: derive@crate::ZeroCopy
    ///
    /// # Safety
    ///
    /// The caller must ensure that the field type `F` is an actual field in
    /// order in the struct being padded and that `align` matches the argument
    /// provided to `#[repr(packed)]` (note that empty means 1). `align` must
    /// be a power of two.
    #[inline]
    pub unsafe fn pad_zero_sized_with<F>(&mut self, align: usize)
    where
        F: ZeroSized,
    {
        unsafe {
            let count = buf::padding_to(self.offset, align.min(align_of::<F>()));
            // zero out padding.
            self.data.as_ptr().add(self.offset).write_bytes(0, count);
            self.offset += count;
        }
    }

    /// Specific method to both pad for a discriminant and load it
    /// simultaneously for inspection.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the field type `D` is the actual type of the
    /// discriminant first in order in the enum being padded and that `D` is a
    /// primitive that does not contain any interior padding.
    #[inline]
    pub unsafe fn pad_discriminant<D>(&mut self) -> D
    where
        D: ZeroCopy,
    {
        unsafe {
            let count = buf::padding_to(self.offset, align_of::<D>());
            // zero out padding.
            self.data.as_ptr().add(self.offset).write_bytes(0, count);
            let at = self.offset + count;
            let value = self.data.as_ptr().add(at).cast::<D>().read_unaligned();
            self.offset = at + size_of::<D>();
            value
        }
    }

    /// Finish writing the current buffer.
    ///
    /// This is typically not called directly, but rather is implemented by the
    /// [`ZeroCopy`] derive.
    ///
    /// [`ZeroCopy`]: derive@crate::ZeroCopy
    ///
    /// # Safety
    ///
    /// Before calling `remaining()`, the caller must ensure that they've called
    /// [`pad::<F>()`] *in order* for every field in a struct being serialized
    /// where `F` is the type of the field. Otherwise we might not have written
    /// the necessary padding to ensure that all bytes related to the struct are
    /// initialized. Failure to do so would result in undefined behavior.
    ///
    /// Fields which are [`ZeroSized`] must be padded with
    /// [`pad_zero_sized::<F>()`] unless their alignment is 1, since they can
    /// affect the offset of the fields that follow them.
    ///
    /// [`pad::<F>()`]: Self::pad
    /// [`pad_zero_sized::<F>()`]: Self::pad_zero_sized
    /// [`ZeroSized`]: crate::traits::ZeroSized
    #[inline]
    pub unsafe fn remaining(self)
    where
        T: Sized,
    {
        debug_assert!(
            self.offset <= size_of::<T>(),
            "padded fields extend past the end of the type"
        );

        unsafe {
            let count = size_of::<T>() - self.offset;
            self.data.as_ptr().add(self.offset).write_bytes(0, count);
        }
    }

    /// Finalize remaining padding based on the size of an unsized value.
    #[inline]
    pub(crate) unsafe fn remaining_unsized(self, value: &T) {
        debug_assert!(
            self.offset <= size_of_val(value),
            "padded fields extend past the end of the value"
        );

        unsafe {
            let count = size_of_val(value) - self.offset;
            self.data.as_ptr().add(self.offset).write_bytes(0, count);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;

    use anyhow::Result;

    use crate::{OwnedBuf, ZeroCopy, buf};

    #[test]
    fn ensure_padding() -> Result<()> {
        #[derive(Debug, PartialEq, Eq, ZeroCopy)]
        #[repr(C)]
        #[zero_copy(crate)]
        struct ZeroPadded(u8, u16, u64);

        let padded = ZeroPadded(
            0x01u8.to_be(),
            0x0203u16.to_be(),
            0x0405060708090a0bu64.to_be(),
        );

        let mut buf = OwnedBuf::new();

        // Note: You're responsible for ensuring that the buffer has enough
        // capacity.
        buf.reserve(size_of::<ZeroPadded>())?;

        // SAFETY: We do not pad beyond known fields and are making sure to
        // initialize all of the buffer.
        unsafe {
            buf::store_unaligned(buf.as_nonnull(), &padded);
            buf.advance(size_of::<ZeroPadded>());
        }

        // Note: The bytes are explicitly convert to big-endian encoding above.
        assert_eq!(
            buf.as_slice(),
            &[1, 0, 2, 3, 0, 0, 0, 0, 4, 5, 6, 7, 8, 9, 10, 11]
        );

        Ok(())
    }
}