Skip to main content

micromegas_transit/
serialize.rs

1use anyhow::{Result, bail};
2
3#[allow(unsafe_code)]
4#[inline(always)]
5pub fn write_any<T>(buffer: &mut Vec<u8>, value: &T) {
6    let ptr = std::ptr::addr_of!(*value).cast::<u8>();
7    let slice = std::ptr::slice_from_raw_parts(ptr, std::mem::size_of::<T>());
8    unsafe {
9        buffer.extend_from_slice(&*slice);
10    }
11}
12
13#[allow(unsafe_code)]
14/// Helper function to read a u* pointer to a value of type T.
15///
16/// # Safety
17/// ptr must be valid it's size and it's memory size must be the size
18/// of T or higher.
19#[inline(always)]
20pub unsafe fn read_any<T>(ptr: *const u8) -> T {
21    unsafe { std::ptr::read_unaligned(ptr.cast::<T>()) }
22}
23
24/// Trusted-path window advance: panics if `offset` exceeds the window length.
25/// Only safe to use on same-process, self-produced buffers (in-proc queue
26/// reads). Payload-derived (untrusted) offsets must use [`try_advance_window`].
27pub fn advance_window(window: &[u8], offset: usize) -> &[u8] {
28    assert!(offset <= window.len());
29    &window[offset..]
30}
31
32/// Checked variant of [`advance_window`]: returns `Err` instead of panicking
33/// when `offset` exceeds the window length. Use for payload- or
34/// metadata-derived offsets, which must be treated as untrusted.
35#[inline(always)]
36pub fn try_advance_window(window: &[u8], offset: usize) -> Result<&[u8]> {
37    if offset > window.len() {
38        bail!(
39            "truncated window: need {offset} bytes, have {}",
40            window.len()
41        );
42    }
43    Ok(&window[offset..])
44}
45
46/// Trusted-path pod read: panics (via `advance_window`'s assert) if the
47/// window is shorter than `size_of::<T>()`. Only safe to use on same-process,
48/// self-produced buffers (in-proc queue reads, e.g.
49/// `InProcSerialize::read_value`). Payload-derived (untrusted) windows must
50/// use [`try_read_consume_pod`].
51pub fn read_consume_pod<T>(window: &mut &[u8]) -> T {
52    let object_size = std::mem::size_of::<T>();
53    let begin: *const u8 = window.as_ptr();
54    *window = advance_window(window, object_size);
55    unsafe { std::ptr::read_unaligned(begin.cast::<T>()) }
56}
57
58/// Checked variant of [`read_consume_pod`]: returns `Err` instead of
59/// panicking when the window is shorter than `size_of::<T>()`. Use for
60/// payload-derived (untrusted) windows.
61#[allow(unsafe_code)]
62#[inline(always)]
63pub fn try_read_consume_pod<T>(window: &mut &[u8]) -> Result<T> {
64    let object_size = std::mem::size_of::<T>();
65    if object_size > window.len() {
66        bail!(
67            "truncated window reading {}: need {object_size} bytes, have {}",
68            std::any::type_name::<T>(),
69            window.len()
70        );
71    }
72    let begin: *const u8 = window.as_ptr();
73    *window = &window[object_size..];
74    Ok(unsafe { std::ptr::read_unaligned(begin.cast::<T>()) })
75}
76
77/// Bounds-checked read of a POD value at `offset` within `window`. Replaces
78/// `read_any(window.as_ptr().add(offset))` on untrusted windows — both
79/// `offset` and `size_of::<T>()` may originate from untrusted stream
80/// metadata, so their sum is validated with `checked_add` before the window
81/// length check.
82#[allow(unsafe_code)]
83#[inline(always)]
84pub fn try_read_pod_at<T>(window: &[u8], offset: usize) -> Result<T> {
85    let object_size = std::mem::size_of::<T>();
86    let end = match offset.checked_add(object_size) {
87        Some(end) => end,
88        None => bail!(
89            "offset {offset} overflows reading {}",
90            std::any::type_name::<T>()
91        ),
92    };
93    if end > window.len() {
94        bail!(
95            "truncated window reading {} at offset {offset}: need {object_size} bytes, have {}",
96            std::any::type_name::<T>(),
97            window.len().saturating_sub(offset)
98        );
99    }
100    Ok(unsafe { std::ptr::read_unaligned(window.as_ptr().add(offset).cast::<T>()) })
101}
102
103/// Helps speed up the serialization of types which size is known at compile time.
104pub enum InProcSize {
105    Const(usize),
106    Dynamic,
107}
108
109// InProcSerialize is used by the heterogeneous queue to write objects in its
110// buffer serialized objects can have references with static lifetimes
111pub trait InProcSerialize: Sized {
112    const IN_PROC_SIZE: InProcSize = InProcSize::Const(std::mem::size_of::<Self>());
113
114    fn get_value_size(&self) -> Option<u32> {
115        // for POD serialization we don't write the size of each instance
116        // the metadata will contain this size
117        None
118    }
119
120    #[inline(always)]
121    fn write_value(&self, buffer: &mut Vec<u8>) {
122        assert!(matches!(Self::IN_PROC_SIZE, InProcSize::Const(_)));
123        #[allow(clippy::needless_borrow)]
124        //clippy complains here but we don't want to move or copy the value
125        write_any::<Self>(buffer, &self);
126    }
127
128    // read_value allows to read objects from the same process they were stored in
129    // i.e. iterating in the heterogenous queue
130    /// # Safety
131    /// This is called from the serializer context that that uses `value_size`
132    /// call to make sure that the proper size is used
133    #[allow(unsafe_code)]
134    #[inline(always)]
135    unsafe fn read_value(mut window: &[u8]) -> Self {
136        let res = read_consume_pod(&mut window);
137        assert_eq!(window.len(), 0);
138        res
139    }
140}