Files
Stalwart/crates/store/src/write/serialize.rs
2026-03-12 18:26:35 +01:00

626 lines
21 KiB
Rust

/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use super::{ARCHIVE_ALIGNMENT, AlignedBytes, Archive, ArchiveVersion, Archiver};
use crate::{Deserialize, Serialize, SerializeInfallible, U32_LEN, U64_LEN, Value};
use compact_str::format_compact;
use rkyv::util::AlignedVec;
use roaring::{RoaringBitmap, RoaringTreemap};
const MAGIC_MARKER: u8 = 1 << 7;
const VERSIONED: u8 = 1 << 6;
const HASHED: u8 = 1 << 5;
const LZ4_COMPRESSED: u8 = 1 << 4;
const COMPRESS_WATERMARK: usize = 8192;
fn validate_marker_and_contents(bytes: &[u8]) -> Option<(bool, &[u8], ArchiveVersion)> {
let (marker, contents) = bytes
.split_last()
.filter(|(marker, _)| (**marker & MAGIC_MARKER) != 0)?;
let is_uncompressed = (marker & LZ4_COMPRESSED) == 0;
if marker & VERSIONED != 0 {
let (contents, change_id) = contents
.split_at_checked(contents.len() - U64_LEN)
.and_then(|(contents, change_id)| {
change_id
.try_into()
.ok()
.map(|change_id| (contents, u64::from_be_bytes(change_id)))
})?;
contents
.split_at_checked(contents.len() - U32_LEN)
.and_then(|(contents, archive_hash)| {
let hash = xxhash_rust::xxh3::xxh3_64(contents) as u32;
if hash.to_be_bytes().as_slice() == archive_hash {
Some((
is_uncompressed,
contents,
ArchiveVersion::Versioned { change_id, hash },
))
} else {
None
}
})
} else if marker & HASHED != 0 {
contents
.split_at_checked(contents.len() - U32_LEN)
.and_then(|(contents, archive_hash)| {
let hash = xxhash_rust::xxh3::xxh3_64(contents) as u32;
if hash.to_be_bytes().as_slice() == archive_hash {
Some((is_uncompressed, contents, ArchiveVersion::Hashed { hash }))
} else {
None
}
})
} else {
Some((is_uncompressed, contents, ArchiveVersion::Unversioned))
}
}
impl Deserialize for Archive<AlignedBytes> {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
let (is_uncompressed, contents, version) =
validate_marker_and_contents(bytes).ok_or_else(|| {
trc::StoreEvent::DataCorruption
.into_err()
.details("Archive integrity compromised")
.ctx(trc::Key::Value, bytes)
.caused_by(trc::location!())
})?;
if is_uncompressed {
let mut bytes = AlignedVec::with_capacity(contents.len());
bytes.extend_from_slice(contents);
Ok(Archive {
version,
inner: AlignedBytes::Aligned(bytes),
})
} else {
aligned_lz4_deflate(contents).map(|inner| Archive { version, inner })
}
}
fn deserialize_owned(mut bytes: Vec<u8>) -> trc::Result<Self> {
let (is_uncompressed, contents, version) = validate_marker_and_contents(&bytes)
.ok_or_else(|| {
trc::StoreEvent::DataCorruption
.into_err()
.details("Archive integrity compromised")
.ctx(trc::Key::Value, bytes.as_slice())
.caused_by(trc::location!())
})?;
if is_uncompressed {
bytes.truncate(contents.len());
if bytes.as_ptr().addr() & (ARCHIVE_ALIGNMENT - 1) == 0 {
Ok(Archive {
version,
inner: AlignedBytes::Vec(bytes),
})
} else {
let mut aligned = AlignedVec::with_capacity(bytes.len());
aligned.extend_from_slice(&bytes);
Ok(Archive {
version,
inner: AlignedBytes::Aligned(aligned),
})
}
} else {
aligned_lz4_deflate(contents).map(|inner| Archive { version, inner })
}
}
}
#[inline]
fn aligned_lz4_deflate(archive: &[u8]) -> trc::Result<AlignedBytes> {
lz4_flex::block::uncompressed_size(archive)
.and_then(|(uncompressed_size, archive)| {
let mut bytes = AlignedVec::with_capacity(uncompressed_size);
unsafe {
// SAFETY: `new_len` is equal to `capacity` and vector is initialized by lz4_flex.
bytes.set_len(uncompressed_size);
}
lz4_flex::decompress_into(archive, &mut bytes)?;
Ok(AlignedBytes::Aligned(bytes))
})
.map_err(|err| {
trc::StoreEvent::DecompressError
.ctx(trc::Key::Value, archive)
.caused_by(trc::location!())
.reason(err)
})
}
impl<T> Serialize for Archiver<T>
where
T: rkyv::Archive
+ for<'a> rkyv::Serialize<
rkyv::api::high::HighSerializer<
rkyv::util::AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::rancor::Error,
>,
>,
{
fn serialize(&self) -> trc::Result<Vec<u8>> {
rkyv::to_bytes::<rkyv::rancor::Error>(&self.inner)
.map_err(|err| {
trc::StoreEvent::DeserializeError
.caused_by(trc::location!())
.reason(err)
})
.map(|input| {
let input = input.as_ref();
let input_len = input.len();
let version_offset = ((self.flags & VERSIONED != 0) as usize) * U64_LEN;
let mut bytes = if input_len > COMPRESS_WATERMARK {
let mut bytes = vec![
self.flags | LZ4_COMPRESSED;
lz4_flex::block::get_maximum_output_size(input_len)
+ (U32_LEN * 2)
+ version_offset
+ 1
];
// Compress the data
let compressed_len =
lz4_flex::compress_into(input, &mut bytes[U32_LEN..]).unwrap();
if compressed_len < input_len {
// Prepend the length of the uncompressed data
bytes[..U32_LEN].copy_from_slice(&(input_len as u32).to_le_bytes());
if self.flags & HASHED != 0 {
// Hash the compressed data including the length
let hash =
xxhash_rust::xxh3::xxh3_64(&bytes[..compressed_len + U32_LEN])
as u32;
// Add the hash
bytes[compressed_len + U32_LEN..compressed_len + (U32_LEN * 2)]
.copy_from_slice(&hash.to_be_bytes());
// Truncate to the actual size
bytes.truncate(compressed_len + (U32_LEN * 2) + version_offset + 1);
} else {
// Truncate to the actual size
bytes.truncate(compressed_len + U32_LEN + 1);
}
return bytes;
}
bytes.clear();
bytes
} else {
Vec::with_capacity(input_len + U32_LEN + version_offset + 1)
};
bytes.extend_from_slice(input);
if self.flags & HASHED != 0 {
bytes.extend_from_slice(
&(xxhash_rust::xxh3::xxh3_64(input) as u32).to_be_bytes(),
);
}
if version_offset != 0 {
bytes.extend_from_slice(0u64.to_be_bytes().as_slice());
}
bytes.push(self.flags);
bytes
})
}
}
impl Archive<AlignedBytes> {
#[inline]
pub fn as_bytes(&self) -> &[u8] {
match &self.inner {
AlignedBytes::Vec(bytes) => bytes.as_slice(),
AlignedBytes::Aligned(bytes) => bytes.as_slice(),
}
}
pub fn unarchive<T>(&self) -> trc::Result<&<T as rkyv::Archive>::Archived>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
let bytes = self.as_bytes();
if self.version != ArchiveVersion::Unversioned {
if bytes.len() >= std::mem::size_of::<T::Archived>() {
// SAFETY: Trusted input with integrity hash
Ok(unsafe { rkyv::access_unchecked::<T::Archived>(bytes) })
} else {
Err(trc::StoreEvent::DataCorruption
.into_err()
.details(format_compact!(
"Archive size mismatch, expected {} bytes but got {} bytes.",
std::mem::size_of::<T::Archived>(),
bytes.len()
))
.ctx(trc::Key::Value, bytes)
.caused_by(trc::location!()))
}
} else {
rkyv::access::<T::Archived, rkyv::rancor::Error>(bytes).map_err(|err| {
trc::StoreEvent::DeserializeError
.ctx(trc::Key::Value, self.as_bytes())
.details("Archive access failed")
.caused_by(trc::location!())
.reason(err)
})
}
}
pub fn unarchive_untrusted<T>(&self) -> trc::Result<&<T as rkyv::Archive>::Archived>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
let bytes = self.as_bytes();
if bytes.len() >= std::mem::size_of::<T::Archived>() {
rkyv::access::<T::Archived, rkyv::rancor::Error>(bytes).map_err(|err| {
trc::StoreEvent::DeserializeError
.ctx(trc::Key::Value, self.as_bytes())
.details("Archive access failed")
.caused_by(trc::location!())
.reason(err)
})
} else {
Err(trc::StoreEvent::DataCorruption
.into_err()
.details(format_compact!(
"Archive size mismatch, expected {} bytes but got {} bytes.",
std::mem::size_of::<T::Archived>(),
bytes.len()
))
.ctx(trc::Key::Value, bytes)
.caused_by(trc::location!()))
}
}
pub fn deserialize<T>(&self) -> trc::Result<T>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
self.unarchive::<T>().and_then(|input| {
rkyv::deserialize(input).map_err(|err| {
trc::StoreEvent::DeserializeError
.ctx(trc::Key::Value, self.as_bytes())
.caused_by(trc::location!())
.reason(err)
})
})
}
pub fn deserialize_untrusted<T>(&self) -> trc::Result<T>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
self.unarchive_untrusted::<T>().and_then(|input| {
rkyv::deserialize(input).map_err(|err| {
trc::StoreEvent::DeserializeError
.ctx(trc::Key::Value, self.as_bytes())
.caused_by(trc::location!())
.reason(err)
})
})
}
pub fn to_unarchived<T>(&self) -> trc::Result<Archive<&<T as rkyv::Archive>::Archived>>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
self.unarchive::<T>().map(|inner| Archive {
version: self.version,
inner,
})
}
pub fn into_deserialized<T>(&self) -> trc::Result<Archive<T>>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
self.deserialize::<T>().map(|inner| Archive {
version: self.version,
inner,
})
}
pub fn into_inner(self) -> Vec<u8> {
let mut bytes = match self.inner {
AlignedBytes::Vec(bytes) => bytes,
AlignedBytes::Aligned(bytes) => bytes.to_vec(),
};
match self.version {
ArchiveVersion::Versioned { change_id, hash } => {
bytes.extend_from_slice(&change_id.to_be_bytes());
bytes.extend_from_slice(&hash.to_be_bytes());
bytes.push(MAGIC_MARKER | VERSIONED | HASHED);
}
ArchiveVersion::Hashed { hash } => {
bytes.extend_from_slice(&hash.to_be_bytes());
bytes.push(MAGIC_MARKER | HASHED);
}
ArchiveVersion::Unversioned => {
bytes.push(MAGIC_MARKER);
}
}
bytes
}
pub fn extract_hash(bytes: &[u8]) -> Option<u32> {
let marker = *bytes.last()?;
if marker & VERSIONED != 0 {
bytes
.get(bytes.len() - U32_LEN - U64_LEN - 1..bytes.len() - U64_LEN - 1)
.and_then(|slice| slice.try_into().ok().map(u32::from_be_bytes))
} else if marker & HASHED != 0 {
bytes
.get(bytes.len() - U32_LEN - 1..bytes.len() - 1)
.and_then(|slice| slice.try_into().ok().map(u32::from_be_bytes))
} else {
None
}
}
}
impl<T> Archiver<T>
where
T: rkyv::Archive
+ for<'a> rkyv::Serialize<
rkyv::api::high::HighSerializer<
rkyv::util::AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::rancor::Error,
>,
>,
{
pub fn new(inner: T) -> Self {
Self {
inner,
flags: MAGIC_MARKER | HASHED,
}
}
pub fn into_inner(self) -> T {
self.inner
}
pub fn with_version(self) -> Self {
Self {
inner: self.inner,
flags: self.flags | VERSIONED,
}
}
pub fn untrusted(self) -> Self {
Self {
inner: self.inner,
flags: MAGIC_MARKER,
}
}
pub fn serialize_versioned(self) -> trc::Result<(u64, Vec<u8>)> {
self.with_version()
.serialize()
.map(|bytes| ((bytes.len() - U64_LEN - 1) as u64, bytes))
}
}
impl<T> Archive<&T>
where
T: rkyv::Portable
+ for<'a> rkyv::bytecheck::CheckBytes<rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>>
+ Sync
+ Send,
{
pub fn to_deserialized<V>(&self) -> trc::Result<Archive<V>>
where
T: rkyv::Deserialize<V, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
rkyv::deserialize::<V, rkyv::rancor::Error>(self.inner)
.map_err(|err| {
trc::StoreEvent::DeserializeError
.caused_by(trc::location!())
.reason(err)
})
.map(|inner| Archive {
version: self.version,
inner,
})
}
pub fn deserialize<V>(&self) -> trc::Result<V>
where
T: rkyv::Deserialize<V, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
rkyv::deserialize::<V, rkyv::rancor::Error>(self.inner).map_err(|err| {
trc::StoreEvent::DeserializeError
.caused_by(trc::location!())
.reason(err)
})
}
}
#[inline]
pub fn rkyv_deserialize<T, V>(input: &T) -> trc::Result<V>
where
T: rkyv::Portable
+ for<'a> rkyv::bytecheck::CheckBytes<rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>>
+ Sync
+ Send
+ rkyv::Deserialize<V, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
rkyv::deserialize::<V, rkyv::rancor::Error>(input).map_err(|err| {
trc::StoreEvent::DeserializeError
.caused_by(trc::location!())
.reason(err)
})
}
pub fn rkyv_unarchive<T>(input: &[u8]) -> trc::Result<&<T as rkyv::Archive>::Archived>
where
T: rkyv::Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>>
+ rkyv::Deserialize<T, rkyv::api::high::HighDeserializer<rkyv::rancor::Error>>,
{
rkyv::access::<T::Archived, rkyv::rancor::Error>(input).map_err(|err| {
trc::StoreEvent::DataCorruption
.caused_by(trc::location!())
.ctx(trc::Key::Value, input)
.reason(err)
})
}
impl SerializeInfallible for u32 {
fn serialize(&self) -> Vec<u8> {
self.to_be_bytes().to_vec()
}
}
impl SerializeInfallible for u64 {
fn serialize(&self) -> Vec<u8> {
self.to_be_bytes().to_vec()
}
}
impl SerializeInfallible for i64 {
fn serialize(&self) -> Vec<u8> {
self.to_be_bytes().to_vec()
}
}
impl SerializeInfallible for u16 {
fn serialize(&self) -> Vec<u8> {
self.to_be_bytes().to_vec()
}
}
impl SerializeInfallible for f64 {
fn serialize(&self) -> Vec<u8> {
self.to_be_bytes().to_vec()
}
}
impl SerializeInfallible for &str {
fn serialize(&self) -> Vec<u8> {
self.as_bytes().to_vec()
}
}
impl Deserialize for String {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
Ok(String::from_utf8_lossy(bytes).into_owned())
}
fn deserialize_owned(bytes: Vec<u8>) -> trc::Result<Self> {
Ok(String::from_utf8(bytes)
.unwrap_or_else(|err| String::from_utf8_lossy(err.as_bytes()).into_owned()))
}
}
impl Deserialize for u64 {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
Ok(u64::from_be_bytes(bytes.try_into().map_err(|_| {
trc::StoreEvent::DataCorruption.caused_by(trc::location!())
})?))
}
}
impl Deserialize for i64 {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
Ok(i64::from_be_bytes(bytes.try_into().map_err(|_| {
trc::StoreEvent::DataCorruption.caused_by(trc::location!())
})?))
}
}
impl Deserialize for u32 {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
Ok(u32::from_be_bytes(bytes.try_into().map_err(|_| {
trc::StoreEvent::DataCorruption.caused_by(trc::location!())
})?))
}
}
impl<T> From<Value<'static>> for Archive<T> {
fn from(_: Value<'static>) -> Self {
unimplemented!()
}
}
impl Default for Archive<AlignedBytes> {
fn default() -> Self {
Archive {
version: ArchiveVersion::Unversioned,
inner: AlignedBytes::Aligned(AlignedVec::new()),
}
}
}
impl Serialize for RoaringBitmap {
fn serialize(&self) -> trc::Result<Vec<u8>> {
let mut bytes = Vec::with_capacity(self.serialized_size());
self.serialize_into(&mut bytes)
.map_err(|err| {
trc::StoreEvent::UnexpectedError
.caused_by(trc::location!())
.reason(err)
})
.map(|_| bytes)
}
}
impl Deserialize for RoaringBitmap {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
RoaringBitmap::deserialize_from(bytes).map_err(|err| {
trc::StoreEvent::DeserializeError
.caused_by(trc::location!())
.reason(err)
})
}
}
impl Serialize for RoaringTreemap {
fn serialize(&self) -> trc::Result<Vec<u8>> {
let mut bytes = Vec::with_capacity(self.serialized_size());
self.serialize_into(&mut bytes)
.map_err(|err| {
trc::StoreEvent::UnexpectedError
.caused_by(trc::location!())
.reason(err)
})
.map(|_| bytes)
}
}
impl Deserialize for RoaringTreemap {
fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
RoaringTreemap::deserialize_from(bytes).map_err(|err| {
trc::StoreEvent::DeserializeError
.caused_by(trc::location!())
.reason(err)
})
}
}