ps_datachunk/
lib.rs

1pub mod aligned;
2pub mod borrowed;
3pub mod encrypted;
4pub mod error;
5pub mod mbuf;
6pub mod owned;
7pub mod serialized;
8pub mod shared;
9pub mod typed;
10pub mod utils;
11pub use aligned::AlignedDataChunk;
12pub use borrowed::BorrowedDataChunk;
13pub use encrypted::EncryptedDataChunk;
14pub use error::PsDataChunkError;
15pub use error::Result;
16pub use mbuf::MbufDataChunk;
17pub use owned::OwnedDataChunk;
18pub use ps_hash::Hash;
19pub use ps_mbuf::Mbuf;
20pub use serialized::SerializedDataChunk;
21pub use shared::SharedDataChunk;
22pub use typed::ToDataChunk;
23pub use typed::ToTypedDataChunk;
24pub use typed::TypedDataChunk;
25
26use std::sync::Arc;
27
28/// represents any representation of a chunk of data
29pub trait DataChunkTrait {
30    fn data_ref(&self) -> &[u8];
31    fn hash_ref(&self) -> &[u8];
32
33    fn hash(&self) -> Arc<Hash> {
34        ps_hash::hash(self.data_ref()).into()
35    }
36
37    fn encrypt(&self) -> Result<EncryptedDataChunk> {
38        self.serialize()?.encrypt()
39    }
40
41    fn decrypt(&self, key: &[u8]) -> Result<SerializedDataChunk> {
42        utils::decrypt(self.data_ref(), key)
43    }
44
45    fn to_datachunk(&self) -> DataChunk {
46        DataChunk::Borrowed(self.borrow())
47    }
48
49    fn borrow(&self) -> BorrowedDataChunk {
50        BorrowedDataChunk::from_parts(self.data_ref(), self.hash())
51    }
52
53    fn to_owned(&self) -> OwnedDataChunk {
54        let data_ref = self.data_ref();
55        let reserved_size = utils::offsets(data_ref.len()).2;
56        let mut data_vec = Vec::with_capacity(reserved_size);
57
58        data_vec.extend_from_slice(data_ref);
59
60        OwnedDataChunk::from_parts(data_vec, self.hash())
61    }
62
63    fn serialize(&self) -> Result<SerializedDataChunk> {
64        SerializedDataChunk::from_parts(self.data_ref(), self.hash())
65    }
66
67    fn try_as<T: rkyv::Archive>(&self) -> Result<TypedDataChunk<T>>
68    where
69        T::Archived:
70            for<'a> rkyv::bytecheck::CheckBytes<rkyv::api::high::HighValidator<'a, rancor::Error>>,
71    {
72        match typed::check_byte_layout::<T>(self.data_ref()) {
73            true => Ok(unsafe { TypedDataChunk::from_chunk_unchecked(self.to_datachunk()) }),
74            false => Err(PsDataChunkError::TypeError),
75        }
76    }
77}
78
79/// represents a chunk of data that is either owned or pointed to
80pub enum DataChunk<'lt> {
81    Aligned(AlignedDataChunk),
82    Borrowed(BorrowedDataChunk<'lt>),
83    Mbuf(MbufDataChunk<'lt>),
84    Owned(OwnedDataChunk),
85    Serialized(SerializedDataChunk),
86    Shared(SharedDataChunk),
87}
88
89impl<'lt> DataChunkTrait for DataChunk<'lt> {
90    fn data_ref(&self) -> &[u8] {
91        self.data_ref()
92    }
93    fn hash_ref(&self) -> &[u8] {
94        self.hash_ref()
95    }
96    fn hash(&self) -> Arc<Hash> {
97        self.hash()
98    }
99}
100
101impl<'lt> From<AlignedDataChunk> for DataChunk<'lt> {
102    fn from(chunk: AlignedDataChunk) -> Self {
103        Self::Aligned(chunk)
104    }
105}
106
107impl<'lt> From<BorrowedDataChunk<'lt>> for DataChunk<'lt> {
108    fn from(chunk: BorrowedDataChunk<'lt>) -> Self {
109        Self::Borrowed(chunk)
110    }
111}
112
113impl<'lt> From<MbufDataChunk<'lt>> for DataChunk<'lt> {
114    fn from(chunk: MbufDataChunk<'lt>) -> Self {
115        Self::Mbuf(chunk)
116    }
117}
118
119impl<'lt> From<OwnedDataChunk> for DataChunk<'lt> {
120    fn from(chunk: OwnedDataChunk) -> Self {
121        Self::Owned(chunk)
122    }
123}
124
125impl<'lt> DataChunk<'lt> {
126    pub fn data_ref(&self) -> &[u8] {
127        match self {
128            Self::Borrowed(borrowed) => borrowed.data_ref(),
129            Self::Aligned(aligned) => aligned.data_ref(),
130            Self::Mbuf(mbuf) => mbuf.data_ref(),
131            Self::Owned(owned) => owned.data_ref(),
132            Self::Serialized(serialized) => serialized.data_ref(),
133            Self::Shared(shared) => shared.data_ref(),
134        }
135    }
136
137    pub fn hash_ref(&self) -> &[u8] {
138        match self {
139            Self::Aligned(aligned) => aligned.hash_ref(),
140            Self::Borrowed(borrowed) => borrowed.hash_ref(),
141            Self::Mbuf(mbuf) => mbuf.hash_ref(),
142            Self::Owned(owned) => owned.hash_ref(),
143            Self::Serialized(serialized) => serialized.hash_ref(),
144            Self::Shared(shared) => shared.hash_ref(),
145        }
146    }
147
148    pub fn hash(&self) -> Arc<Hash> {
149        match self {
150            Self::Aligned(aligned) => aligned.hash(),
151            Self::Borrowed(borrowed) => borrowed.hash(),
152            Self::Mbuf(mbuf) => mbuf.hash(),
153            Self::Owned(owned) => owned.hash(),
154            Self::Serialized(serialized) => serialized.hash(),
155            Self::Shared(shared) => shared.hash(),
156        }
157    }
158
159    #[inline(always)]
160    /// Decrypts this [DataChunk] with a given key.
161    pub fn decrypt(&self, key: &[u8]) -> Result<SerializedDataChunk> {
162        let decrypted = match self {
163            Self::Borrowed(borrowed) => borrowed.decrypt(key),
164            Self::Mbuf(mbuf) => mbuf.decrypt(key),
165            Self::Owned(chunk) => chunk.decrypt(key),
166            Self::Aligned(aligned) => aligned.decrypt(key),
167            Self::Serialized(serialized) => serialized.decrypt(key),
168            Self::Shared(shared) => shared.decrypt(key),
169        }?;
170
171        Ok(decrypted)
172    }
173
174    #[inline(always)]
175    /// Encrypts this [DataChunk].
176    pub fn encrypt(&self) -> Result<EncryptedDataChunk> {
177        match self {
178            DataChunk::Owned(owned) => owned.encrypt(),
179            DataChunk::Aligned(aligned) => {
180                OwnedDataChunk::encrypt_serialized_bytes(aligned.serialize()?.serialized_bytes())
181            }
182            DataChunk::Serialized(serialized) => serialized.encrypt(),
183            DataChunk::Shared(shared) => shared.encrypt(),
184            _ => self.serialize()?.encrypt(),
185        }
186    }
187
188    #[inline(always)]
189    /// Encrypts this [DataChunk] using `self.data` if owned.
190    pub fn encrypt_mut(&mut self) -> Result<EncryptedDataChunk> {
191        match self {
192            DataChunk::Owned(chunk) => chunk.encrypt_mut(),
193            _ => self.encrypt(),
194        }
195    }
196
197    pub fn guarantee_alignment<T>(self) -> Result<DataChunk<'lt>> {
198        let align_size = std::mem::align_of::<T>();
199        let remainder = self.data_ref().as_ptr() as usize % align_size;
200
201        let chunk = if remainder == 0 {
202            self
203        } else {
204            DataChunk::Serialized(self.serialize()?)
205        };
206
207        Ok(chunk)
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn test_encryption_decryption() -> Result<()> {
217        let original_data = "Neboť tak Bůh miluje svět, že dal [svého] jediného Syna, aby žádný, kdo v něho věří, nezahynul, ale měl život věčný. Vždyť Bůh neposlal [svého] Syna na svět, aby svět odsoudil, ale aby byl svět skrze něj zachráněn.".as_bytes().to_owned();
218
219        let data_chunk = DataChunk::Owned(OwnedDataChunk::from_data_ref(&original_data));
220
221        let encrypted_chunk = data_chunk.encrypt()?;
222        let decrypted_chunk = encrypted_chunk.decrypt()?;
223
224        assert_eq!(decrypted_chunk.data_ref(), original_data);
225
226        Ok(())
227    }
228
229    #[test]
230    fn test_serialization() -> Result<()> {
231        let original_data = vec![1, 2, 3, 4, 5];
232        let hash = ps_hash::hash(&original_data).into();
233        let owned_chunk = OwnedDataChunk::from_parts(original_data.to_vec(), hash);
234        let data_chunk = DataChunk::Owned(owned_chunk);
235
236        let serialized = data_chunk.serialize()?;
237        let deserialized = OwnedDataChunk::deserialize(&serialized.into_buffer())?;
238
239        assert_eq!(deserialized.data_ref(), original_data);
240
241        Ok(())
242    }
243}
ps_datachunk/lib.rs

ps_datachunk/
lib.rs
