flatbuffers_util/
ownedfb.rs

1use flatbuffers::{Follow, InvalidFlatbuffer, Verifiable};
2
3/// Stores the owned bytes of the flatbuffer type
4/// and can access the actual type.
5pub struct OwnedFB<T> {
6    buf: Vec<u8>,
7    index: usize,
8    _phantom: std::marker::PhantomData<T>,
9}
10
11impl<T> OwnedFB<T> {
12    pub fn new<'a>(buf: &'a [u8]) -> Result<OwnedFB<T>, InvalidFlatbuffer>
13    where
14        T: Verifiable + Follow<'a> + 'a,
15    {
16        check_flatbuffer::<T>(buf, 0)?;
17        Ok(unsafe { Self::new_from_vec_unchecked(buf.to_owned(), 0) })
18    }
19
20    /// # Safety
21    /// Caller is responsible for verifying the buffer and align the type T.
22    pub unsafe fn new_from_vec_unchecked(buf: Vec<u8>, index: usize) -> Self {
23        Self {
24            buf,
25            index,
26            _phantom: std::marker::PhantomData,
27        }
28    }
29
30    /// # Safety
31    /// Caller is responsible for verifying the buffer and align the type T.
32    pub unsafe fn new_from_builder_collapse(pair: (Vec<u8>, usize)) -> Self {
33        unsafe { Self::new_from_vec_unchecked(pair.0, pair.1) }
34    }
35
36    pub fn new_from_vec(buf: Vec<u8>, index: usize) -> Result<OwnedFB<T>, InvalidFlatbuffer>
37    where
38        T: Verifiable + Follow<'static> + 'static,
39    {
40        check_flatbuffer::<T>(&buf, index)?;
41
42        Ok(unsafe { Self::new_from_vec_unchecked(buf, index) })
43    }
44
45    /// This in practice is not zero copy for tonic.
46    pub fn new_from_bytes(buf: bytes::Bytes) -> Result<OwnedFB<T>, InvalidFlatbuffer>
47    where
48        T: Verifiable + Follow<'static> + 'static,
49    {
50        match buf.try_into_mut() {
51            // This is zero copy if the Bytes has the full ownership of the vec.
52            Ok(vec) => Self::new_from_vec(vec.into(), 0),
53            // This will make a copy of the bytes.
54            Err(bytes) => {
55                // This happens in tonic DecodeBuf because it is a shared BytesMut.
56                // So in practice tonic forces us to make a copy here.
57                // debug_assert!(false, "zero copy failed.");
58                Self::new_from_vec(bytes.to_vec(), 0)
59            }
60        }
61    }
62
63    pub fn get_ref<'a>(&'a self) -> <T as Follow<'a>>::Inner
64    where
65        T: Follow<'a>,
66    {
67        // Safety: We have already verified the buffer in `new_owned_fb`.
68        unsafe { get_ref_flatbuffer_unchecked::<'a, T>(&self.buf, self.index) }
69    }
70
71    pub fn get_slice(&self) -> &[u8] {
72        &self.buf[self.index..]
73    }
74
75    /// This may be zero copy if the vec capacity equals to length and index is zero.
76    pub fn into_bytes(self) -> bytes::Bytes {
77        // This is zero copy if vec cap == len.
78        debug_assert_eq!(self.buf.capacity(), self.buf.len());
79        let full_bytes = bytes::Bytes::from(self.buf);
80        // adjust the offset should be zero copy.
81        full_bytes.slice(self.index..)
82    }
83}
84
85/// Generic check.
86pub fn check_flatbuffer<'a, T>(buf: &[u8], index: usize) -> Result<(), InvalidFlatbuffer>
87where
88    T: Verifiable + Follow<'a> + 'a,
89{
90    let opts = flatbuffers::VerifierOptions::default();
91    let mut v = flatbuffers::Verifier::new(&opts, buf);
92    <flatbuffers::ForwardsUOffset<T>>::run_verifier(&mut v, index)?;
93    Ok(())
94}
95
96/// # Safety
97/// Caller is responsible for verifying the buffer.
98pub unsafe fn get_ref_flatbuffer_unchecked<'a, T>(buf: &'a [u8], index: usize) -> T::Inner
99where
100    T: Follow<'a> + 'a,
101{
102    // Safety: We have already verified the buffer in `check_flatbuffer`.
103    unsafe { <flatbuffers::ForwardsUOffset<T>>::follow(buf, index) }
104}
flatbuffers_util/ownedfb.rs

flatbuffers_util/
ownedfb.rs
