Skip to main content

crabka_protocol/records/
owned.rs

1//! Owned `RecordBatch`, `Record`, and `RecordHeader` types.
2
3use bytes::{Buf, BufMut, Bytes, BytesMut};
4use zerocopy::FromBytes as _;
5
6use crate::{
7    primitives::varint::{
8        get_varint, get_varlong, put_varint, put_varlong, varint_len, varlong_len,
9    },
10    records::{
11        RecordsError,
12        crc::{crc32c, crc32c_append},
13        header::{Attributes, HEADER_LEN},
14    },
15};
16
17#[derive(Debug, Clone, PartialEq, Eq, Default)]
18pub struct RecordHeader {
19    pub key: String,
20    pub value: Option<Bytes>,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, Default)]
24pub struct Record {
25    pub attributes: i8,
26    pub timestamp_delta: i64,
27    pub offset_delta: i32,
28    pub key: Option<Bytes>,
29    pub value: Option<Bytes>,
30    pub headers: Vec<RecordHeader>,
31}
32
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct RecordBatch {
35    pub base_offset: i64,
36    pub partition_leader_epoch: i32,
37    pub attributes: Attributes,
38    pub last_offset_delta: i32,
39    pub base_timestamp: i64,
40    pub max_timestamp: i64,
41    pub producer_id: i64,
42    pub producer_epoch: i16,
43    pub base_sequence: i32,
44    pub records: Vec<Record>,
45}
46
47impl Default for RecordBatch {
48    fn default() -> Self {
49        Self {
50            base_offset: 0,
51            partition_leader_epoch: 0,
52            attributes: Attributes::default(),
53            last_offset_delta: 0,
54            base_timestamp: 0,
55            max_timestamp: 0,
56            producer_id: -1, // sentinel: non-idempotent
57            producer_epoch: -1,
58            base_sequence: -1,
59            records: Vec::new(),
60        }
61    }
62}
63
64impl Record {
65    /// Encode a single record (varlong length prefix + fields) into `buf`.
66    pub fn encode<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
67        let body_len = self.body_len();
68        put_varlong(
69            buf,
70            i64::try_from(body_len)
71                .map_err(|_| RecordsError::RecordParse("record body length overflow".into()))?,
72        );
73        self.encode_body(buf)
74    }
75
76    /// Predicted total length of this record on the wire (length-prefix + body).
77    pub fn encoded_len(&self) -> usize {
78        let body = self.body_len();
79        #[allow(clippy::cast_possible_wrap, clippy::cast_possible_truncation)]
80        let body_i64 = body as i64;
81        varlong_len(body_i64) + body
82    }
83
84    fn body_len(&self) -> usize {
85        let mut n = 1; // attributes (i8)
86        n += varlong_len(self.timestamp_delta);
87        n += varint_len(self.offset_delta);
88        n += match &self.key {
89            None => varint_len(-1),
90            Some(k) => varint_len(i32::try_from(k.len()).unwrap_or(i32::MAX)) + k.len(),
91        };
92        n += match &self.value {
93            None => varint_len(-1),
94            Some(v) => varint_len(i32::try_from(v.len()).unwrap_or(i32::MAX)) + v.len(),
95        };
96        n += varint_len(i32::try_from(self.headers.len()).unwrap_or(i32::MAX));
97        for h in &self.headers {
98            let key_bytes = h.key.as_bytes();
99            n += varint_len(i32::try_from(key_bytes.len()).unwrap_or(i32::MAX)) + key_bytes.len();
100            n += match &h.value {
101                None => varint_len(-1),
102                Some(v) => varint_len(i32::try_from(v.len()).unwrap_or(i32::MAX)) + v.len(),
103            };
104        }
105        n
106    }
107
108    fn encode_body<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
109        buf.put_i8(self.attributes);
110        put_varlong(buf, self.timestamp_delta);
111        put_varint(buf, self.offset_delta);
112        match &self.key {
113            None => put_varint(buf, -1),
114            Some(k) => {
115                put_varint(
116                    buf,
117                    i32::try_from(k.len()).map_err(|_| {
118                        RecordsError::RecordParse("record key length overflow".into())
119                    })?,
120                );
121                buf.put_slice(k);
122            }
123        }
124        match &self.value {
125            None => put_varint(buf, -1),
126            Some(v) => {
127                put_varint(
128                    buf,
129                    i32::try_from(v.len()).map_err(|_| {
130                        RecordsError::RecordParse("record value length overflow".into())
131                    })?,
132                );
133                buf.put_slice(v);
134            }
135        }
136        put_varint(
137            buf,
138            i32::try_from(self.headers.len())
139                .map_err(|_| RecordsError::RecordParse("record header count overflow".into()))?,
140        );
141        for h in &self.headers {
142            let key_bytes = h.key.as_bytes();
143            put_varint(
144                buf,
145                i32::try_from(key_bytes.len())
146                    .map_err(|_| RecordsError::RecordParse("header key length overflow".into()))?,
147            );
148            buf.put_slice(key_bytes);
149            match &h.value {
150                None => put_varint(buf, -1),
151                Some(v) => {
152                    put_varint(
153                        buf,
154                        i32::try_from(v.len()).map_err(|_| {
155                            RecordsError::RecordParse("header value length overflow".into())
156                        })?,
157                    );
158                    buf.put_slice(v);
159                }
160            }
161        }
162        Ok(())
163    }
164
165    /// Decode a single record. `buf` must be positioned at the record's
166    /// varlong length prefix.
167    pub fn decode<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
168        let body_len = get_varlong(buf)
169            .map_err(|e| RecordsError::RecordParse(format!("record length: {e}")))?;
170        let body_len = usize::try_from(body_len).map_err(|_| {
171            RecordsError::RecordParse(format!("record length negative or too large: {body_len}"))
172        })?;
173        if buf.remaining() < body_len {
174            return Err(RecordsError::BodyTooShort {
175                needed: body_len - buf.remaining(),
176            });
177        }
178        // Restrict to body_len bytes so a malformed inner field doesn't run
179        // past the record boundary.
180        let mut body = buf.take(body_len);
181        let r = Self::decode_body(&mut body)?;
182        // Trailing bytes inside the record's claimed length — protocol corruption.
183        if body.has_remaining() {
184            return Err(RecordsError::RecordParse(format!(
185                "trailing bytes inside record (left={})",
186                body.remaining()
187            )));
188        }
189        Ok(r)
190    }
191
192    fn decode_body<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
193        if buf.remaining() == 0 {
194            return Err(RecordsError::RecordParse("record body empty".into()));
195        }
196        let attributes = buf.get_i8();
197        let timestamp_delta = get_varlong(buf)
198            .map_err(|e| RecordsError::RecordParse(format!("timestamp_delta: {e}")))?;
199        let offset_delta =
200            get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("offset_delta: {e}")))?;
201
202        let key = decode_nullable_bytes(buf, "key")?;
203        let value = decode_nullable_bytes(buf, "value")?;
204
205        let header_count =
206            get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("header_count: {e}")))?;
207        if header_count < 0 {
208            return Err(RecordsError::RecordParse(format!(
209                "negative header count {header_count}"
210            )));
211        }
212        #[allow(clippy::cast_sign_loss)] // checked < 0 above
213        let header_count_usize = header_count as usize;
214        // Bound pre-allocation: a record header is at least 1 byte on the wire,
215        // so an honest `header_count` can never exceed the bytes left in the
216        // record body. Clamp the capacity hint to reject huge declared counts
217        // without affecting the loop (it stops at EOF anyway).
218        let mut headers = Vec::with_capacity(header_count_usize.min(buf.remaining()));
219        for i in 0..header_count {
220            headers.push(
221                decode_record_header(buf)
222                    .map_err(|e| RecordsError::RecordParse(format!("header[{i}]: {e}")))?,
223            );
224        }
225
226        Ok(Self {
227            attributes,
228            timestamp_delta,
229            offset_delta,
230            key,
231            value,
232            headers,
233        })
234    }
235}
236
237fn decode_nullable_bytes<B: Buf>(buf: &mut B, label: &str) -> Result<Option<Bytes>, RecordsError> {
238    let len =
239        get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("{label} length: {e}")))?;
240    if len < 0 {
241        Ok(None)
242    } else {
243        #[allow(clippy::cast_sign_loss)] // checked < 0 above
244        let n = len as usize;
245        if buf.remaining() < n {
246            return Err(RecordsError::BodyTooShort {
247                needed: n - buf.remaining(),
248            });
249        }
250        let mut v = vec![0u8; n];
251        buf.copy_to_slice(&mut v);
252        Ok(Some(Bytes::from(v)))
253    }
254}
255
256fn decode_record_header<B: Buf>(buf: &mut B) -> Result<RecordHeader, String> {
257    let key_len = get_varint(buf).map_err(|e| format!("key length: {e}"))?;
258    if key_len < 0 {
259        return Err(format!("non-nullable key has negative length {key_len}"));
260    }
261    #[allow(clippy::cast_sign_loss)] // checked < 0 above
262    let n = key_len as usize;
263    if buf.remaining() < n {
264        return Err(format!("key truncated (need {} more)", n - buf.remaining()));
265    }
266    let mut kv = vec![0u8; n];
267    buf.copy_to_slice(&mut kv);
268    let key = String::from_utf8(kv).map_err(|e| format!("key utf-8: {e}"))?;
269
270    let value_len = get_varint(buf).map_err(|e| format!("value length: {e}"))?;
271    let value = if value_len < 0 {
272        None
273    } else {
274        #[allow(clippy::cast_sign_loss)] // checked < 0 above
275        let n = value_len as usize;
276        if buf.remaining() < n {
277            return Err(format!(
278                "value truncated (need {} more)",
279                n - buf.remaining()
280            ));
281        }
282        let mut vv = vec![0u8; n];
283        buf.copy_to_slice(&mut vv);
284        Some(Bytes::from(vv))
285    };
286
287    Ok(RecordHeader { key, value })
288}
289
290#[cfg(test)]
291mod record_tests {
292    use assert2::assert;
293    use bytes::BytesMut;
294
295    use super::*;
296
297    fn fixture_minimal_record() -> Record {
298        Record {
299            attributes: 0,
300            timestamp_delta: 0,
301            offset_delta: 0,
302            key: None,
303            value: None,
304            headers: vec![],
305        }
306    }
307
308    fn fixture_keyed_record() -> Record {
309        Record {
310            attributes: 0,
311            timestamp_delta: 17,
312            offset_delta: 2,
313            key: Some(Bytes::from_static(b"the-key")),
314            value: Some(Bytes::from_static(b"hello kafka")),
315            headers: vec![
316                RecordHeader {
317                    key: "trace-id".to_string(),
318                    value: Some(Bytes::from_static(b"abc")),
319                },
320                RecordHeader {
321                    key: "null-val".to_string(),
322                    value: None,
323                },
324            ],
325        }
326    }
327
328    fn fixture_large_payload_record() -> Record {
329        Record {
330            attributes: 0,
331            timestamp_delta: 1_000_000,
332            offset_delta: 999,
333            key: Some(Bytes::from(vec![b'k'; 128])),
334            value: Some(Bytes::from(vec![b'v'; 4096])),
335            headers: vec![],
336        }
337    }
338
339    macro_rules! roundtrip {
340        ($name:ident, $fixture:ident) => {
341            #[test]
342            fn $name() {
343                let r = $fixture();
344                let mut buf = BytesMut::new();
345                r.encode(&mut buf).unwrap();
346                assert!(buf.len() == r.encoded_len(), "predicted len mismatch");
347
348                let mut cur: &[u8] = &buf[..];
349                let decoded = Record::decode(&mut cur).unwrap();
350                assert!(decoded == r);
351                assert!(cur.is_empty(), "trailing bytes after decode");
352            }
353        };
354    }
355
356    roundtrip!(minimal, fixture_minimal_record);
357    roundtrip!(keyed_with_headers, fixture_keyed_record);
358    roundtrip!(large_payload, fixture_large_payload_record);
359
360    #[test]
361    fn decode_rejects_negative_header_count() {
362        let mut buf = BytesMut::new();
363        // body: attributes(1) + timestamp_delta(1) + offset_delta(1) + key=-1(1)
364        //       + value=-1(1) + headers=-1(1) = 6 bytes body
365        put_varlong(&mut buf, 6); // body length 6 bytes
366        buf.put_i8(0); // attributes
367        put_varlong(&mut buf, 0); // timestamp_delta = 0  (1 byte)
368        put_varint(&mut buf, 0); // offset_delta = 0     (1 byte)
369        put_varint(&mut buf, -1); // key len               (1 byte)
370        put_varint(&mut buf, -1); // value len             (1 byte)
371        put_varint(&mut buf, -1); // negative header count (1 byte)
372
373        let mut cur: &[u8] = &buf[..];
374        match Record::decode(&mut cur) {
375            Err(RecordsError::RecordParse(msg)) => {
376                assert!(msg.contains("negative header count"), "got: {msg}");
377            }
378            other => panic!("expected RecordParse, got {other:?}"),
379        }
380    }
381
382    #[test]
383    fn decode_huge_header_count_does_not_overallocate() {
384        // A record that declares ~1 billion headers but supplies none. The
385        // capacity hint must be clamped to the (tiny) body remaining, and the
386        // decode must fail cleanly on EOF rather than attempting a multi-GB
387        // allocation.
388        let mut inner = BytesMut::new();
389        inner.put_i8(0); // attributes
390        put_varlong(&mut inner, 0); // timestamp_delta
391        put_varint(&mut inner, 0); // offset_delta
392        put_varint(&mut inner, -1); // key len = null
393        put_varint(&mut inner, -1); // value len = null
394        put_varint(&mut inner, 1_000_000_000); // absurd header count
395
396        let mut buf = BytesMut::new();
397        put_varlong(&mut buf, i64::try_from(inner.len()).unwrap());
398        buf.extend_from_slice(&inner);
399
400        let mut cur: &[u8] = &buf[..];
401        // Must return an Err (EOF reading the first header), not OOM.
402        assert!(Record::decode(&mut cur).is_err());
403    }
404}
405
406impl RecordBatch {
407    /// KIP-534 delete horizon, if the delete-horizon attribute bit is set.
408    /// `base_timestamp` is repurposed to carry it (no separate wire field).
409    #[must_use]
410    pub fn delete_horizon_ms(&self) -> Option<i64> {
411        self.attributes
412            .has_delete_horizon()
413            .then_some(self.base_timestamp)
414    }
415
416    /// Stamp the delete horizon: set bit 6, move the horizon into
417    /// `base_timestamp`, and rewrite every record's `timestamp_delta` so
418    /// reconstructed absolute timestamps (`base_timestamp + delta`) are
419    /// unchanged.
420    #[must_use]
421    pub fn with_delete_horizon(mut self, horizon_ms: i64) -> Self {
422        let old_base = self.base_timestamp;
423        for r in &mut self.records {
424            // Reconstruct the original absolute timestamp, then re-base it onto
425            // the new `base_timestamp`. Deltas are `i64`, so absolute timestamps
426            // round-trip exactly across re-basing; `saturating_*` only guards
427            // pathological inputs from panicking.
428            let abs = old_base.saturating_add(r.timestamp_delta);
429            r.timestamp_delta = abs.saturating_sub(horizon_ms);
430        }
431        self.base_timestamp = horizon_ms;
432        self.attributes = self.attributes.with_delete_horizon(true);
433        self
434    }
435
436    /// Decode a complete v2 record batch from `buf`. Reads from the start of
437    /// the header.
438    pub fn decode<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
439        // batch_length field semantics: bytes after itself.
440        // Header tail = partition_leader_epoch(4) + magic(1) + crc(4) +
441        //   attributes(2) + last_offset_delta(4) + base_timestamp(8) +
442        //   max_timestamp(8) + producer_id(8) + producer_epoch(2) +
443        //   base_sequence(4) + records_count(4) = 49 bytes.
444        const HEADER_TAIL_LEN: i32 = 49;
445
446        // Need the full header before doing anything.
447        if buf.remaining() < HEADER_LEN {
448            return Err(RecordsError::HeaderTooShort {
449                needed: HEADER_LEN - buf.remaining(),
450            });
451        }
452        // Copy out the header to a stack buffer so we can use zerocopy.
453        let mut hdr_bytes = [0u8; HEADER_LEN];
454        buf.copy_to_slice(&mut hdr_bytes);
455
456        let hdr = crate::records::header::RecordBatchHeader::ref_from_bytes(&hdr_bytes[..])
457            .map_err(|_| RecordsError::ZerocopyFailure)?;
458
459        if hdr.magic != 2 {
460            return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
461        }
462
463        // body_len = batch_length - HEADER_TAIL_LEN
464        let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
465            .and_then(|n| usize::try_from(n).ok())
466            .ok_or_else(|| {
467                RecordsError::RecordParse("negative or oversized batch_length".into())
468            })?;
469
470        if buf.remaining() < body_len {
471            return Err(RecordsError::BodyTooShort {
472                needed: body_len - buf.remaining(),
473            });
474        }
475
476        // Read the (possibly compressed) body.
477        let mut body = vec![0u8; body_len];
478        buf.copy_to_slice(&mut body);
479
480        // CRC is computed over: header bytes 21..HEADER_LEN (attributes through
481        // records_count), then the body bytes.
482        let expected_crc = hdr.crc.get();
483        let mut computed = crc32c(&hdr_bytes[21..HEADER_LEN]);
484        computed = crc32c_append(computed, &body);
485        if computed != expected_crc {
486            return Err(RecordsError::CrcMismatch {
487                expected: expected_crc,
488                computed,
489            });
490        }
491
492        let attributes = Attributes(hdr.attributes.get());
493        let codec = attributes.compression();
494
495        // Decompress body if needed.
496        let body_for_records: Bytes = if codec == crabka_compression::CompressionType::None {
497            Bytes::from(body)
498        } else {
499            // Bound decompressed output: generous vs. legit ratios, but finite.
500            // A small compressed batch must not be able to expand to gigabytes
501            // and OOM the broker (decompression bomb).
502            const DECOMPRESS_MIN_CAP: usize = 16 * 1024 * 1024; // 16 MiB floor (small inputs)
503            const DECOMPRESS_MAX_RATIO: usize = 100; // ≤100x the compressed size
504            const DECOMPRESS_ABSOLUTE_CEILING: usize = 1024 * 1024 * 1024; // 1 GiB hard ceiling
505            let max_output = body
506                .len()
507                .saturating_mul(DECOMPRESS_MAX_RATIO)
508                .clamp(DECOMPRESS_MIN_CAP, DECOMPRESS_ABSOLUTE_CEILING);
509            crabka_compression::decompress(codec, &body, max_output)?
510        };
511
512        // Parse records.
513        let count = hdr.records_count.get();
514        if count < 0 {
515            return Err(RecordsError::RecordParse(format!(
516                "negative records_count {count}"
517            )));
518        }
519        let mut body_cur: &[u8] = &body_for_records[..];
520        // Bound pre-allocation: each record is at least 1 byte in the
521        // decompressed body, so an honest `records_count` can never exceed the
522        // body length. Clamp the capacity hint to reject huge declared counts.
523        #[allow(clippy::cast_sign_loss)] // checked < 0 above
524        let mut records = Vec::with_capacity((count as usize).min(body_for_records.len()));
525        for i in 0..count {
526            // Parse the record into borrowed slices, then materialise its
527            // key / value / header-values as zero-copy `Bytes` views into
528            // `body_for_records`. The whole batch shares that single backing
529            // allocation instead of copying every field into its own `Bytes`,
530            // which removes ~2 heap allocations + memcpys per record on the
531            // common (header-free) path.
532            let r = crate::records::borrowed::parse_one_record(&mut body_cur)
533                .map_err(|e| RecordsError::RecordParse(format!("record[{i}]: {e}")))?;
534            records.push(Record {
535                attributes: r.attributes,
536                timestamp_delta: r.timestamp_delta,
537                offset_delta: r.offset_delta,
538                key: r.key.map(|s| body_for_records.slice_ref(s)),
539                value: r.value.map(|s| body_for_records.slice_ref(s)),
540                headers: r
541                    .headers
542                    .into_iter()
543                    .map(|h| RecordHeader {
544                        key: h.key.to_string(),
545                        value: h.value.map(|s| body_for_records.slice_ref(s)),
546                    })
547                    .collect(),
548            });
549        }
550        if !body_cur.is_empty() {
551            return Err(RecordsError::RecordParse(format!(
552                "trailing bytes after records (left={})",
553                body_cur.len()
554            )));
555        }
556
557        Ok(Self {
558            base_offset: hdr.base_offset.get(),
559            partition_leader_epoch: hdr.partition_leader_epoch.get(),
560            attributes,
561            last_offset_delta: hdr.last_offset_delta.get(),
562            base_timestamp: hdr.base_timestamp.get(),
563            max_timestamp: hdr.max_timestamp.get(),
564            producer_id: hdr.producer_id.get(),
565            producer_epoch: hdr.producer_epoch.get(),
566            base_sequence: hdr.base_sequence.get(),
567            records,
568        })
569    }
570
571    /// Encode this batch into `buf`.
572    pub fn encode<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
573        const HEADER_TAIL_LEN: i32 = 49;
574
575        // 1. Encode records into a temporary buffer.
576        let mut raw_body =
577            BytesMut::with_capacity(self.records.iter().map(Record::encoded_len).sum());
578        for r in &self.records {
579            r.encode(&mut raw_body)?;
580        }
581        let raw_body = raw_body.freeze();
582
583        // 2. Compress if needed.
584        let codec = self.attributes.compression();
585        let body: Bytes = if codec == crabka_compression::CompressionType::None {
586            raw_body
587        } else {
588            crabka_compression::compress(codec, &raw_body)?
589        };
590
591        // 3. batch_length = HEADER_TAIL_LEN + body_len
592        let batch_length = HEADER_TAIL_LEN
593            + i32::try_from(body.len())
594                .map_err(|_| RecordsError::RecordParse("body length exceeds i32".into()))?;
595
596        // 4. Build the CRC-covered header portion (attributes through records_count = 40 bytes).
597        let mut covered = BytesMut::with_capacity(40);
598        covered.put_i16(self.attributes.0);
599        covered.put_i32(self.last_offset_delta);
600        covered.put_i64(self.base_timestamp);
601        covered.put_i64(self.max_timestamp);
602        covered.put_i64(self.producer_id);
603        covered.put_i16(self.producer_epoch);
604        covered.put_i32(self.base_sequence);
605        covered.put_i32(
606            i32::try_from(self.records.len())
607                .map_err(|_| RecordsError::RecordParse("records_count exceeds i32".into()))?,
608        );
609        let covered_head = covered.freeze();
610
611        // 5. Compute CRC over covered_head then body.
612        let mut crc = crc32c(&covered_head);
613        crc = crc32c_append(crc, &body);
614
615        // 6. Emit the full header then body.
616        buf.put_i64(self.base_offset);
617        buf.put_i32(batch_length);
618        buf.put_i32(self.partition_leader_epoch);
619        buf.put_i8(2); // magic v2
620        buf.put_u32(crc);
621        buf.put_slice(&covered_head);
622        buf.put_slice(&body);
623        Ok(())
624    }
625
626    /// Predicted total bytes that `encode` will write (uncompressed; for
627    /// compressed batches the actual size will differ).
628    pub fn encoded_len(&self) -> usize {
629        let body: usize = self.records.iter().map(Record::encoded_len).sum();
630        HEADER_LEN + body
631    }
632}
633
634#[cfg(test)]
635mod batch_tests {
636    use assert2::{assert, check};
637    use crabka_compression::CompressionType;
638
639    use super::*;
640
641    fn fixture_empty_batch() -> RecordBatch {
642        RecordBatch::default()
643    }
644
645    fn fixture_single_record_batch() -> RecordBatch {
646        RecordBatch {
647            records: vec![Record {
648                key: Some(Bytes::from_static(b"k1")),
649                value: Some(Bytes::from_static(b"v1")),
650                ..Default::default()
651            }],
652            ..RecordBatch::default()
653        }
654    }
655
656    fn fixture_multi_record_batch() -> RecordBatch {
657        RecordBatch {
658            base_offset: 42,
659            partition_leader_epoch: 5,
660            last_offset_delta: 2,
661            base_timestamp: 1_700_000_000,
662            max_timestamp: 1_700_000_500,
663            producer_id: 100,
664            producer_epoch: 3,
665            base_sequence: 7,
666            records: vec![
667                Record {
668                    offset_delta: 0,
669                    timestamp_delta: 0,
670                    key: Some(Bytes::from_static(b"a")),
671                    value: Some(Bytes::from_static(b"1")),
672                    ..Default::default()
673                },
674                Record {
675                    offset_delta: 1,
676                    timestamp_delta: 100,
677                    key: Some(Bytes::from_static(b"b")),
678                    value: Some(Bytes::from_static(b"2")),
679                    ..Default::default()
680                },
681                Record {
682                    offset_delta: 2,
683                    timestamp_delta: 500,
684                    key: None,
685                    value: Some(Bytes::from_static(b"3")),
686                    headers: vec![RecordHeader {
687                        key: "h".to_string(),
688                        value: Some(Bytes::from_static(b"hv")),
689                    }],
690                    ..Default::default()
691                },
692            ],
693            ..RecordBatch::default()
694        }
695    }
696
697    macro_rules! roundtrip_uncompressed {
698        ($name:ident, $fixture:ident) => {
699            #[test]
700            fn $name() {
701                let mut b = $fixture();
702                b.attributes = b.attributes.with_compression(CompressionType::None);
703
704                let mut buf = BytesMut::new();
705                b.encode(&mut buf).unwrap();
706                assert!(buf.len() == b.encoded_len());
707
708                let mut cur: &[u8] = &buf[..];
709                let decoded = RecordBatch::decode(&mut cur).unwrap();
710                assert!(decoded == b);
711                assert!(cur.is_empty());
712            }
713        };
714    }
715
716    roundtrip_uncompressed!(uncompressed_empty, fixture_empty_batch);
717    roundtrip_uncompressed!(uncompressed_single, fixture_single_record_batch);
718    roundtrip_uncompressed!(uncompressed_multi, fixture_multi_record_batch);
719
720    #[test]
721    fn rejects_pre_v2_magic() {
722        let mut buf = BytesMut::new();
723        buf.put_i64(0); // base_offset
724        buf.put_i32(49); // batch_length
725        buf.put_i32(0); // partition_leader_epoch
726        buf.put_i8(1); // magic = 1 (v1, deprecated)
727        buf.put_u32(0); // crc (irrelevant; we reject on magic first)
728        for _ in 21..HEADER_LEN {
729            buf.put_u8(0);
730        }
731        let mut cur: &[u8] = &buf[..];
732        assert!(matches!(
733            RecordBatch::decode(&mut cur),
734            Err(RecordsError::UnsupportedMagic { found: 1 })
735        ));
736    }
737
738    #[test]
739    fn rejects_bad_crc() {
740        let b = fixture_single_record_batch();
741        let mut buf = BytesMut::new();
742        b.encode(&mut buf).unwrap();
743        // Corrupt the CRC bytes (offsets 17..21).
744        buf[17] ^= 0xFF;
745        let mut cur: &[u8] = &buf[..];
746        assert!(matches!(
747            RecordBatch::decode(&mut cur),
748            Err(RecordsError::CrcMismatch { .. })
749        ));
750    }
751
752    macro_rules! roundtrip_compressed {
753        ($name:ident, $codec:expr) => {
754            #[test]
755            fn $name() {
756                let mut b = fixture_multi_record_batch();
757                b.attributes = b.attributes.with_compression($codec);
758
759                let mut buf = BytesMut::new();
760                b.encode(&mut buf).unwrap();
761                let mut cur: &[u8] = &buf[..];
762                let decoded = RecordBatch::decode(&mut cur).unwrap();
763                assert!(decoded == b);
764                assert!(cur.is_empty());
765            }
766        };
767    }
768
769    roundtrip_compressed!(compressed_gzip, CompressionType::Gzip);
770    roundtrip_compressed!(compressed_snappy, CompressionType::Snappy);
771    roundtrip_compressed!(compressed_lz4, CompressionType::Lz4);
772    roundtrip_compressed!(compressed_zstd, CompressionType::Zstd);
773
774    #[test]
775    fn with_delete_horizon_stamps_and_preserves_record_timestamps() {
776        // base 1000, two records at deltas [0, 5] → absolutes [1000, 1005].
777        let b = RecordBatch {
778            base_timestamp: 1000,
779            records: vec![
780                Record {
781                    timestamp_delta: 0,
782                    ..Default::default()
783                },
784                Record {
785                    timestamp_delta: 5,
786                    ..Default::default()
787                },
788            ],
789            ..RecordBatch::default()
790        };
791
792        let stamped = b.with_delete_horizon(9999);
793
794        check!(stamped.attributes.has_delete_horizon());
795        check!(stamped.base_timestamp == 9999);
796        check!(stamped.delete_horizon_ms() == Some(9999));
797
798        // Reconstructed absolutes (base + delta) must equal the ORIGINALS.
799        let absolutes: Vec<i64> = stamped
800            .records
801            .iter()
802            .map(|r| stamped.base_timestamp + r.timestamp_delta)
803            .collect();
804        assert!(absolutes == vec![1000, 1005]);
805    }
806
807    #[test]
808    fn delete_horizon_round_trips_through_encode_decode() {
809        // base 1000, two keyed records at deltas [0, 5] → absolutes [1000, 1005].
810        let b = RecordBatch {
811            base_timestamp: 1000,
812            last_offset_delta: 1,
813            records: vec![
814                Record {
815                    timestamp_delta: 0,
816                    offset_delta: 0,
817                    key: Some(Bytes::from_static(b"k1")),
818                    value: Some(Bytes::from_static(b"v1")),
819                    ..Default::default()
820                },
821                Record {
822                    timestamp_delta: 5,
823                    offset_delta: 1,
824                    key: Some(Bytes::from_static(b"k2")),
825                    value: Some(Bytes::from_static(b"v2")),
826                    ..Default::default()
827                },
828            ],
829            ..RecordBatch::default()
830        }
831        .with_delete_horizon(9999);
832
833        let mut buf = BytesMut::new();
834        b.encode(&mut buf).unwrap();
835
836        let mut cur: &[u8] = &buf[..];
837        let decoded = RecordBatch::decode(&mut cur).unwrap();
838        assert!(cur.is_empty());
839
840        assert!(decoded.delete_horizon_ms() == Some(9999));
841
842        let absolutes: Vec<i64> = decoded
843            .records
844            .iter()
845            .map(|r| decoded.base_timestamp + r.timestamp_delta)
846            .collect();
847        assert!(absolutes == vec![1000, 1005]);
848    }
849
850    #[test]
851    fn decode_huge_records_count_does_not_overallocate() {
852        // Encode an empty uncompressed batch, then overwrite records_count with
853        // an absurd value (~1 billion) and fix up the CRC. The capacity hint
854        // must be clamped to the (empty) decompressed body, and decode must
855        // fail cleanly on EOF rather than attempting a multi-GB allocation.
856        let mut b = fixture_empty_batch();
857        b.attributes = b.attributes.with_compression(CompressionType::None);
858        let mut buf = BytesMut::new();
859        b.encode(&mut buf).unwrap();
860
861        // records_count is the last 4 bytes of the fixed header.
862        let rc_off = HEADER_LEN - 4;
863        buf[rc_off..HEADER_LEN].copy_from_slice(&1_000_000_000i32.to_be_bytes());
864
865        // Recompute CRC over header[21..HEADER_LEN] + body (body is empty here).
866        let body = &buf[HEADER_LEN..];
867        let mut computed = crc32c(&buf[21..HEADER_LEN]);
868        computed = crc32c_append(computed, body);
869        buf[17..21].copy_from_slice(&computed.to_be_bytes());
870
871        let mut cur: &[u8] = &buf[..];
872        // Must return an Err (EOF reading the first record), not OOM.
873        assert!(RecordBatch::decode(&mut cur).is_err());
874    }
875}
876
877impl crate::Encode for RecordBatch {
878    fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
879        RecordBatch::encode(self, buf).map_err(Into::into)
880    }
881
882    fn encoded_len(&self, _version: i16) -> usize {
883        RecordBatch::encoded_len(self)
884    }
885}
886
887impl crate::Decode<'_> for RecordBatch {
888    fn decode<B: Buf>(buf: &mut B, _version: i16) -> Result<Self, crate::ProtocolError> {
889        RecordBatch::decode(buf).map_err(Into::into)
890    }
891}