weavatrix_memory/codec/
compression.rs

1use super::Codec;
2use crate::error::{MemoryError, Result};
3
4const HEADER: &[u8; 8] = b"WMEMLZ01";
5const HEADER_LEN: usize = 25;
6const MODE_RAW: u8 = 0;
7const MODE_LZ4: u8 = 1;
8
9/// Size-bounded LZ4 wrapper for any existing codec.
10#[derive(Debug, Clone)]
11pub struct Lz4Codec<C> {
12    inner: C,
13    max_decoded_bytes: usize,
14}
15
16impl<C> Lz4Codec<C> {
17    /// Wraps a codec and caps the allocation accepted during decode.
18    ///
19    /// # Errors
20    ///
21    /// Rejects a zero decoded-size limit.
22    pub fn new(inner: C, max_decoded_bytes: usize) -> Result<Self> {
23        if max_decoded_bytes == 0 {
24            return Err(invalid("must be greater than zero"));
25        }
26        Ok(Self {
27            inner,
28            max_decoded_bytes,
29        })
30    }
31
32    #[must_use]
33    pub fn inner(&self) -> &C {
34        &self.inner
35    }
36}
37
38impl<T, C> Codec<T> for Lz4Codec<C>
39where
40    C: Codec<T>,
41{
42    fn encode(&self, value: &T) -> Result<Vec<u8>> {
43        let raw = self.inner.encode(value)?;
44        if raw.len() > self.max_decoded_bytes {
45            return Err(invalid("encoded value exceeds max_decoded_bytes"));
46        }
47        let capacity = HEADER_LEN
48            .checked_add(lz4_flex::block::get_maximum_output_size(raw.len()))
49            .ok_or(MemoryError::CapacityOverflow)?;
50        let mut output = vec![0; capacity];
51        let compressed_len = lz4_flex::block::compress_into(&raw, &mut output[HEADER_LEN..])
52            .map_err(|_| codec("LZ4 output capacity was insufficient"))?;
53        let (mode, payload_len) = if compressed_len < raw.len() {
54            (MODE_LZ4, compressed_len)
55        } else {
56            output[HEADER_LEN..HEADER_LEN + raw.len()].copy_from_slice(&raw);
57            (MODE_RAW, raw.len())
58        };
59        let raw_len = u64::try_from(raw.len()).map_err(|_| MemoryError::CapacityOverflow)?;
60        let stored_len = u64::try_from(payload_len).map_err(|_| MemoryError::CapacityOverflow)?;
61        output[..8].copy_from_slice(HEADER);
62        output[8] = mode;
63        output[9..17].copy_from_slice(&raw_len.to_le_bytes());
64        output[17..25].copy_from_slice(&stored_len.to_le_bytes());
65        output.truncate(HEADER_LEN + payload_len);
66        Ok(output)
67    }
68
69    fn decode(&self, bytes: &[u8]) -> Result<T> {
70        if bytes.len() < HEADER_LEN || &bytes[..8] != HEADER {
71            return Err(codec("unsupported LZ4 envelope"));
72        }
73        let raw_len = usize::try_from(u64::from_le_bytes(bytes[9..17].try_into().unwrap()))
74            .map_err(|_| codec("decoded length exceeds platform capacity"))?;
75        if raw_len > self.max_decoded_bytes {
76            return Err(codec("decoded value exceeds configured size limit"));
77        }
78        let stored_len = usize::try_from(u64::from_le_bytes(bytes[17..25].try_into().unwrap()))
79            .map_err(|_| codec("stored length exceeds platform capacity"))?;
80        let envelope_len = HEADER_LEN
81            .checked_add(stored_len)
82            .ok_or(MemoryError::CapacityOverflow)?;
83        if bytes.len() != envelope_len {
84            return Err(codec("compressed envelope length mismatch"));
85        }
86        let payload = &bytes[HEADER_LEN..];
87        let raw = match bytes[8] {
88            MODE_RAW if payload.len() == raw_len => payload.to_vec(),
89            MODE_RAW => return Err(codec("raw payload length mismatch")),
90            MODE_LZ4 => {
91                let mut raw = vec![0; raw_len];
92                let written = lz4_flex::block::decompress_into(payload, &mut raw)
93                    .map_err(|_| codec("invalid LZ4 payload"))?;
94                if written != raw_len {
95                    return Err(codec("decoded payload length mismatch"));
96                }
97                raw
98            }
99            _ => return Err(codec("unsupported compression mode")),
100        };
101        if raw.len() != raw_len {
102            return Err(codec("decoded payload length mismatch"));
103        }
104        self.inner.decode(&raw)
105    }
106}
107
108fn codec(message: &str) -> MemoryError {
109    MemoryError::Codec {
110        message: message.to_owned(),
111    }
112}
113
114fn invalid(reason: &'static str) -> MemoryError {
115    MemoryError::InvalidValue {
116        field: "max_decoded_bytes",
117        reason,
118    }
119}
weavatrix_memory/codec/compression.rs

weavatrix_memory/codec/
compression.rs
