weavatrix_memory/codec/
compression.rs

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

weavatrix_memory/codec/
compression.rs
