Skip to main content

krabka_records_legacy/set/
decode.rs

1//! Decoding a v0/v1 `MessageSet`: walking the `(offset, size, message)`
2//! frames, and unwrapping the single compression layer the format allows.
3//!
4//! The recursion into a compressed wrapper is what makes this the larger half
5//! of the codec. Unwrapping is where nested compression is rejected, where the
6//! decompression policy is applied, and where the KIP-32 relative inner
7//! offsets of a v1 wrapper are rewritten back to absolute ones.
8
9use {
10    super::ParsedRecord,
11    crate::{
12        error::LegacyRecordsError,
13        message::{Magic, Message, compression_from_attrs},
14    },
15    bytes::Buf,
16    krabka_compression::{CompressionType, RecordDecompressionPolicy},
17    krabka_ids::Offset,
18    krabka_units::prelude::{ByteSize, ByteSizeExt as _},
19};
20
21/// The length of a wire slice as a byte count.
22///
23/// This function saturates and does not wrap. A `usize` above `u64::MAX`
24/// cannot occur on any target Krabka builds for.
25fn size_of_slice(slice: &[u8]) -> ByteSize {
26    ByteSize::from_bytes(u64::try_from(slice.len()).unwrap_or(u64::MAX))
27}
28
29/// Decode a flat, uncompressed `MessageSet` from `buf`.
30///
31/// This function consumes exactly `set_size_bytes` bytes from the buffer.
32/// It unwraps a compressed wrapper message at the top level one time. It
33/// rejects nested compression, that is, a compressed wrapper inside a
34/// compressed wrapper.
35/// # Errors
36/// Returns `Truncated` if `buf` holds fewer than `set_size_bytes` bytes.
37/// Returns `Truncated` if an entry header or an entry body stops early.
38/// Returns `NegativeLength` if an entry carries a negative `message_size`.
39/// Returns `NestedCompression` if a compressed wrapper holds another compressed wrapper.
40/// Returns the error from `Message::decode_from` for a malformed or corrupt message frame.
41/// Returns a compression error if the wrapper value does not decompress within the default policy.
42pub fn decode_message_set<B: Buf>(
43    buf: &mut B,
44    set_size_bytes: usize,
45) -> Result<Vec<ParsedRecord>, LegacyRecordsError> {
46    decode_message_set_with_policy(buf, set_size_bytes, RecordDecompressionPolicy::default())
47}
48
49/// Decode a legacy `MessageSet` with explicit decompression limits.
50///
51/// # Errors
52///
53/// Returns the legacy records error for malformed input, for truncated
54/// input, for corrupt input, for nested compression, or for input above
55/// the limits in `policy`.
56pub fn decode_message_set_with_policy<B: Buf>(
57    buf: &mut B,
58    set_size_bytes: usize,
59    policy: RecordDecompressionPolicy,
60) -> Result<Vec<ParsedRecord>, LegacyRecordsError> {
61    if buf.remaining() < set_size_bytes {
62        return Err(LegacyRecordsError::Truncated {
63            needed: set_size_bytes - buf.remaining(),
64        });
65    }
66    let mut region = vec![0u8; set_size_bytes];
67    buf.copy_to_slice(&mut region);
68    let mut out = Vec::new();
69    decode_into(
70        &region, &mut out, /* allow_compression = */ true, policy,
71    )?;
72    Ok(out)
73}
74
75fn decode_into(
76    bytes: &[u8],
77    out: &mut Vec<ParsedRecord>,
78    allow_compression: bool,
79    policy: RecordDecompressionPolicy,
80) -> Result<(), LegacyRecordsError> {
81    let mut cur = bytes;
82    while !cur.is_empty() {
83        if cur.len() < 12 {
84            return Err(LegacyRecordsError::Truncated {
85                needed: 12 - cur.len(),
86            });
87        }
88        let offset = cur.get_i64();
89        let size = cur.get_i32();
90        let Ok(size) = usize::try_from(size) else {
91            return Err(LegacyRecordsError::NegativeLength {
92                label: "message_size",
93                len: size,
94            });
95        };
96        if cur.remaining() < size {
97            return Err(LegacyRecordsError::Truncated {
98                needed: size - cur.remaining(),
99            });
100        }
101        let msg = Message::decode_from(&mut cur, size)?;
102        let codec = compression_from_attrs(msg.attributes)?;
103        if codec == CompressionType::None {
104            out.push(ParsedRecord {
105                offset: Offset(offset),
106                timestamp: msg.timestamp,
107                key: msg.key,
108                value: msg.value,
109            });
110        } else {
111            if !allow_compression {
112                return Err(LegacyRecordsError::NestedCompression);
113            }
114            let inner_compressed = msg.value.ok_or_else(|| {
115                LegacyRecordsError::Malformed("compressed wrapper has null value".into())
116            })?;
117            // Bound decompressed output to guard against a decompression bomb
118            // in a legacy compressed wrapper.
119            let max_output = policy.output_limit(size_of_slice(&inner_compressed));
120            let inner_bytes = krabka_compression::decompress(codec, &inner_compressed, max_output)?;
121
122            // Parse the inner set (no nested compression allowed).
123            let start_len = out.len();
124            decode_into(&inner_bytes, out, false, policy)?;
125
126            // KIP-32 anchors the last relative inner offset at the wrapper
127            // offset. Retained records can have gaps after compaction.
128            // An empty wrapper has no offsets to reconstruct.
129            if matches!(msg.magic, Magic::V1)
130                && let Some(inner) = out.get_mut(start_len..)
131                && let Some(last) = inner.last()
132            {
133                let base_abs = i128::from(offset) - i128::from(last.offset.0);
134                for rec in inner {
135                    let absolute = base_abs + i128::from(rec.offset.0);
136                    rec.offset = Offset(i64::try_from(absolute).map_err(|_| {
137                        LegacyRecordsError::Malformed("absolute inner offset exceeds i64".into())
138                    })?);
139                }
140            }
141        }
142    }
143    Ok(())
144}
145
146#[cfg(test)]
147mod tests {
148    use {
149        super::*,
150        crate::{
151            message::attrs_with_compression,
152            set::{
153                encode_compressed_message_set, encode_flat_message_set,
154                test_support::{compressed_bytes, decode_bytes, sample_records_v1},
155            },
156        },
157        bytes::{BufMut, Bytes, BytesMut},
158        krabka_compression::CompressionError,
159        krabka_units::prelude::{bytes, kibibytes},
160        std::iter,
161    };
162
163    fn record_at(offset: i64) -> ParsedRecord {
164        ParsedRecord {
165            offset: Offset(offset),
166            timestamp: Some(10),
167            key: None,
168            value: Some(Bytes::from_static(b"v")),
169        }
170    }
171
172    fn v1_wrapper(inner_offsets: &[i64], wrapper_offset: i64) -> BytesMut {
173        let inner: Vec<_> = inner_offsets.iter().copied().map(record_at).collect();
174        let mut flat = BytesMut::new();
175        encode_flat_message_set(inner, Magic::V1, &mut flat);
176        let wrapper = Message {
177            magic: Magic::V1,
178            attributes: attrs_with_compression(0, CompressionType::Gzip),
179            timestamp: Some(10),
180            key: None,
181            value: Some(krabka_compression::compress(CompressionType::Gzip, &flat).unwrap()),
182        };
183        let mut wire = BytesMut::new();
184        wire.put_i64(wrapper_offset);
185        wire.put_i32(i32::try_from(wrapper.encoded_len()).unwrap());
186        wrapper.encode_into(&mut wire);
187        wire
188    }
189
190    #[yare::parameterized(
191        first_sparse_gap = { [0, 2], 102, [100, 102] },
192        second_sparse_gap = { [3, 5], 102, [100, 102] },
193        extreme_sparse_gap = { [i64::MIN, i64::MIN + 2], i64::MAX, [i64::MAX - 2, i64::MAX] },
194    )]
195    fn sparse_v1_wrapper_anchors_last_inner_offset_without_closing_gaps(
196        inner: [i64; 2],
197        outer: i64,
198        absolute: [i64; 2],
199    ) {
200        // Construct the inner MessageSet independently of the compressed-set
201        // encoder: Kafka retains relative gaps when compaction removes a row.
202        let wire = v1_wrapper(&inner, outer);
203        let decoded = decode_bytes(&wire);
204        let expected: Vec<_> = absolute.into_iter().map(record_at).collect();
205        assert2::assert!(decoded == expected);
206    }
207
208    #[test]
209    fn v1_wrapper_rejects_unrepresentable_offsets() {
210        let wire = v1_wrapper(&[0, 1], i64::MIN);
211        assert2::assert!(matches!(
212            decode_message_set(&mut &wire[..], wire.len()),
213            Err(LegacyRecordsError::Malformed(_))
214        ));
215    }
216
217    #[yare::parameterized(
218        minimum = { i64::MIN },
219        maximum = { i64::MAX },
220    )]
221    fn empty_v1_wrapper_accepts_extreme_outer_offsets(outer: i64) {
222        let wire = v1_wrapper(&[], outer);
223        assert2::assert!(decode_bytes(&wire).is_empty());
224    }
225
226    #[test]
227    fn slice_length_lifts_to_a_byte_count() {
228        assert2::check!(size_of_slice(&[]) == ByteSize::ZERO);
229        assert2::check!(size_of_slice(&[0u8; 4096]) == kibibytes(4));
230    }
231
232    krabka_macros::legacy_policy_fixture!(policy_wire, crate);
233
234    #[test]
235    fn decompression_policy_limits_legacy_decode() {
236        let wire = policy_wire();
237
238        decode_message_set(&mut &wire[..], wire.len()).unwrap();
239
240        let policy =
241            RecordDecompressionPolicy::new(krabka_units::fraction(1.0), bytes(1), bytes(32))
242                .unwrap();
243        assert2::assert!(matches!(
244            decode_message_set_with_policy(&mut &wire[..], wire.len(), policy),
245            Err(LegacyRecordsError::Compression(
246                CompressionError::TooLarge { limit: 32 }
247            ))
248        ));
249    }
250
251    #[test]
252    fn rejects_nested_compression() {
253        // Construct a wrapper containing a wrapper.
254        let inner_recs = sample_records_v1();
255        let mut inner_outer = BytesMut::new();
256        encode_compressed_message_set(
257            &inner_recs,
258            Magic::V1,
259            CompressionType::Gzip,
260            &mut inner_outer,
261        )
262        .unwrap();
263
264        // Now wrap that bytestream as the value of another compressed wrapper.
265        let outer_compressed =
266            krabka_compression::compress(CompressionType::Gzip, &inner_outer).unwrap();
267        let outer_msg = Message {
268            magic: Magic::V1,
269            attributes: attrs_with_compression(0, CompressionType::Gzip),
270            timestamp: Some(0),
271            key: None,
272            value: Some(outer_compressed),
273        };
274        let mut wire = BytesMut::new();
275        let outer_len = outer_msg.encoded_len();
276        wire.put_i64(0);
277        wire.put_i32(i32::try_from(outer_len).unwrap());
278        outer_msg.encode_into(&mut wire);
279
280        let mut cur: &[u8] = &wire[..];
281        let err = decode_message_set(&mut cur, wire.len()).unwrap_err();
282        assert2::assert!(matches!(err, LegacyRecordsError::NestedCompression));
283    }
284
285    // --- mutation-coverage tests --------------------------------------------
286    //
287    // Round-trips above don't exercise the malformed/boundary paths or pin
288    // exact framing. These do: precise `needed` counts, error-variant
289    // boundaries, the decompression-cap floor, and the v1 inner-offset rewrite.
290    //
291
292    #[test]
293    fn decode_message_set_short_buffer_reports_needed() {
294        // 4 bytes available, caller claims a 12-byte set: needed = 12 - 4 = 8.
295        let data = [0u8; 4];
296        let mut cur: &[u8] = &data;
297        let err = decode_message_set(&mut cur, 12).unwrap_err();
298        assert2::assert!(matches!(err, LegacyRecordsError::Truncated { needed: 8 }));
299    }
300
301    #[test]
302    fn entry_header_truncated_reports_needed() {
303        // 8 bytes where a 12-byte (offset+size) entry header is required:
304        // needed = 12 - 8 = 4.
305        let data = [0u8; 8];
306        let mut cur: &[u8] = &data;
307        let err = decode_message_set(&mut cur, 8).unwrap_err();
308        assert2::assert!(matches!(err, LegacyRecordsError::Truncated { needed: 4 }));
309    }
310
311    fn entry_error(size: i32, trailing: &[u8]) -> LegacyRecordsError {
312        let mut data = BytesMut::new();
313        data.put_i64(0);
314        data.put_i32(size);
315        data.put_slice(trailing);
316        decode_message_set(&mut &data[..], data.len()).unwrap_err()
317    }
318
319    #[test]
320    fn entry_zero_message_size_is_malformed() {
321        // offset(8) + size(0): clears the `< 12` and `size < 0` guards, then
322        // Message::decode_from rejects the 0-byte frame as Malformed (< 6).
323        // Distinguishes the `<` boundaries from `<=`/`==`.
324        let err = entry_error(0, &[]);
325        assert2::assert!(matches!(err, LegacyRecordsError::Malformed(_)));
326    }
327
328    #[test]
329    fn entry_negative_message_size_rejected() {
330        let err = entry_error(-1, &[0u8; 4]); // keep region >= 12 bytes
331        assert2::assert!(matches!(err, LegacyRecordsError::NegativeLength {
332            label: "message_size",
333            len: -1
334        }));
335    }
336
337    #[test]
338    fn entry_message_body_truncated_reports_needed() {
339        // Entry claims a 10-byte message but only 2 bytes follow:
340        // needed = 10 - 2 = 8.
341        let err = entry_error(10, &[0u8; 2]);
342        assert2::assert!(matches!(err, LegacyRecordsError::Truncated { needed: 8 }));
343    }
344
345    #[test]
346    fn compressed_wrapper_allows_large_decompressed_output() {
347        // The 16 MiB decompression-cap floor must let a ~2 MiB wrapper through
348        // even though the compressed size is tiny. Shrinking the floor (a `*`
349        // flip in `16 * 1024 * 1024`) would reject this round-trip as TooLarge.
350        let big = vec![0x7Eu8; 2 * 1024 * 1024];
351        let recs = vec![ParsedRecord {
352            offset: Offset(0),
353            timestamp: Some(5),
354            key: None,
355            value: Some(Bytes::from(big)),
356        }];
357        let buf = compressed_bytes(&recs, Magic::V1, CompressionType::Gzip);
358        let decoded = decode_bytes(&buf);
359        assert2::assert!(decoded == recs);
360    }
361
362    #[test]
363    fn v1_inner_offset_rewrite_leaves_preceding_flat_records_unchanged() {
364        // A set with a flat record FOLLOWED by a compressed v1 wrapper: when the
365        // wrapper is decoded, `out` already holds the flat record (start_len >
366        // 0). Offset reconstruction applies only to the wrapper's inner set.
367        let mut buf = BytesMut::new();
368        let flat = ParsedRecord {
369            offset: Offset(50),
370            timestamp: Some(1),
371            key: None,
372            value: Some(Bytes::from_static(b"flat")),
373        };
374        encode_flat_message_set(vec![flat.clone()], Magic::V1, &mut buf);
375        let inner = vec![
376            ParsedRecord {
377                offset: Offset(100),
378                timestamp: Some(2),
379                key: None,
380                value: Some(Bytes::from_static(b"x")),
381            },
382            ParsedRecord {
383                offset: Offset(101),
384                timestamp: Some(3),
385                key: None,
386                value: Some(Bytes::from_static(b"y")),
387            },
388        ];
389        encode_compressed_message_set(&inner, Magic::V1, CompressionType::Gzip, &mut buf).unwrap();
390
391        let decoded = decode_bytes(&buf);
392        let expected = iter::once(flat).chain(inner).collect::<Vec<_>>();
393        assert2::assert!(decoded == expected);
394    }
395}