fidl_next_codec/wire/

envelope.rs

// Copyright 2024 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use core::mem::{ManuallyDrop, MaybeUninit};
use core::ptr::addr_of_mut;

use munge::munge;

use crate::decoder::InternalHandleDecoder;
use crate::encoder::InternalHandleEncoder;
use crate::{
    CHUNK_SIZE, Constrained, Decode, DecodeError, Decoder, DecoderExt as _, Encode, EncodeError,
    Encoder, EncoderExt as _, Slot, ValidationError, Wire, wire,
};

#[derive(Clone, Copy)]
#[repr(C)]
struct Encoded {
    maybe_num_bytes: wire::Uint32,
    num_handles: wire::Uint16,
    flags: wire::Uint16,
}

const INLINE_SIZE: usize = 4;

/// A FIDL envelope
#[repr(C, align(8))]
pub union Envelope {
    zero: [u8; 8],
    encoded: Encoded,
    decoded_inline: [MaybeUninit<u8>; INLINE_SIZE],
    decoded_out_of_line: *mut (),
}

unsafe impl Send for Envelope {}
unsafe impl Sync for Envelope {}

impl Constrained for Envelope {
    type Constraint = ();

    fn validate(_: Slot<'_, Self>, _: Self::Constraint) -> Result<(), ValidationError> {
        Ok(())
    }
}

unsafe impl Wire for Envelope {
    type Narrowed<'de> = Self;

    fn zero_padding(_: &mut MaybeUninit<Self>) {}
}

impl Envelope {
    const IS_INLINE_BIT: u16 = 1;

    /// Encodes a zero envelope into a slot.
    #[inline]
    pub fn encode_zero(out: &mut MaybeUninit<Self>) {
        out.write(Envelope { zero: [0; 8] });
    }

    /// Encodes a `'static` value into an envelope with an encoder.
    #[inline]
    pub fn encode_value_static<W: Wire, E: InternalHandleEncoder + ?Sized>(
        value: impl Encode<W, E>,
        encoder: &mut E,
        out: &mut MaybeUninit<Self>,
        constraint: W::Constraint,
    ) -> Result<(), EncodeError> {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let encoded = unsafe {
            munge!(let Self { encoded } = out);
            encoded
        };
        munge! {
            let Encoded {
                maybe_num_bytes,
                num_handles,
                flags,
            } = encoded;
        }

        let handles_before = encoder.__internal_handle_count();

        let encoded_size = size_of::<W>();
        if encoded_size <= INLINE_SIZE {
            // If the encoded inline value is less than 4 bytes long, we need to zero out the part
            // that won't get written over
            unsafe {
                maybe_num_bytes
                    .as_mut_ptr()
                    .cast::<u8>()
                    .add(encoded_size)
                    .write_bytes(0, INLINE_SIZE - encoded_size);
            }
        } else {
            return Err(EncodeError::ExpectedInline(encoded_size));
        }

        let value_out = unsafe { &mut *maybe_num_bytes.as_mut_ptr().cast() };
        W::zero_padding(value_out);
        value.encode(encoder, value_out, constraint)?;

        flags.write(wire::Uint16(Self::IS_INLINE_BIT));

        let handle_count = (encoder.__internal_handle_count() - handles_before).try_into().unwrap();
        num_handles.write(wire::Uint16(handle_count));

        Ok(())
    }

    /// Encodes a value into an envelope with an encoder.
    #[inline]
    pub fn encode_value<W: Wire, E: Encoder + ?Sized>(
        value: impl Encode<W, E>,
        encoder: &mut E,
        out: &mut MaybeUninit<Self>,
        constraint: W::Constraint,
    ) -> Result<(), EncodeError> {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let encoded = unsafe {
            munge!(let Self { encoded } = out);
            encoded
        };
        munge! {
            let Encoded {
                maybe_num_bytes,
                num_handles,
                flags,
            } = encoded;
        }

        let handles_before = encoder.__internal_handle_count();

        let encoded_size = size_of::<W>();
        if encoded_size <= INLINE_SIZE {
            // If the encoded inline value is less than 4 bytes long, we need to zero out the part
            // that won't get written over
            unsafe {
                maybe_num_bytes
                    .as_mut_ptr()
                    .cast::<u8>()
                    .add(encoded_size)
                    .write_bytes(0, INLINE_SIZE - encoded_size);
            }
            let value_out = unsafe { &mut *maybe_num_bytes.as_mut_ptr().cast() };
            W::zero_padding(value_out);
            value.encode(encoder, value_out, constraint)?;
            flags.write(wire::Uint16(Self::IS_INLINE_BIT));
        } else {
            let bytes_before = encoder.bytes_written();

            encoder.encode_next_with_constraint(value, constraint)?;

            let bytes_count = (encoder.bytes_written() - bytes_before).try_into().unwrap();
            maybe_num_bytes.write(wire::Uint32(bytes_count));
            flags.write(wire::Uint16(0));
        }

        let handle_count = (encoder.__internal_handle_count() - handles_before).try_into().unwrap();
        num_handles.write(wire::Uint16(handle_count));

        Ok(())
    }

    /// Returns the zero envelope.
    #[inline]
    pub fn zero() -> Self {
        Self { zero: [0; 8] }
    }

    /// Returns whether a envelope slot is encoded as zero.
    #[inline]
    pub fn is_encoded_zero(slot: Slot<'_, Self>) -> bool {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let zero = unsafe {
            munge!(let Self { zero } = slot);
            zero
        };
        *zero == [0; 8]
    }

    /// Returns whether an envelope is zero.
    #[inline]
    pub fn is_zero(&self) -> bool {
        unsafe { self.zero == [0; 8] }
    }

    #[inline]
    fn out_of_line_chunks(
        maybe_num_bytes: Slot<'_, wire::Uint32>,
        flags: Slot<'_, wire::Uint16>,
    ) -> Result<Option<usize>, DecodeError> {
        match **flags {
            Self::IS_INLINE_BIT => Ok(None),
            0 => {
                let num_bytes = **maybe_num_bytes;
                if !(num_bytes as usize).is_multiple_of(CHUNK_SIZE) {
                    Err(DecodeError::InvalidEnvelopeSize(num_bytes))
                } else if num_bytes <= INLINE_SIZE as u32 {
                    Err(DecodeError::OutOfLineValueTooSmall(num_bytes))
                } else {
                    Ok(Some(num_bytes as usize / CHUNK_SIZE))
                }
            }
            _ => Err(DecodeError::InvalidEnvelopeFlags(**flags)),
        }
    }

    /// Decodes and discards a static type in an envelope.
    #[inline]
    pub fn decode_unknown_static<D: InternalHandleDecoder + ?Sized>(
        slot: Slot<'_, Self>,
        decoder: &mut D,
    ) -> Result<(), DecodeError> {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let encoded = unsafe {
            munge!(let Self { encoded } = slot);
            encoded
        };
        munge! {
            let Encoded {
                maybe_num_bytes,
                num_handles,
                flags,
            } = encoded;
        }

        if let Some(count) = Self::out_of_line_chunks(maybe_num_bytes, flags)? {
            return Err(DecodeError::ExpectedInline(count * CHUNK_SIZE));
        }

        decoder.__internal_take_handles(**num_handles as usize)?;

        Ok(())
    }

    /// Decodes and discards an unknown value in an envelope.
    #[inline]
    pub fn decode_unknown<'de, D: Decoder<'de> + ?Sized>(
        slot: Slot<'_, Self>,
        decoder: &mut D,
    ) -> Result<(), DecodeError> {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let encoded = unsafe {
            munge!(let Self { encoded } = slot);
            encoded
        };
        munge! {
            let Encoded {
                maybe_num_bytes,
                num_handles,
                flags,
            } = encoded;
        }

        if let Some(count) = Self::out_of_line_chunks(maybe_num_bytes, flags)? {
            decoder.take_chunks(count)?;
        }

        decoder.__internal_take_handles(**num_handles as usize)?;

        Ok(())
    }

    /// Decodes a value of a known type from an envelope.
    #[inline]
    pub fn decode_as_static<D: InternalHandleDecoder + ?Sized, T: Decode<D>>(
        mut slot: Slot<'_, Self>,
        decoder: &mut D,
        constraint: T::Constraint,
    ) -> Result<(), DecodeError> {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let encoded = unsafe {
            munge!(let Self { encoded } = slot.as_mut());
            encoded
        };
        munge! {
            let Encoded {
                maybe_num_bytes,
                num_handles,
                flags,
            } = encoded;
        }

        let handles_before = decoder.__internal_handles_remaining();
        let num_handles = **num_handles as usize;

        if let Some(count) = Self::out_of_line_chunks(maybe_num_bytes, flags)? {
            return Err(DecodeError::ExpectedInline(count * CHUNK_SIZE));
        }

        // Decode inline value
        if size_of::<T>() > INLINE_SIZE {
            return Err(DecodeError::InlineValueTooBig(size_of::<T>()));
        }
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let mut decoded_inline = unsafe {
            munge!(let Self { decoded_inline } = slot);
            decoded_inline
        };
        let mut slot = unsafe { Slot::<T>::new_unchecked(decoded_inline.as_mut_ptr().cast()) };
        T::decode(slot.as_mut(), decoder, constraint)?;

        let handles_consumed = handles_before - decoder.__internal_handles_remaining();
        if handles_consumed != num_handles {
            return Err(DecodeError::IncorrectNumberOfHandlesConsumed {
                expected: num_handles,
                actual: handles_consumed,
            });
        }

        Ok(())
    }

    /// Decodes a value of a known type from an envelope.
    #[inline]
    pub fn decode_as<'de, D: Decoder<'de> + ?Sized, T: Decode<D>>(
        mut slot: Slot<'_, Self>,
        decoder: &mut D,
        constraint: T::Constraint,
    ) -> Result<(), DecodeError> {
        // `unsafe` block required in the next version of munge
        #[allow(unused_unsafe)]
        let encoded = unsafe {
            munge!(let Self { encoded } = slot.as_mut());
            encoded
        };
        munge! {
            let Encoded {
                mut maybe_num_bytes,
                num_handles,
                flags,
            } = encoded;
        }

        let handles_before = decoder.__internal_handles_remaining();
        let num_handles = **num_handles as usize;

        let out_of_line_chunks = Self::out_of_line_chunks(maybe_num_bytes.as_mut(), flags)?;
        if let Some(_count) = out_of_line_chunks {
            // Decode out-of-line value
            // TODO: set cap on decoder to make sure that the envelope doesn't decode more bytes
            // than it claims that it will
            let mut value_slot = decoder.take_slot::<T>()?;
            let value_ptr = value_slot.as_mut_ptr();
            T::decode(value_slot, decoder, constraint)?;

            // `unsafe` block required in the next version of munge
            #[allow(unused_unsafe)]
            let mut decoded_out_of_line = unsafe {
                munge!(let Self { decoded_out_of_line } = slot);
                decoded_out_of_line
            };
            // SAFETY: Identical to `ptr.write(value_ptr.cast())`, but raw
            // pointers don't currently implement `IntoBytes`.
            unsafe { decoded_out_of_line.as_mut_ptr().write(value_ptr.cast()) };
        } else {
            // Decode inline value
            if size_of::<T>() > INLINE_SIZE {
                return Err(DecodeError::InlineValueTooBig(size_of::<T>()));
            }
            // `unsafe` block required in the next version of munge
            #[allow(unused_unsafe)]
            let mut decoded_inline = unsafe {
                munge!(let Self { decoded_inline } = slot);
                decoded_inline
            };
            let mut slot = unsafe { Slot::<T>::new_unchecked(decoded_inline.as_mut_ptr().cast()) };
            T::decode(slot.as_mut(), decoder, constraint)?;
        }

        let handles_consumed = handles_before - decoder.__internal_handles_remaining();
        if handles_consumed != num_handles {
            return Err(DecodeError::IncorrectNumberOfHandlesConsumed {
                expected: num_handles,
                actual: handles_consumed,
            });
        }

        Ok(())
    }

    #[inline]
    unsafe fn as_ptr<T>(this: *mut Self) -> *mut T {
        if size_of::<T>() <= INLINE_SIZE {
            let inline = unsafe { addr_of_mut!((*this).decoded_inline) };
            inline.cast()
        } else {
            unsafe { (*this).decoded_out_of_line.cast() }
        }
    }

    /// Returns a reference to the contained `T`.
    ///
    /// # Safety
    ///
    /// The envelope must have been successfully decoded as a `T`.
    #[inline]
    pub unsafe fn deref_unchecked<T>(&self) -> &T {
        let ptr = unsafe { Self::as_ptr::<T>((self as *const Self).cast_mut()).cast_const() };
        unsafe { &*ptr }
    }

    /// Returns an `Owned` to the contained `T`.
    ///
    /// # Safety
    ///
    ///  The envelope must have been successfully decoded as a `T`.
    #[inline]
    pub unsafe fn read_unchecked<T>(&self) -> T {
        // SAFETY: `into_raw(this)` is guaranteed to return a pointer that is non-null, properly
        // properly aligned, and valid for reads and writes.
        unsafe { Self::as_ptr::<T>((self as *const Self).cast_mut()).read() }
    }

    /// Clones the envelope, assuming that it contains an inline `T`.
    ///
    /// # Safety
    ///
    /// The envelope must have been successfully decoded inline as a `T`.
    #[inline]
    pub unsafe fn clone_inline_unchecked<T: Clone>(&self) -> Self {
        debug_assert!(size_of::<T>() <= INLINE_SIZE);

        union ClonedToDecodedInline<T> {
            cloned: ManuallyDrop<T>,
            decoded_inline: [MaybeUninit<u8>; INLINE_SIZE],
        }

        let cloned = unsafe { self.deref_unchecked::<T>().clone() };
        unsafe {
            Self {
                decoded_inline: ClonedToDecodedInline { cloned: ManuallyDrop::new(cloned) }
                    .decoded_inline,
            }
        }
    }
}