Skip to main content

krabka_protocol/records/
owned.rs

1//! Owned `RecordBatch`, `Record`, and `RecordHeader` types.
2
3use bytes::{Buf, BufMut, Bytes, BytesMut};
4use krabka_compression::RecordDecompressionPolicy;
5use krabka_units::{ByteSize, convert::ByteSizeExt as _};
6use zerocopy::FromBytes as _;
7
8use crate::{
9    primitives::varint::{
10        get_varint, get_varlong, put_varint, put_varlong, varint_len, varlong_len,
11    },
12    records::{
13        RecordsError,
14        borrowed::parse_one_record,
15        crc::{crc32c, crc32c_append},
16        header::{Attributes, CRC_COVERAGE_START, HEADER_LEN, RecordBatchHeader},
17    },
18};
19
20#[derive(Debug, Clone, PartialEq, Eq, Default)]
21pub struct RecordHeader {
22    pub key: String,
23    pub value: Option<Bytes>,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Default)]
27pub struct Record {
28    pub attributes: i8,
29    pub timestamp_delta: i64,
30    pub offset_delta: i32,
31    pub key: Option<Bytes>,
32    pub value: Option<Bytes>,
33    pub headers: Vec<RecordHeader>,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct RecordBatch {
38    pub base_offset: i64,
39    pub partition_leader_epoch: i32,
40    pub attributes: Attributes,
41    pub last_offset_delta: i32,
42    pub base_timestamp: i64,
43    pub max_timestamp: i64,
44    pub producer_id: i64,
45    pub producer_epoch: i16,
46    pub base_sequence: i32,
47    pub records: Vec<Record>,
48}
49
50impl Default for RecordBatch {
51    fn default() -> Self {
52        Self {
53            base_offset: 0,
54            partition_leader_epoch: 0,
55            attributes: Attributes::default(),
56            last_offset_delta: 0,
57            base_timestamp: 0,
58            max_timestamp: 0,
59            producer_id: -1, // sentinel: non-idempotent
60            producer_epoch: -1,
61            base_sequence: -1,
62            records: Vec::new(),
63        }
64    }
65}
66
67impl Record {
68    /// Encode a single record into `buf` as a varlong length prefix and fields.
69    ///
70    /// # Errors
71    ///
72    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
73    pub fn encode<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
74        let body_len = self.body_len();
75        put_varlong(
76            buf,
77            i64::try_from(body_len)
78                .map_err(|_| RecordsError::RecordParse("record body length overflow".into()))?,
79        );
80        self.encode_body(buf)
81    }
82
83    /// Predicted total length of this record on the wire, length-prefix plus body.
84    pub fn encoded_len(&self) -> usize {
85        let body = self.body_len();
86        // `encode` rejects a body longer than `i64::MAX`; the prediction
87        // saturates there.
88        let body_i64 = i64::try_from(body).unwrap_or(i64::MAX);
89        varlong_len(body_i64) + body
90    }
91
92    /// Size of this record's body in bytes: what Kafka's `DefaultRecord` calls
93    /// `sizeOfBodyInBytes`, and what its varlong length prefix declares.
94    ///
95    /// The body is the attributes byte, the two deltas, and the key, the value
96    /// and the headers, each behind its own length prefix. A null key, value or
97    /// header value counts as the one-byte `-1` marker. The length prefix
98    /// itself is not counted: [`Record::encoded_len`] is the prefix and this.
99    #[must_use]
100    pub fn body_len(&self) -> usize {
101        body_len_of(
102            self.timestamp_delta,
103            self.offset_delta,
104            self.key.as_ref().map(Bytes::len),
105            self.value.as_ref().map(Bytes::len),
106            self.headers
107                .iter()
108                .map(|h| (h.key.len(), h.value.as_ref().map(Bytes::len))),
109        )
110    }
111
112    fn encode_body<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
113        // Metadata needs at most 21 bytes. Include short keys in the same write.
114        const SMALL_KEY_LEN: usize = 32;
115        const PREFIX_LEN: usize = 21 + SMALL_KEY_LEN;
116        let mut prefix = [0_u8; PREFIX_LEN];
117        let mut fields = prefix.as_mut_slice();
118        fields.put_i8(self.attributes);
119        put_varlong(&mut fields, self.timestamp_delta);
120        put_varint(&mut fields, self.offset_delta);
121        match &self.key {
122            None => put_varint(&mut fields, -1),
123            Some(k) => {
124                let Ok(len) = i32::try_from(k.len()) else {
125                    // Preserve the fields already written before an invalid key.
126                    let unused = fields.len();
127                    buf.put_slice(written_prefix(&prefix, unused)?);
128                    return Err(RecordsError::RecordParse(
129                        "record key length overflow".into(),
130                    ));
131                };
132                put_varint(&mut fields, len);
133            }
134        }
135        if let Some(key) = &self.key
136            && key.len() <= SMALL_KEY_LEN
137        {
138            fields.put_slice(key);
139        }
140        let unused = fields.len();
141        buf.put_slice(written_prefix(&prefix, unused)?);
142        if let Some(key) = &self.key
143            && key.len() > SMALL_KEY_LEN
144        {
145            buf.put_slice(key);
146        }
147        match &self.value {
148            None => put_varint(buf, -1),
149            Some(v) => {
150                put_varint(
151                    buf,
152                    i32::try_from(v.len()).map_err(|_| {
153                        RecordsError::RecordParse("record value length overflow".into())
154                    })?,
155                );
156                buf.put_slice(v);
157            }
158        }
159        put_varint(
160            buf,
161            i32::try_from(self.headers.len())
162                .map_err(|_| RecordsError::RecordParse("record header count overflow".into()))?,
163        );
164        for h in &self.headers {
165            let key_bytes = h.key.as_bytes();
166            put_varint(
167                buf,
168                i32::try_from(key_bytes.len())
169                    .map_err(|_| RecordsError::RecordParse("header key length overflow".into()))?,
170            );
171            buf.put_slice(key_bytes);
172            match &h.value {
173                None => put_varint(buf, -1),
174                Some(v) => {
175                    put_varint(
176                        buf,
177                        i32::try_from(v.len()).map_err(|_| {
178                            RecordsError::RecordParse("header value length overflow".into())
179                        })?,
180                    );
181                    buf.put_slice(v);
182                }
183            }
184        }
185        Ok(())
186    }
187
188    /// Decode a single record.
189    ///
190    /// `buf` must be positioned at the record's varlong length prefix.
191    ///
192    /// # Errors
193    ///
194    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
195    pub fn decode<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
196        let body_len = get_varlong(buf)
197            .map_err(|e| RecordsError::RecordParse(format!("record length: {e}")))?;
198        let body_len = usize::try_from(body_len).map_err(|_| {
199            RecordsError::RecordParse(format!("record length negative or too large: {body_len}"))
200        })?;
201        if buf.remaining() < body_len {
202            return Err(RecordsError::BodyTooShort {
203                needed: body_len - buf.remaining(),
204            });
205        }
206        // Restrict to body_len bytes so a malformed inner field doesn't run
207        // past the record boundary.
208        let mut body = buf.take(body_len);
209        let r = Self::decode_body(&mut body)?;
210        // Trailing bytes inside the record's claimed length — protocol corruption.
211        if body.has_remaining() {
212            return Err(RecordsError::RecordParse(format!(
213                "trailing bytes inside record (left={})",
214                body.remaining()
215            )));
216        }
217        Ok(r)
218    }
219
220    fn decode_body<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
221        if buf.remaining() == 0 {
222            return Err(RecordsError::RecordParse("record body empty".into()));
223        }
224        let attributes = buf.get_i8();
225        let timestamp_delta = get_varlong(buf)
226            .map_err(|e| RecordsError::RecordParse(format!("timestamp_delta: {e}")))?;
227        let offset_delta =
228            get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("offset_delta: {e}")))?;
229
230        let key = decode_nullable_bytes(buf, "key")?;
231        let value = decode_nullable_bytes(buf, "value")?;
232
233        let header_count =
234            get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("header_count: {e}")))?;
235        let Ok(header_count_usize) = usize::try_from(header_count) else {
236            return Err(RecordsError::RecordParse(format!(
237                "negative header count {header_count}"
238            )));
239        };
240        // Bound pre-allocation: a record header is at least 1 byte on the wire,
241        // so an honest `header_count` can never exceed the bytes left in the
242        // record body. Clamp the capacity hint to reject huge declared counts
243        // without affecting the loop (it stops at EOF anyway).
244        let mut headers = Vec::with_capacity(header_count_usize.min(buf.remaining()));
245        for i in 0..header_count {
246            headers.push(
247                decode_record_header(buf)
248                    .map_err(|e| RecordsError::RecordParse(format!("header[{i}]: {e}")))?,
249            );
250        }
251
252        Ok(Self {
253            attributes,
254            timestamp_delta,
255            offset_delta,
256            key,
257            value,
258            headers,
259        })
260    }
261}
262
263/// The body size of a record with these field lengths, shared by the owned and
264/// the borrowed `Record::body_len`. `key` and `value` are `None` for null, and
265/// `headers` yields each header's key length and optional value length.
266pub fn body_len_of(
267    timestamp_delta: i64,
268    offset_delta: i32,
269    key: Option<usize>,
270    value: Option<usize>,
271    headers: impl ExactSizeIterator<Item = (usize, Option<usize>)>,
272) -> usize {
273    let field = |len: Option<usize>| {
274        len.map_or_else(
275            || varint_len(-1),
276            |len| varint_len(i32::try_from(len).unwrap_or(i32::MAX)) + len,
277        )
278    };
279    let header_count = varint_len(i32::try_from(headers.len()).unwrap_or(i32::MAX));
280    let headers: usize = headers
281        .map(|(key, value)| field(Some(key)) + field(value))
282        .sum();
283    1 // attributes (i8)
284        + varlong_len(timestamp_delta)
285        + varint_len(offset_delta)
286        + field(key)
287        + field(value)
288        + header_count
289        + headers
290}
291
292fn decode_nullable_bytes<B: Buf>(buf: &mut B, label: &str) -> Result<Option<Bytes>, RecordsError> {
293    let len =
294        get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("{label} length: {e}")))?;
295    // A negative length is null.
296    let Ok(n) = usize::try_from(len) else {
297        return Ok(None);
298    };
299    if buf.remaining() < n {
300        return Err(RecordsError::BodyTooShort {
301            needed: n - buf.remaining(),
302        });
303    }
304    let mut v = vec![0u8; n];
305    buf.copy_to_slice(&mut v);
306    Ok(Some(Bytes::from(v)))
307}
308
309/// The bytes written to the front of `prefix` when `unused` bytes of it are
310/// left over.
311fn written_prefix(prefix: &[u8], unused: usize) -> Result<&[u8], RecordsError> {
312    prefix
313        .len()
314        .checked_sub(unused)
315        .and_then(|used| prefix.get(..used))
316        .ok_or_else(|| RecordsError::RecordParse("record prefix overflow".into()))
317}
318
319fn decode_record_header<B: Buf>(buf: &mut B) -> Result<RecordHeader, String> {
320    let key_len = get_varint(buf).map_err(|e| format!("key length: {e}"))?;
321    let Ok(n) = usize::try_from(key_len) else {
322        return Err(format!("non-nullable key has negative length {key_len}"));
323    };
324    if buf.remaining() < n {
325        return Err(format!("key truncated (need {} more)", n - buf.remaining()));
326    }
327    let mut kv = vec![0u8; n];
328    buf.copy_to_slice(&mut kv);
329    let key = String::from_utf8(kv).map_err(|e| format!("key utf-8: {e}"))?;
330
331    let value_len = get_varint(buf).map_err(|e| format!("value length: {e}"))?;
332    // A negative length is null.
333    let value = if let Ok(n) = usize::try_from(value_len) {
334        if buf.remaining() < n {
335            return Err(format!(
336                "value truncated (need {} more)",
337                n - buf.remaining()
338            ));
339        }
340        let mut vv = vec![0u8; n];
341        buf.copy_to_slice(&mut vv);
342        Some(Bytes::from(vv))
343    } else {
344        None
345    };
346
347    Ok(RecordHeader { key, value })
348}
349
350#[cfg(test)]
351mod record_tests {
352
353    use std::iter;
354
355    use bytes::BytesMut;
356
357    use super::*;
358
359    fn fixture_minimal_record() -> Record {
360        Record {
361            attributes: 0,
362            timestamp_delta: 0,
363            offset_delta: 0,
364            key: None,
365            value: None,
366            headers: vec![],
367        }
368    }
369
370    fn fixture_keyed_record() -> Record {
371        Record {
372            attributes: 0,
373            timestamp_delta: 17,
374            offset_delta: 2,
375            key: Some(Bytes::from_static(b"the-key")),
376            value: Some(Bytes::from_static(b"hello kafka")),
377            headers: vec![
378                RecordHeader {
379                    key: "trace-id".to_owned(),
380                    value: Some(Bytes::from_static(b"abc")),
381                },
382                RecordHeader {
383                    key: "null-val".to_owned(),
384                    value: None,
385                },
386            ],
387        }
388    }
389
390    fn fixture_large_payload_record() -> Record {
391        Record {
392            attributes: 0,
393            timestamp_delta: 1_000_000,
394            offset_delta: 999,
395            key: Some(Bytes::from(vec![b'k'; 128])),
396            value: Some(Bytes::from(vec![b'v'; 4096])),
397            headers: vec![],
398        }
399    }
400
401    #[test]
402    fn record_prefix_golden_bytes_across_destination_splits() {
403        let mut cases = vec![
404            (Record::default(), vec![0x0c, 0, 0, 0, 1, 1, 0]),
405            (
406                Record {
407                    attributes: i8::MIN,
408                    timestamp_delta: i64::MIN,
409                    offset_delta: i32::MIN,
410                    key: Some(Bytes::from_static(b"k")),
411                    value: Some(Bytes::from_static(b"v")),
412                    headers: vec![],
413                },
414                vec![
415                    0x2a, 0x80, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01, 0xff,
416                    0xff, 0xff, 0xff, 0x0f, 0x02, b'k', 0x02, b'v', 0,
417                ],
418            ),
419            (
420                Record {
421                    offset_delta: 129,
422                    key: Some(Bytes::from_static(b"k")),
423                    value: Some(Bytes::from_static(b"v")),
424                    ..Default::default()
425                },
426                vec![0x12, 0, 0, 0x82, 0x02, 0x02, b'k', 0x02, b'v', 0],
427            ),
428        ];
429        // Null and empty keys differ; 32/33 cover the short-key boundary,
430        // and 63/64 cross a zigzag length-prefix boundary.
431        for (len, length_prefix, key_prefix) in [
432            (0, vec![12], vec![0]),
433            (32, vec![76], vec![64]),
434            (33, vec![78], vec![66]),
435            (63, vec![138, 1], vec![126]),
436            (64, vec![142, 1], vec![128, 1]),
437        ] {
438            let mut expected = length_prefix;
439            expected.extend_from_slice(&[0, 0, 0]);
440            expected.extend(key_prefix);
441            expected.extend(iter::repeat_n(b'k', len));
442            expected.extend_from_slice(&[1, 0]);
443            cases.push((
444                Record {
445                    key: Some(Bytes::from(vec![b'k'; len])),
446                    ..Default::default()
447                },
448                expected,
449            ));
450        }
451        for (record, expected) in cases {
452            for split in 0..=expected.len() {
453                let mut encoded = vec![0xcd; expected.len()];
454                let (first, second) = encoded.split_at_mut(split);
455                let mut destination = first.chain_mut(second);
456                record.encode(&mut destination).unwrap();
457                assert2::assert!(destination.remaining_mut() == 0);
458                assert2::assert!(encoded == expected);
459                let mut cursor = encoded.as_slice();
460                assert2::assert!(Record::decode(&mut cursor).unwrap() == record);
461                assert2::assert!(cursor.is_empty());
462            }
463        }
464    }
465
466    #[test]
467    fn record_roundtrip_cases() {
468        type TestCase1<'a> = (&'a str, fn() -> Record);
469        let cases: [TestCase1<'_>; 3] = [
470            ("minimal", fixture_minimal_record),
471            ("keyed with headers", fixture_keyed_record),
472            ("large payload", fixture_large_payload_record),
473        ];
474        for (_case, fixture) in cases {
475            let record = fixture();
476            let mut buf = BytesMut::new();
477            record.encode(&mut buf).unwrap();
478            assert2::assert!(buf.len() == record.encoded_len());
479
480            let mut cur: &[u8] = &buf[..];
481            let decoded = Record::decode(&mut cur).unwrap();
482            assert2::assert!((decoded, cur.is_empty()) == (record, true));
483        }
484    }
485
486    /// A record with zero deltas, a null key, no headers and a `value_len`-byte
487    /// value has a body of `1 + 1 + 1 + 1 + prefix + value_len + 1` bytes: the
488    /// attributes byte, two one-byte deltas, the one-byte `-1` key, the value
489    /// behind its own zigzag-varint length prefix, and a one-byte header count.
490    /// The record's length prefix is a zigzag varlong too, so it grows with the
491    /// body: one byte up to 63, two bytes from 64 to 8191, three from 8192. The
492    /// value's length prefix grows the same way, one byte up to 63, two from 64
493    /// to 8191 and three from 8192.
494    #[test]
495    fn body_len_is_the_size_the_length_prefix_declares() {
496        for (value_len, body, prefix) in [
497            (0, 6, 1),
498            (57, 63, 1),
499            (58, 64, 2),
500            (63, 69, 2),
501            (64, 71, 2),
502            (8_184, 8_191, 2),
503            (8_185, 8_192, 3),
504            (8_192, 8_200, 3),
505            (70_000, 70_008, 3),
506        ] {
507            let record = Record {
508                value: Some(Bytes::from(vec![7_u8; value_len])),
509                ..Default::default()
510            };
511            let mut buf = BytesMut::new();
512            record.encode(&mut buf).unwrap();
513
514            let mut cur: &[u8] = &buf[..];
515            let declared = get_varlong(&mut cur).unwrap();
516            let case = format!("value_len={value_len}");
517            assert2::check!(record.body_len() == body, "{case}");
518            assert2::check!(record.encoded_len() == prefix + body, "{case}");
519            assert2::check!(buf.len() == prefix + body, "{case}");
520            assert2::check!(declared == i64::try_from(body).unwrap(), "{case}");
521            assert2::check!(buf.len() - cur.len() == prefix, "{case}");
522        }
523    }
524
525    /// Every field counts toward the body, and a null key, value or header
526    /// value counts as its one-byte `-1` marker. The sizes are worked out by
527    /// hand from the fields, and match what `encode` declares.
528    #[test]
529    fn body_len_counts_every_field() {
530        type Fixture<'a> = (&'a str, fn() -> Record, usize);
531        let cases: [Fixture<'_>; 3] = [
532            // attributes, three one-byte deltas and markers, header count
533            ("minimal", fixture_minimal_record, 6),
534            // 1 attributes + 1 timestamp delta + 1 offset delta
535            // + (1 + 7) key + (1 + 11) value + 1 header count
536            // + (1 + 8) + (1 + 3) first header + (1 + 8) + 1 second header
537            ("keyed with headers", fixture_keyed_record, 47),
538            // 1 attributes + 3 timestamp delta + 2 offset delta
539            // + (2 + 128) key + (2 + 4096) value + 1 header count
540            ("large payload", fixture_large_payload_record, 4_235),
541        ];
542        for (case, fixture, body) in cases {
543            let record = fixture();
544            let mut buf = BytesMut::new();
545            record.encode(&mut buf).unwrap();
546            let declared = get_varlong(&mut &buf[..]).unwrap();
547            assert2::check!(record.body_len() == body, "{case}");
548            assert2::check!(declared == i64::try_from(body).unwrap(), "{case}");
549        }
550    }
551
552    #[test]
553    fn decode_rejects_negative_header_count() {
554        let mut buf = BytesMut::new();
555        // body: attributes(1) + timestamp_delta(1) + offset_delta(1) + key=-1(1)
556        //       + value=-1(1) + headers=-1(1) = 6 bytes body
557        put_varlong(&mut buf, 6); // body length 6 bytes
558        buf.put_i8(0); // attributes
559        put_varlong(&mut buf, 0); // timestamp_delta = 0  (1 byte)
560        put_varint(&mut buf, 0); // offset_delta = 0     (1 byte)
561        put_varint(&mut buf, -1); // key len               (1 byte)
562        put_varint(&mut buf, -1); // value len             (1 byte)
563        put_varint(&mut buf, -1); // negative header count (1 byte)
564
565        let mut cur: &[u8] = &buf[..];
566        match Record::decode(&mut cur) {
567            Err(RecordsError::RecordParse(msg)) => {
568                assert2::assert!(msg.contains("negative header count"));
569            }
570            other => panic!("expected RecordParse, got {other:?}"),
571        }
572    }
573
574    #[test]
575    fn decode_huge_header_count_does_not_overallocate() {
576        // A record that declares ~1 billion headers but supplies none. The
577        // capacity hint must be clamped to the (tiny) body remaining, and the
578        // decode must fail cleanly on EOF rather than attempting a multi-GB
579        // allocation.
580        let mut inner = BytesMut::new();
581        inner.put_i8(0); // attributes
582        put_varlong(&mut inner, 0); // timestamp_delta
583        put_varint(&mut inner, 0); // offset_delta
584        put_varint(&mut inner, -1); // key len = null
585        put_varint(&mut inner, -1); // value len = null
586        put_varint(&mut inner, 1_000_000_000); // absurd header count
587
588        let mut buf = BytesMut::new();
589        put_varlong(&mut buf, i64::try_from(inner.len()).unwrap());
590        buf.extend_from_slice(&inner);
591
592        let mut cur: &[u8] = &buf[..];
593        // Must return an Err (EOF reading the first header), not OOM.
594        assert2::assert!(Record::decode(&mut cur).is_err());
595    }
596}
597
598impl RecordBatch {
599    /// KIP-534 delete horizon, if the delete-horizon attribute bit is set.
600    ///
601    /// `base_timestamp` carries the horizon. There is no separate wire field.
602    #[must_use]
603    pub fn delete_horizon_ms(&self) -> Option<i64> {
604        self.attributes
605            .has_delete_horizon()
606            .then_some(self.base_timestamp)
607    }
608
609    /// Stamp the delete horizon onto this batch.
610    ///
611    /// This method sets bit 6 and moves the horizon into `base_timestamp`. It
612    /// then rewrites every record's `timestamp_delta`, so the reconstructed
613    /// absolute timestamps `base_timestamp + delta` do not change.
614    #[must_use]
615    pub fn with_delete_horizon(mut self, horizon_ms: i64) -> Self {
616        let old_base = self.base_timestamp;
617        for r in &mut self.records {
618            // Reconstruct the original absolute timestamp, then re-base it onto
619            // the new `base_timestamp`. Deltas are `i64`, so absolute timestamps
620            // round-trip exactly across re-basing; `saturating_*` only guards
621            // pathological inputs from panicking.
622            let abs = old_base.saturating_add(r.timestamp_delta);
623            r.timestamp_delta = abs.saturating_sub(horizon_ms);
624        }
625        self.base_timestamp = horizon_ms;
626        self.attributes = self.attributes.with_delete_horizon(true);
627        self
628    }
629
630    /// Decode a complete v2 record batch from `buf`.
631    ///
632    /// This method reads from the start of the header. Unlike
633    /// [`Decode::decode`](crate::Decode::decode), it takes no API version and
634    /// returns the typed [`RecordsError`].
635    ///
636    /// # Errors
637    ///
638    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
639    pub fn decode_batch<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
640        Self::decode_with_policy(buf, RecordDecompressionPolicy::default())
641    }
642
643    /// Decode a complete v2 record batch with explicit decompression limits.
644    ///
645    /// # Errors
646    ///
647    /// Returns the underlying records error for malformed, truncated, corrupt,
648    /// or over-limit input.
649    pub fn decode_with_policy<B: Buf>(
650        buf: &mut B,
651        policy: RecordDecompressionPolicy,
652    ) -> Result<Self, RecordsError> {
653        // Need the full header before doing anything.
654        if buf.remaining() < HEADER_LEN {
655            return Err(RecordsError::HeaderTooShort {
656                needed: HEADER_LEN - buf.remaining(),
657            });
658        }
659        // When the whole batch sits in the current chunk, its covered header
660        // and body are one contiguous region: CRC it in a single pass.
661        let contiguous_crc = contiguous_covered_crc(buf.chunk());
662        // Copy out the header to a stack buffer so we can use zerocopy.
663        let mut hdr_bytes = [0u8; HEADER_LEN];
664        buf.copy_to_slice(&mut hdr_bytes);
665
666        let hdr = RecordBatchHeader::ref_from_bytes(&hdr_bytes[..])
667            .map_err(|_| RecordsError::ZerocopyFailure)?;
668
669        if hdr.magic != 2 {
670            return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
671        }
672
673        let body_len = declared_body_len(hdr.batch_length.get()).ok_or_else(|| {
674            RecordsError::RecordParse("negative or oversized batch_length".into())
675        })?;
676
677        if buf.remaining() < body_len {
678            return Err(RecordsError::BodyTooShort {
679                needed: body_len - buf.remaining(),
680            });
681        }
682
683        // CRC is computed over: header bytes 21..HEADER_LEN (attributes through
684        // records_count), then the body bytes. Fragmented input continues the
685        // header CRC over the separately staged body.
686        let expected_crc = hdr.crc.get();
687        let verify_crc = |body: &[u8]| {
688            let computed = contiguous_crc.unwrap_or_else(|| {
689                crc32c_append(crc32c(&hdr_bytes[CRC_COVERAGE_START..HEADER_LEN]), body)
690            });
691            if computed == expected_crc {
692                Ok(())
693            } else {
694                Err(RecordsError::CrcMismatch {
695                    expected: expected_crc,
696                    computed,
697                })
698            }
699        };
700
701        let attributes = Attributes(hdr.attributes.get());
702        let codec = attributes.compression();
703
704        let decompress = |body: &[u8]| -> Result<Bytes, RecordsError> {
705            verify_crc(body)?;
706            Ok(krabka_compression::decompress(
707                codec,
708                body,
709                policy.output_limit(ByteSize::from_bytes(body.len() as u64)),
710            )?)
711        };
712
713        // A compressed body can be read directly from contiguous input. The
714        // decoded records own the decompressed output, so they retain no input
715        // buffer. Consume the body even on CRC/decompression errors, as the
716        // original staged read did. Fragmented input still needs staging;
717        // uncompressed records keep their own body instead of a larger frame.
718        let body_for_records = if codec != krabka_compression::CompressionType::None
719            && let Some(chunk) = buf.chunk().get(..body_len)
720        {
721            let decoded = decompress(chunk);
722            buf.advance(body_len);
723            decoded?
724        } else {
725            let mut body = vec![0u8; body_len];
726            buf.copy_to_slice(&mut body);
727            if codec == krabka_compression::CompressionType::None {
728                verify_crc(&body)?;
729                Bytes::from(body)
730            } else {
731                decompress(&body)?
732            }
733        };
734
735        // Parse records.
736        let count = hdr.records_count.get();
737        let Ok(count_usize) = usize::try_from(count) else {
738            return Err(RecordsError::RecordParse(format!(
739                "negative records_count {count}"
740            )));
741        };
742        let mut body_cur: &[u8] = &body_for_records[..];
743        // Bound pre-allocation: each record is at least 1 byte in the
744        // decompressed body, so an honest `records_count` can never exceed the
745        // body length. Clamp the capacity hint to reject huge declared counts.
746        let record_capacity = count_usize.min(body_for_records.len());
747        let mut records = Vec::with_capacity(record_capacity);
748        for i in 0..count {
749            // Parse the record into borrowed slices, then materialise its
750            // key / value / header-values as zero-copy `Bytes` views into
751            // `body_for_records`. The whole batch shares that single backing
752            // allocation instead of copying every field into its own `Bytes`,
753            // which removes ~2 heap allocations + memcpys per record on the
754            // common (header-free) path.
755            let r = parse_one_record(&mut body_cur)
756                .map_err(|e| RecordsError::RecordParse(format!("record[{i}]: {e}")))?;
757            records.push(Record {
758                attributes: r.attributes,
759                timestamp_delta: r.timestamp_delta,
760                offset_delta: r.offset_delta,
761                key: r.key.map(|s| body_for_records.slice_ref(s)),
762                value: r.value.map(|s| body_for_records.slice_ref(s)),
763                headers: r
764                    .headers
765                    .into_iter()
766                    .map(|h| RecordHeader {
767                        key: h.key.to_owned(),
768                        value: h.value.map(|s| body_for_records.slice_ref(s)),
769                    })
770                    .collect(),
771            });
772        }
773        if !body_cur.is_empty() {
774            return Err(RecordsError::RecordParse(format!(
775                "trailing bytes after records (left={})",
776                body_cur.len()
777            )));
778        }
779
780        Ok(Self {
781            base_offset: hdr.base_offset.get(),
782            partition_leader_epoch: hdr.partition_leader_epoch.get(),
783            attributes,
784            last_offset_delta: hdr.last_offset_delta.get(),
785            base_timestamp: hdr.base_timestamp.get(),
786            max_timestamp: hdr.max_timestamp.get(),
787            producer_id: hdr.producer_id.get(),
788            producer_epoch: hdr.producer_epoch.get(),
789            base_sequence: hdr.base_sequence.get(),
790            records,
791        })
792    }
793
794    /// Encode this batch into `buf`.
795    ///
796    /// Unlike [`Encode::encode`](crate::Encode::encode), this method takes no
797    /// API version and returns the typed [`RecordsError`].
798    ///
799    /// # Errors
800    ///
801    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
802    pub fn encode_batch<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
803        self.encode_with_compression_level(buf, None)
804    }
805
806    /// Encode this batch into an owned wire buffer.
807    ///
808    /// Uncompressed records are written directly behind the header, avoiding
809    /// a temporary body buffer and its copy into the output.
810    ///
811    /// # Errors
812    ///
813    /// Returns the same encoding errors as [`Self::encode_batch`].
814    pub fn encode_to_bytes(&self) -> Result<Bytes, RecordsError> {
815        if self.attributes.compression() == krabka_compression::CompressionType::None {
816            return Ok(self.encode_uncompressed()?.freeze());
817        }
818        let mut wire = BytesMut::with_capacity(self.encoded_batch_len());
819        self.encode_batch(&mut wire)?;
820        Ok(wire.freeze())
821    }
822
823    /// Encodes an uncompressed batch: the records go directly behind a
824    /// reserved header, then the header is filled in and the CRC taken over
825    /// the contiguous covered region in one pass.
826    fn encode_uncompressed(&self) -> Result<BytesMut, RecordsError> {
827        let mut wire = BytesMut::with_capacity(self.encoded_batch_len());
828        wire.resize(HEADER_LEN, 0);
829        for record in &self.records {
830            record.encode(&mut wire)?;
831        }
832        let body_len = wire.len() - HEADER_LEN;
833        let batch_length = batch_length_for(body_len)?;
834        let covered_head = self.covered_head()?;
835        let short = || RecordsError::RecordParse("batch buffer shorter than its header".into());
836        let (prefix, covered) = wire
837            .split_at_mut_checked(CRC_COVERAGE_START)
838            .ok_or_else(short)?;
839        let (covered_head_dst, _) = covered
840            .split_first_chunk_mut::<{ HEADER_LEN - CRC_COVERAGE_START }>()
841            .ok_or_else(short)?;
842        *covered_head_dst = covered_head;
843        let crc = crc32c(covered);
844        self.put_header_prefix(&mut &mut *prefix, batch_length, crc);
845        Ok(wire)
846    }
847
848    /// Encode this batch into `buf`, and compress the records at `level`.
849    ///
850    /// `None` uses the default level of the codec, as [`Self::encode_batch`] does.
851    /// A producer passes its `compression.<codec>.level` setting (KIP-390).
852    /// See [`krabka_compression::compress_with_level`] for how each codec
853    /// applies the level.
854    ///
855    /// # Errors
856    ///
857    /// Returns [`RecordsError::Compression`] when `level` is not a level that
858    /// Kafka accepts for the codec of the batch attributes, including any
859    /// level for a batch with no compression. Returns the other errors of
860    /// [`Self::encode_batch`].
861    pub fn encode_with_compression_level<B: BufMut>(
862        &self,
863        buf: &mut B,
864        level: Option<i32>,
865    ) -> Result<(), RecordsError> {
866        let codec = self.attributes.compression();
867        if level.is_none() && codec == krabka_compression::CompressionType::None {
868            buf.put_slice(&self.encode_uncompressed()?);
869            return Ok(());
870        }
871
872        // 1. Encode records into a temporary buffer.
873        let mut raw_body =
874            BytesMut::with_capacity(self.records.iter().map(Record::encoded_len).sum());
875        for r in &self.records {
876            r.encode(&mut raw_body)?;
877        }
878        let raw_body = raw_body.freeze();
879
880        // 2. Compress.
881        let body: Bytes = match level {
882            Some(level) => krabka_compression::compress_with_level(codec, &raw_body, level)?,
883            None => krabka_compression::compress(codec, &raw_body)?,
884        };
885
886        self.encode_header(buf, &body)?;
887        buf.put_slice(&body);
888        Ok(())
889    }
890
891    /// Writes the header for a `body` held apart from it, such as a
892    /// compressor's output: the CRC continues from the covered header into
893    /// the body.
894    fn encode_header<B: BufMut>(&self, buf: &mut B, body: &[u8]) -> Result<(), RecordsError> {
895        let batch_length = batch_length_for(body.len())?;
896        let covered_head = self.covered_head()?;
897        let crc = crc32c_append(crc32c(&covered_head), body);
898        self.put_header_prefix(buf, batch_length, crc);
899        buf.put_slice(&covered_head);
900        Ok(())
901    }
902
903    /// The CRC-covered header fields, attributes through `records_count`.
904    fn covered_head(&self) -> Result<[u8; HEADER_LEN - CRC_COVERAGE_START], RecordsError> {
905        let mut covered_head = [0_u8; HEADER_LEN - CRC_COVERAGE_START];
906        let mut covered = covered_head.as_mut_slice();
907        covered.put_i16(self.attributes.0);
908        covered.put_i32(self.last_offset_delta);
909        covered.put_i64(self.base_timestamp);
910        covered.put_i64(self.max_timestamp);
911        covered.put_i64(self.producer_id);
912        covered.put_i16(self.producer_epoch);
913        covered.put_i32(self.base_sequence);
914        covered.put_i32(
915            i32::try_from(self.records.len())
916                .map_err(|_| RecordsError::RecordParse("records_count exceeds i32".into()))?,
917        );
918        Ok(covered_head)
919    }
920
921    /// Writes the header fields before the CRC-covered region: base offset,
922    /// batch length, leader epoch, magic and CRC.
923    fn put_header_prefix<B: BufMut>(&self, buf: &mut B, batch_length: i32, crc: u32) {
924        buf.put_i64(self.base_offset);
925        buf.put_i32(batch_length);
926        buf.put_i32(self.partition_leader_epoch);
927        buf.put_i8(2); // magic v2
928        buf.put_u32(crc);
929    }
930
931    /// Predicted total bytes that [`Self::encode_batch`] will write, for an
932    /// uncompressed batch. Unlike [`Encode::encoded_len`](crate::Encode::encoded_len),
933    /// it takes no API version.
934    ///
935    /// For a compressed batch, the actual size will differ.
936    pub fn encoded_batch_len(&self) -> usize {
937        let body: usize = self.records.iter().map(Record::encoded_len).sum();
938        HEADER_LEN + body
939    }
940}
941
942/// Header fields counted by `batch_length`, which counts the bytes after
943/// itself: `partition_leader_epoch`(4) + magic(1) + crc(4) + attributes(2) +
944/// `last_offset_delta`(4) + `base_timestamp`(8) + `max_timestamp`(8) +
945/// `producer_id`(8) + `producer_epoch`(2) + `base_sequence`(4) +
946/// `records_count`(4) = 49 bytes.
947const HEADER_TAIL_LEN: i32 = 49;
948
949/// The body length a `batch_length` declares, or `None` when it is shorter
950/// than the header tail.
951fn declared_body_len(batch_length: i32) -> Option<usize> {
952    let n = i32::checked_sub(batch_length, HEADER_TAIL_LEN)?;
953    usize::try_from(n).ok()
954}
955
956/// The `batch_length` of a batch with a `body_len`-byte body.
957fn batch_length_for(body_len: usize) -> Result<i32, RecordsError> {
958    i32::try_from(body_len)
959        .ok()
960        .and_then(|body_len| HEADER_TAIL_LEN.checked_add(body_len))
961        .ok_or_else(|| RecordsError::RecordParse("body length exceeds i32".into()))
962}
963
964/// The CRC-32C of a v2 batch's covered region, when `chunk` starts with the
965/// whole batch; `None` otherwise, or when `chunk` is no v2 batch.
966fn contiguous_covered_crc(chunk: &[u8]) -> Option<u32> {
967    let hdr = RecordBatchHeader::ref_from_bytes(chunk.get(..HEADER_LEN)?).ok()?;
968    if hdr.magic != 2 {
969        return None;
970    }
971    let body_len = declared_body_len(hdr.batch_length.get())?;
972    let covered = chunk.get(CRC_COVERAGE_START..HEADER_LEN.checked_add(body_len)?)?;
973    Some(crc32c(covered))
974}
975
976#[cfg(test)]
977mod batch_tests {
978    use assert2::check;
979    use krabka_compression::{CompressionError, CompressionType, RecordDecompressionPolicy};
980    use krabka_units::{bytes, fraction};
981
982    use super::*;
983
984    fn fixture_empty_batch() -> RecordBatch {
985        RecordBatch::default()
986    }
987
988    fn fixture_single_record_batch() -> RecordBatch {
989        RecordBatch {
990            records: vec![Record {
991                key: Some(Bytes::from_static(b"k1")),
992                value: Some(Bytes::from_static(b"v1")),
993                ..Default::default()
994            }],
995            ..RecordBatch::default()
996        }
997    }
998
999    fn fixture_multi_record_batch() -> RecordBatch {
1000        RecordBatch {
1001            base_offset: 42,
1002            partition_leader_epoch: 5,
1003            last_offset_delta: 2,
1004            base_timestamp: 1_700_000_000,
1005            max_timestamp: 1_700_000_500,
1006            producer_id: 100,
1007            producer_epoch: 3,
1008            base_sequence: 7,
1009            records: vec![
1010                Record {
1011                    offset_delta: 0,
1012                    timestamp_delta: 0,
1013                    key: Some(Bytes::from_static(b"a")),
1014                    value: Some(Bytes::from_static(b"1")),
1015                    ..Default::default()
1016                },
1017                Record {
1018                    offset_delta: 1,
1019                    timestamp_delta: 100,
1020                    key: Some(Bytes::from_static(b"b")),
1021                    value: Some(Bytes::from_static(b"2")),
1022                    ..Default::default()
1023                },
1024                Record {
1025                    offset_delta: 2,
1026                    timestamp_delta: 500,
1027                    key: None,
1028                    value: Some(Bytes::from_static(b"3")),
1029                    headers: vec![RecordHeader {
1030                        key: "h".to_owned(),
1031                        value: Some(Bytes::from_static(b"hv")),
1032                    }],
1033                    ..Default::default()
1034                },
1035            ],
1036            ..RecordBatch::default()
1037        }
1038    }
1039
1040    #[test]
1041    fn uncompressed_roundtrip_cases() {
1042        type TestCase2 = (&'static str, fn() -> RecordBatch);
1043        let cases: [TestCase2; 3] = [
1044            ("empty", fixture_empty_batch),
1045            ("single", fixture_single_record_batch),
1046            ("multiple", fixture_multi_record_batch),
1047        ];
1048        for (_case, fixture) in cases {
1049            let mut batch = fixture();
1050            batch.attributes = batch.attributes.with_compression(CompressionType::None);
1051            let mut buf = BytesMut::new();
1052            batch.encode_batch(&mut buf).unwrap();
1053            assert2::assert!(buf.len() == batch.encoded_batch_len());
1054            let mut cur: &[u8] = &buf[..];
1055            let decoded = RecordBatch::decode_batch(&mut cur).unwrap();
1056            assert2::assert!((decoded, cur.is_empty()) == (batch, true));
1057        }
1058    }
1059
1060    /// `encode` and `encode_to_bytes` write the same bytes, and the stored
1061    /// CRC is the reference CRC-32C of everything after the CRC field.
1062    #[test]
1063    fn encodings_agree_and_store_the_reference_crc() {
1064        for codec in [
1065            CompressionType::None,
1066            CompressionType::Gzip,
1067            CompressionType::Lz4,
1068        ] {
1069            for fixture in [
1070                fixture_empty_batch,
1071                fixture_single_record_batch,
1072                fixture_multi_record_batch,
1073            ] {
1074                let mut batch = fixture();
1075                batch.attributes = batch.attributes.with_compression(codec);
1076                let mut encoded = BytesMut::new();
1077                batch.encode_batch(&mut encoded).unwrap();
1078                let to_bytes = batch.encode_to_bytes().unwrap();
1079                let stored = u32::from_be_bytes(encoded[17..21].try_into().unwrap());
1080                assert2::assert!(
1081                    (
1082                        to_bytes[..] == encoded[..],
1083                        stored == crc32c::crc32c(&encoded[CRC_COVERAGE_START..]),
1084                    ) == (true, true),
1085                    "{codec:?}"
1086                );
1087            }
1088        }
1089    }
1090
1091    /// A batch split across `Buf` chunks decodes as the contiguous batch does,
1092    /// and a corrupt one reports the same CRC mismatch, wherever the split
1093    /// falls: in the uncovered header, in the covered header, at the body or
1094    /// inside it.
1095    #[test]
1096    fn fragmented_input_decodes_like_contiguous_input() {
1097        for codec in [CompressionType::None, CompressionType::Lz4] {
1098            let mut batch = fixture_multi_record_batch();
1099            batch.attributes = batch.attributes.with_compression(codec);
1100            let mut encoded = BytesMut::new();
1101            batch.encode_batch(&mut encoded).unwrap();
1102            let mut corrupt = encoded.to_vec();
1103            *corrupt.last_mut().unwrap() ^= 0xFF;
1104            let contiguous_error = RecordBatch::decode_batch(&mut &corrupt[..]).unwrap_err();
1105            let RecordsError::CrcMismatch { expected, computed } = contiguous_error else {
1106                panic!("{codec:?}: {contiguous_error:?}");
1107            };
1108            assert2::assert!(computed == crc32c::crc32c(&corrupt[CRC_COVERAGE_START..]));
1109
1110            for split in [0, 10, 30, HEADER_LEN, HEADER_LEN + 1, encoded.len() - 1] {
1111                let mut fragmented = (&encoded[..split]).chain(&encoded[split..]);
1112                let decoded = RecordBatch::decode_batch(&mut fragmented).unwrap();
1113                assert2::assert!(
1114                    (decoded == batch, fragmented.has_remaining()) == (true, false),
1115                    "{codec:?} split at {split}"
1116                );
1117
1118                let mut fragmented = (&corrupt[..split]).chain(&corrupt[split..]);
1119                let error = RecordBatch::decode_batch(&mut fragmented).unwrap_err();
1120                assert2::assert!(
1121                    matches!(
1122                        error,
1123                        RecordsError::CrcMismatch { expected: e, computed: c }
1124                            if (e, c) == (expected, computed)
1125                    ),
1126                    "{codec:?} split at {split}: {error:?}"
1127                );
1128            }
1129        }
1130    }
1131
1132    /// The level reaches the codec. A gzip batch at level 1 and at level 9
1133    /// has different bytes, both decode to the batch, and the default level
1134    /// gives the bytes of `encode`. A level that Kafka rejects fails the
1135    /// encode.
1136    #[test]
1137    fn encode_with_compression_level_compresses_at_the_level() {
1138        let value: Vec<u8> = (0..16 * 1024u32)
1139            .map(|i| u8::try_from((i * 7 + i / 13) % 251).unwrap_or(0))
1140            .collect();
1141        let mut batch = RecordBatch {
1142            records: vec![Record {
1143                value: Some(Bytes::from(value)),
1144                ..Default::default()
1145            }],
1146            ..RecordBatch::default()
1147        };
1148        batch.attributes = batch.attributes.with_compression(CompressionType::Gzip);
1149        let encoded = |level: Option<i32>| {
1150            let mut buf = BytesMut::new();
1151            batch
1152                .encode_with_compression_level(&mut buf, level)
1153                .map(|()| buf.freeze())
1154        };
1155        let default = encoded(None).unwrap();
1156        let fastest = encoded(Some(1)).unwrap();
1157        let best = encoded(Some(9)).unwrap();
1158        let decoded = [&fastest, &best]
1159            .map(|bytes| RecordBatch::decode_batch(&mut &bytes[..]).unwrap() == batch);
1160
1161        assert2::assert!(
1162            (
1163                encoded(Some(-1)).unwrap() == default,
1164                fastest != best,
1165                decoded,
1166            ) == (true, true, [true, true])
1167        );
1168        assert2::assert!(matches!(
1169            encoded(Some(10)),
1170            Err(RecordsError::Compression(CompressionError::InvalidLevel {
1171                level: 10,
1172                ..
1173            }))
1174        ));
1175    }
1176
1177    #[test]
1178    fn rejects_pre_v2_magic() {
1179        let mut buf = BytesMut::new();
1180        buf.put_i64(0); // base_offset
1181        buf.put_i32(49); // batch_length
1182        buf.put_i32(0); // partition_leader_epoch
1183        buf.put_i8(1); // magic = 1 (v1, deprecated)
1184        buf.put_u32(0); // crc (irrelevant; we reject on magic first)
1185        for _ in 21..HEADER_LEN {
1186            buf.put_u8(0);
1187        }
1188        let mut cur: &[u8] = &buf[..];
1189        assert2::assert!(matches!(
1190            RecordBatch::decode_batch(&mut cur),
1191            Err(RecordsError::UnsupportedMagic { found: 1 })
1192        ));
1193    }
1194
1195    #[test]
1196    fn rejects_bad_crc() {
1197        let b = fixture_single_record_batch();
1198        let mut buf = BytesMut::new();
1199        b.encode_batch(&mut buf).unwrap();
1200        // Corrupt the CRC bytes (offsets 17..21).
1201        buf[17] ^= 0xFF;
1202        let mut cur: &[u8] = &buf[..];
1203        assert2::assert!(matches!(
1204            RecordBatch::decode_batch(&mut cur),
1205            Err(RecordsError::CrcMismatch { .. })
1206        ));
1207    }
1208
1209    #[test]
1210    fn compressed_roundtrip_cases() {
1211        for (_case, codec) in [
1212            ("gzip", CompressionType::Gzip),
1213            ("snappy", CompressionType::Snappy),
1214            ("lz4", CompressionType::Lz4),
1215            ("zstd", CompressionType::Zstd),
1216        ] {
1217            let mut batch = fixture_multi_record_batch();
1218            batch.attributes = batch.attributes.with_compression(codec);
1219            let mut buf = BytesMut::new();
1220            batch.encode_batch(&mut buf).unwrap();
1221            let mut cur: &[u8] = &buf[..];
1222            let decoded = RecordBatch::decode_batch(&mut cur).unwrap();
1223            assert2::assert!((decoded, cur.is_empty()) == (batch, true));
1224        }
1225    }
1226
1227    #[test]
1228    fn decompression_policy_limits_owned_decode() {
1229        let mut batch = fixture_single_record_batch();
1230        batch.records[0].value = Some(Bytes::from(vec![b'x'; 4096]));
1231        batch.attributes = batch.attributes.with_compression(CompressionType::Lz4);
1232        let mut wire = BytesMut::new();
1233        batch.encode_batch(&mut wire).unwrap();
1234
1235        let mut default_cur = &wire[..];
1236        RecordBatch::decode_batch(&mut default_cur).unwrap();
1237
1238        let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
1239        let mut limited_cur = &wire[..];
1240        assert2::assert!(matches!(
1241            RecordBatch::decode_with_policy(&mut limited_cur, policy),
1242            Err(RecordsError::Compression(CompressionError::TooLarge {
1243                limit: 32
1244            }))
1245        ));
1246    }
1247
1248    #[test]
1249    fn with_delete_horizon_stamps_and_preserves_record_timestamps() {
1250        // base 1000, two records at deltas [0, 5] → absolutes [1000, 1005].
1251        let b = RecordBatch {
1252            base_timestamp: 1000,
1253            records: vec![
1254                Record {
1255                    timestamp_delta: 0,
1256                    ..Default::default()
1257                },
1258                Record {
1259                    timestamp_delta: 5,
1260                    ..Default::default()
1261                },
1262            ],
1263            ..RecordBatch::default()
1264        };
1265
1266        let stamped = b.with_delete_horizon(9999);
1267
1268        check!(stamped.attributes.has_delete_horizon());
1269        check!(stamped.base_timestamp == 9999);
1270        check!(stamped.delete_horizon_ms() == Some(9999));
1271
1272        // Reconstructed absolutes (base + delta) must equal the ORIGINALS.
1273        let absolutes: Vec<i64> = stamped
1274            .records
1275            .iter()
1276            .map(|r| stamped.base_timestamp + r.timestamp_delta)
1277            .collect();
1278        assert2::assert!(absolutes == vec![1000, 1005]);
1279    }
1280
1281    #[test]
1282    fn delete_horizon_round_trips_through_encode_decode() {
1283        // base 1000, two keyed records at deltas [0, 5] → absolutes [1000, 1005].
1284        let b = RecordBatch {
1285            base_timestamp: 1000,
1286            last_offset_delta: 1,
1287            records: vec![
1288                Record {
1289                    timestamp_delta: 0,
1290                    offset_delta: 0,
1291                    key: Some(Bytes::from_static(b"k1")),
1292                    value: Some(Bytes::from_static(b"v1")),
1293                    ..Default::default()
1294                },
1295                Record {
1296                    timestamp_delta: 5,
1297                    offset_delta: 1,
1298                    key: Some(Bytes::from_static(b"k2")),
1299                    value: Some(Bytes::from_static(b"v2")),
1300                    ..Default::default()
1301                },
1302            ],
1303            ..RecordBatch::default()
1304        }
1305        .with_delete_horizon(9999);
1306
1307        let mut buf = BytesMut::new();
1308        b.encode_batch(&mut buf).unwrap();
1309
1310        let mut cur: &[u8] = &buf[..];
1311        let decoded = RecordBatch::decode_batch(&mut cur).unwrap();
1312        assert2::assert!(cur.is_empty());
1313
1314        assert2::assert!(decoded.delete_horizon_ms() == Some(9999));
1315
1316        let absolutes: Vec<i64> = decoded
1317            .records
1318            .iter()
1319            .map(|r| decoded.base_timestamp + r.timestamp_delta)
1320            .collect();
1321        assert2::assert!(absolutes == vec![1000, 1005]);
1322    }
1323
1324    #[test]
1325    fn decode_huge_records_count_does_not_overallocate() {
1326        // Encode an empty uncompressed batch, then overwrite records_count with
1327        // an absurd value (~1 billion) and fix up the CRC. The capacity hint
1328        // must be clamped to the (empty) decompressed body, and decode must
1329        // fail cleanly on EOF rather than attempting a multi-GB allocation.
1330        let mut b = fixture_empty_batch();
1331        b.attributes = b.attributes.with_compression(CompressionType::None);
1332        let mut buf = BytesMut::new();
1333        b.encode_batch(&mut buf).unwrap();
1334
1335        // records_count is the last 4 bytes of the fixed header.
1336        let rc_off = HEADER_LEN - 4;
1337        buf[rc_off..HEADER_LEN].copy_from_slice(&1_000_000_000i32.to_be_bytes());
1338
1339        // Recompute CRC over header[21..HEADER_LEN] + body (body is empty here).
1340        let computed = crc32c(&buf[CRC_COVERAGE_START..]);
1341        buf[17..21].copy_from_slice(&computed.to_be_bytes());
1342
1343        let mut cur: &[u8] = &buf[..];
1344        // Must return an Err (EOF reading the first record), not OOM.
1345        assert2::assert!(RecordBatch::decode_batch(&mut cur).is_err());
1346    }
1347}
1348
1349impl crate::Encode for RecordBatch {
1350    fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
1351        self.encode_batch(buf).map_err(Into::into)
1352    }
1353
1354    fn encoded_len(&self, _version: i16) -> usize {
1355        self.encoded_batch_len()
1356    }
1357}
1358
1359impl crate::Decode<'_> for RecordBatch {
1360    fn decode<B: Buf>(buf: &mut B, _version: i16) -> Result<Self, crate::ProtocolError> {
1361        Self::decode_batch(buf).map_err(Into::into)
1362    }
1363}