Skip to main content

krabka_protocol/records/
borrowed.rs

1//! Borrowed `RecordBatch<'a>`, `Record<'a>`, and `RecordHeader<'a>`.
2
3use bytes::Bytes;
4use krabka_compression::RecordDecompressionPolicy;
5use krabka_units::{ByteSize, convert::ByteSizeExt as _};
6use zerocopy::FromBytes as _;
7
8use crate::{
9    primitives::varint::{get_varint, get_varlong},
10    records::{
11        RecordsError,
12        crc::crc32c,
13        header::{Attributes, HEADER_LEN, RecordBatchHeader},
14    },
15};
16
17// batch_length field semantics: bytes after itself.
18// Header tail = partition_leader_epoch(4) + magic(1) + crc(4) +
19//   attributes(2) + last_offset_delta(4) + base_timestamp(8) +
20//   max_timestamp(8) + producer_id(8) + producer_epoch(2) +
21//   base_sequence(4) + records_count(4) = 49 bytes.
22const HEADER_TAIL_LEN: i32 = 49;
23
24pub struct RecordBatch<'a> {
25    pub(crate) header: &'a RecordBatchHeader,
26    pub(crate) body: RecordBody<'a>,
27}
28
29pub(crate) enum RecordBody<'a> {
30    Borrowed(&'a [u8]),
31    Owned(Bytes),
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct Record<'a> {
36    pub attributes: i8,
37    pub timestamp_delta: i64,
38    pub offset_delta: i32,
39    pub key: Option<&'a [u8]>,
40    pub value: Option<&'a [u8]>,
41    pub headers: Vec<RecordHeader<'a>>,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct RecordHeader<'a> {
46    pub key: &'a str,
47    pub value: Option<&'a [u8]>,
48}
49
50impl Record<'_> {
51    /// Size of this record's body in bytes: what Kafka's `DefaultRecord` calls
52    /// `sizeOfBodyInBytes`, and what its varlong length prefix declares.
53    ///
54    /// The body is the attributes byte, the two deltas, and the key, the value
55    /// and the headers, each behind its own length prefix. A null key, value or
56    /// header value counts as the one-byte `-1` marker. The length prefix
57    /// itself is not counted. This is the size the owned
58    /// [`crate::records::Record::body_len`] reports for the same record.
59    #[must_use]
60    pub fn body_len(&self) -> usize {
61        super::owned::body_len_of(
62            self.timestamp_delta,
63            self.offset_delta,
64            self.key.map(<[u8]>::len),
65            self.value.map(<[u8]>::len),
66            self.headers
67                .iter()
68                .map(|h| (h.key.len(), h.value.map(<[u8]>::len))),
69        )
70    }
71}
72
73impl RecordBatch<'_> {
74    #[must_use]
75    pub fn header(&self) -> &RecordBatchHeader {
76        self.header
77    }
78
79    #[must_use]
80    pub fn attributes(&self) -> Attributes {
81        Attributes(self.header.attributes.get())
82    }
83}
84
85impl<'a> Default for RecordBatch<'a> {
86    /// Returns an empty record batch with a zeroed header and no records body.
87    /// Use it in generated `Default` impls and round-trip tests. It is not
88    /// suitable for a real Kafka record batch.
89    fn default() -> Self {
90        use zerocopy::FromZeros as _;
91        // RecordBatchHeader derives zerocopy::FromZeros (via FromBytes), so zeroing is safe.
92        let header: &'a RecordBatchHeader = Box::leak(Box::new(RecordBatchHeader::new_zeroed()));
93        Self {
94            header,
95            body: RecordBody::Owned(bytes::Bytes::new()),
96        }
97    }
98}
99
100// ── Decode ────────────────────────────────────────────────────────────────────
101
102impl<'de> crate::DecodeBorrow<'de> for RecordBatch<'de> {
103    fn decode_borrow(buf: &mut &'de [u8], _version: i16) -> Result<Self, crate::ProtocolError> {
104        Self::decode_borrow_with_policy(buf, RecordDecompressionPolicy::default())
105            .map_err(Into::into)
106    }
107}
108
109impl<'de> RecordBatch<'de> {
110    /// Decodes a borrowed v2 record batch with explicit decompression limits.
111    ///
112    /// # Errors
113    ///
114    /// Returns the underlying records error for malformed, truncated, corrupt,
115    /// or over-limit input.
116    pub fn decode_borrow_with_policy(
117        buf: &mut &'de [u8],
118        policy: RecordDecompressionPolicy,
119    ) -> Result<Self, RecordsError> {
120        decode_borrow_impl(buf, policy)
121    }
122}
123
124fn decode_borrow_impl<'de>(
125    buf: &mut &'de [u8],
126    policy: RecordDecompressionPolicy,
127) -> Result<RecordBatch<'de>, RecordsError> {
128    if buf.len() < HEADER_LEN {
129        return Err(RecordsError::HeaderTooShort {
130            needed: HEADER_LEN - buf.len(),
131        });
132    }
133    // Split off the header slice — both slices remain tied to 'de.
134    let (hdr_slice, rest) = buf.split_at(HEADER_LEN);
135    let hdr: &'de RecordBatchHeader =
136        RecordBatchHeader::ref_from_bytes(hdr_slice).map_err(|_| RecordsError::ZerocopyFailure)?;
137    if hdr.magic != 2 {
138        return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
139    }
140
141    let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
142        .and_then(|n| usize::try_from(n).ok())
143        .ok_or_else(|| RecordsError::RecordParse("negative or oversized batch_length".into()))?;
144
145    if rest.len() < body_len {
146        return Err(RecordsError::BodyTooShort {
147            needed: body_len - rest.len(),
148        });
149    }
150    let (raw_body, after) = rest.split_at(body_len);
151    // The covered header and raw body are contiguous in the input batch.
152    let computed = crc32c(&buf[21..HEADER_LEN + body_len]);
153    *buf = after;
154
155    let expected = hdr.crc.get();
156    if computed != expected {
157        return Err(RecordsError::CrcMismatch { expected, computed });
158    }
159
160    let attributes = Attributes(hdr.attributes.get());
161    let codec = attributes.compression();
162    let body = if codec == krabka_compression::CompressionType::None {
163        RecordBody::Borrowed(raw_body)
164    } else {
165        let decompressed = krabka_compression::decompress(
166            codec,
167            raw_body,
168            policy.output_limit(ByteSize::from_bytes(raw_body.len() as u64)),
169        )?;
170        RecordBody::Owned(decompressed)
171    };
172
173    Ok(RecordBatch { header: hdr, body })
174}
175
176/// A single complete v2 batch located within a larger buffer, together with
177/// its already-validated header.
178///
179/// [`validate_one_v2_batch`] returns this type.
180#[derive(Debug)]
181pub struct ValidatedBatch<'a> {
182    /// The fixed 61-byte header, reinterpreted in place, zero-copy.
183    pub header: &'a RecordBatchHeader,
184    /// Total on-disk and wire length of this batch in bytes, which is
185    /// `12 + batch_length`, that is the header plus the body.
186    pub total_len: usize,
187    raw_body: &'a [u8],
188}
189
190impl ValidatedBatch<'_> {
191    /// Validates that the body contains exactly the number of structurally
192    /// valid records declared by the batch header.
193    ///
194    /// Uncompressed records are parsed directly from the input slice. A
195    /// compressed body is decompressed once under `policy`, parsed, and then
196    /// discarded; the original compressed bytes remain available to callers
197    /// for verbatim append.
198    ///
199    /// # Errors
200    ///
201    /// Returns a records or compression error when the body is malformed,
202    /// exceeds the decompression policy, contains fewer records than declared,
203    /// or has trailing bytes after the declared records.
204    pub fn validate_records(&self, policy: RecordDecompressionPolicy) -> Result<(), RecordsError> {
205        self.validate_records_with(policy, |_| {})
206    }
207
208    /// Validates the records and calls `visit` for each borrowed record.
209    ///
210    /// This combines structural validation with callers' metadata checks while
211    /// keeping compressed batches on the borrowed path.
212    ///
213    /// # Errors
214    ///
215    /// Returns the same errors as [`Self::validate_records`].
216    pub fn validate_records_with(
217        &self,
218        policy: RecordDecompressionPolicy,
219        mut visit: impl FnMut(&Record<'_>),
220    ) -> Result<(), RecordsError> {
221        let count = self.header.records_count.get();
222        if count < 0 {
223            return Err(RecordsError::RecordParse(format!(
224                "negative records_count {count}"
225            )));
226        }
227
228        let codec = Attributes(self.header.attributes.get()).compression();
229        let decompressed = if codec == krabka_compression::CompressionType::None {
230            None
231        } else {
232            Some(krabka_compression::decompress(
233                codec,
234                self.raw_body,
235                policy.output_limit(ByteSize::from_bytes(self.raw_body.len() as u64)),
236            )?)
237        };
238        let mut remaining = decompressed.as_deref().unwrap_or(self.raw_body);
239        for i in 0..count {
240            let record = parse_one_record(&mut remaining)
241                .map_err(|error| RecordsError::RecordParse(format!("record[{i}]: {error}")))?;
242            visit(&record);
243        }
244        if !remaining.is_empty() {
245            return Err(RecordsError::RecordParse(format!(
246                "trailing bytes after records (left={})",
247                remaining.len()
248            )));
249        }
250        Ok(())
251    }
252}
253
254/// Validates the header, declared length, and CRC of one v2 record batch at
255/// the start of `buf`.
256///
257/// This function materializes no records and decompresses no body.
258///
259/// This is the first stage of the produce passthrough fast path. It
260/// reinterprets the fixed header in place, zero-copy, checks `magic == 2`, and
261/// verifies the producer's CRC over `header[21..61] ++ raw_body`, which is the
262/// compressed body bytes exactly as stored. Call
263/// [`ValidatedBatch::validate_records`] before accepting producer input.
264///
265/// Returns the CRC-validated header, which the caller needs for offset
266/// stamping and for idempotent and transactional gating, and the batch's total
267/// byte length, which lets the caller slice out the verbatim bytes.
268///
269/// # Errors
270///
271/// - [`RecordsError::HeaderTooShort`] or [`RecordsError::BodyTooShort`]
272///   when `buf` does not contain a whole batch.
273/// - [`RecordsError::UnsupportedMagic`] for a non-v2 batch.
274/// - [`RecordsError::CrcMismatch`] when the stored CRC does not match.
275pub fn validate_one_v2_batch(buf: &[u8]) -> Result<ValidatedBatch<'_>, RecordsError> {
276    if buf.len() < HEADER_LEN {
277        return Err(RecordsError::HeaderTooShort {
278            needed: HEADER_LEN - buf.len(),
279        });
280    }
281    let (hdr_slice, rest) = buf.split_at(HEADER_LEN);
282    let hdr: &RecordBatchHeader =
283        RecordBatchHeader::ref_from_bytes(hdr_slice).map_err(|_| RecordsError::ZerocopyFailure)?;
284    if hdr.magic != 2 {
285        return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
286    }
287    let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
288        .and_then(|n| usize::try_from(n).ok())
289        .ok_or_else(|| RecordsError::RecordParse("negative or oversized batch_length".into()))?;
290    if rest.len() < body_len {
291        return Err(RecordsError::BodyTooShort {
292            needed: body_len - rest.len(),
293        });
294    }
295    let raw_body = &rest[..body_len];
296
297    let expected = hdr.crc.get();
298    let computed = crc32c(&buf[21..HEADER_LEN + body_len]);
299    if computed != expected {
300        return Err(RecordsError::CrcMismatch { expected, computed });
301    }
302
303    Ok(ValidatedBatch {
304        header: hdr,
305        total_len: HEADER_LEN + body_len,
306        raw_body,
307    })
308}
309
310/// Sums the `records_count` header field of every concatenated v2 batch in
311/// `buf`, and decompresses or parses no record body.
312///
313/// This is intentionally a best-effort metrics helper. Non-v2 input returns 0.
314/// Malformed or truncated input contributes every complete v2 batch before the
315/// first bad boundary. A negative `records_count` contributes 0. Real append
316/// validation still goes through [`validate_one_v2_batch`].
317#[must_use]
318pub fn count_records_in_v2_batches(buf: &[u8]) -> u64 {
319    // Magic byte sits at offset 16; only v2 (magic == 2) carries a
320    // `records_count` header. Legacy slices contribute 0 here.
321    const MAGIC_OFFSET: usize = 16;
322    if buf.len() <= MAGIC_OFFSET || buf[MAGIC_OFFSET] != 2 {
323        return 0;
324    }
325
326    let mut total = 0u64;
327    let mut remaining = buf;
328    while remaining.len() >= HEADER_LEN {
329        if remaining[MAGIC_OFFSET] != 2 {
330            break;
331        }
332        let Ok(hdr) = RecordBatchHeader::ref_from_bytes(&remaining[..HEADER_LEN]) else {
333            break;
334        };
335        let Ok(after_len) = usize::try_from(hdr.batch_length.get()) else {
336            break;
337        };
338        let total_len = 12 + after_len;
339        if total_len < HEADER_LEN || total_len > remaining.len() {
340            break;
341        }
342        total += u64::try_from(hdr.records_count.get().max(0)).unwrap_or(0);
343        remaining = &remaining[total_len..];
344    }
345    total
346}
347
348// ── Iteration ─────────────────────────────────────────────────────────────────
349
350impl RecordBatch<'_> {
351    /// Iterates over the records and parses each one lazily.
352    ///
353    /// The returned `Record<'b>` items borrow from `self` with lifetime `'b`,
354    /// not from the original input buffer. For uncompressed batches the
355    /// backing memory is the input buffer. For compressed batches it is the
356    /// batch's internal decompressed `Bytes`.
357    /// # Panics
358    /// Panics if a value previously validated by the protocol type no longer satisfies its encoded-length or field-range invariant.
359    pub fn iter(&self) -> RecordIter<'_> {
360        let body: &[u8] = match &self.body {
361            RecordBody::Borrowed(s) => s,
362            RecordBody::Owned(b) => b.as_ref(),
363        };
364        let count = usize::try_from(self.header.records_count.get().max(0))
365            .expect("non-negative record count must fit in usize");
366        RecordIter {
367            remaining: body,
368            count,
369            index: 0,
370        }
371    }
372}
373
374impl<'a> IntoIterator for &'a RecordBatch<'_> {
375    type Item = Result<Record<'a>, RecordsError>;
376    type IntoIter = RecordIter<'a>;
377
378    fn into_iter(self) -> Self::IntoIter {
379        self.iter()
380    }
381}
382
383pub struct RecordIter<'a> {
384    remaining: &'a [u8],
385    count: usize,
386    index: usize,
387}
388
389impl<'a> Iterator for RecordIter<'a> {
390    type Item = Result<Record<'a>, RecordsError>;
391
392    fn next(&mut self) -> Option<Self::Item> {
393        if self.index >= self.count {
394            return None;
395        }
396        self.index += 1;
397        Some(parse_one_record(&mut self.remaining))
398    }
399}
400
401#[inline]
402pub(crate) fn parse_one_record<'a>(buf: &mut &'a [u8]) -> Result<Record<'a>, RecordsError> {
403    let body_len =
404        get_varlong(buf).map_err(|e| RecordsError::RecordParse(format!("record length: {e}")))?;
405    let body_len = usize::try_from(body_len).map_err(|_| {
406        RecordsError::RecordParse(format!("record length negative or too large: {body_len}"))
407    })?;
408    if buf.len() < body_len {
409        return Err(RecordsError::BodyTooShort {
410            needed: body_len - buf.len(),
411        });
412    }
413    let (body, rest) = buf.split_at(body_len);
414    *buf = rest;
415    let mut body_cur = body;
416    let r = parse_body(&mut body_cur)?;
417    if !body_cur.is_empty() {
418        return Err(RecordsError::RecordParse(format!(
419            "trailing bytes inside record (left={})",
420            body_cur.len()
421        )));
422    }
423    Ok(r)
424}
425
426fn parse_body<'a>(buf: &mut &'a [u8]) -> Result<Record<'a>, RecordsError> {
427    if buf.is_empty() {
428        return Err(RecordsError::RecordParse("record body empty".into()));
429    }
430    let attributes = i8::from_ne_bytes([buf[0]]);
431    *buf = &buf[1..];
432    let timestamp_delta =
433        get_varlong(buf).map_err(|e| RecordsError::RecordParse(format!("timestamp_delta: {e}")))?;
434    let offset_delta =
435        get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("offset_delta: {e}")))?;
436
437    let key = read_nullable_slice(buf, "key")?;
438    let value = read_nullable_slice(buf, "value")?;
439
440    let header_count =
441        get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("header_count: {e}")))?;
442    if header_count < 0 {
443        return Err(RecordsError::RecordParse(format!(
444            "negative header count {header_count}"
445        )));
446    }
447    // Bound pre-allocation: a record header is at least 1 byte on the wire, so
448    // an honest `header_count` can never exceed the bytes left in the record
449    // body. Clamp the capacity hint to reject huge declared counts.
450    let header_capacity = usize::try_from(header_count)
451        .expect("non-negative header count must fit in usize")
452        .min(buf.len());
453    let mut headers = Vec::with_capacity(header_capacity);
454    for i in 0..header_count {
455        let key_len = get_varint(buf)
456            .map_err(|e| RecordsError::RecordParse(format!("header[{i}] key length: {e}")))?;
457        if key_len < 0 {
458            return Err(RecordsError::RecordParse(format!(
459                "header[{i}] negative key length"
460            )));
461        }
462        let n = usize::try_from(key_len).expect("non-negative key length must fit in usize");
463        if buf.len() < n {
464            return Err(RecordsError::BodyTooShort {
465                needed: n - buf.len(),
466            });
467        }
468        let (key_bytes, rest) = buf.split_at(n);
469        *buf = rest;
470        let key_str = std::str::from_utf8(key_bytes)
471            .map_err(|e| RecordsError::RecordParse(format!("header[{i}] key utf-8: {e}")))?;
472
473        let value = read_nullable_slice(buf, &format!("header[{i}] value"))?;
474        headers.push(RecordHeader {
475            key: key_str,
476            value,
477        });
478    }
479
480    Ok(Record {
481        attributes,
482        timestamp_delta,
483        offset_delta,
484        key,
485        value,
486        headers,
487    })
488}
489
490fn read_nullable_slice<'a>(
491    buf: &mut &'a [u8],
492    label: &str,
493) -> Result<Option<&'a [u8]>, RecordsError> {
494    let len =
495        get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("{label} length: {e}")))?;
496    if len < 0 {
497        Ok(None)
498    } else {
499        let n = usize::try_from(len).expect("non-negative value length must fit in usize");
500        if buf.len() < n {
501            return Err(RecordsError::BodyTooShort {
502                needed: n - buf.len(),
503            });
504        }
505        let (head, rest) = buf.split_at(n);
506        *buf = rest;
507        Ok(Some(head))
508    }
509}
510
511// ── to_owned bridge ───────────────────────────────────────────────────────────
512
513impl RecordBatch<'_> {
514    /// Materializes an owned `RecordBatch` and copies every byte slice into
515    /// `Bytes` or `String`.
516    /// # Errors
517    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
518    pub fn to_owned(&self) -> Result<super::owned::RecordBatch, RecordsError> {
519        let mut records = Vec::new();
520        for r in self {
521            let r = r?;
522            records.push(super::owned::Record {
523                attributes: r.attributes,
524                timestamp_delta: r.timestamp_delta,
525                offset_delta: r.offset_delta,
526                key: r.key.map(Bytes::copy_from_slice),
527                value: r.value.map(Bytes::copy_from_slice),
528                headers: r
529                    .headers
530                    .into_iter()
531                    .map(|h| super::owned::RecordHeader {
532                        key: h.key.to_string(),
533                        value: h.value.map(Bytes::copy_from_slice),
534                    })
535                    .collect(),
536            });
537        }
538        Ok(super::owned::RecordBatch {
539            base_offset: self.header.base_offset.get(),
540            partition_leader_epoch: self.header.partition_leader_epoch.get(),
541            attributes: self.attributes(),
542            last_offset_delta: self.header.last_offset_delta.get(),
543            base_timestamp: self.header.base_timestamp.get(),
544            max_timestamp: self.header.max_timestamp.get(),
545            producer_id: self.header.producer_id.get(),
546            producer_epoch: self.header.producer_epoch.get(),
547            base_sequence: self.header.base_sequence.get(),
548            records,
549        })
550    }
551}
552
553// ── Debug / Clone / PartialEq / Eq ────────────────────────────────────────────
554//
555// RecordBatch<'a> holds a header reference and a body slice/bytes, so we can't
556// #[derive] these. We provide hand-rolled impls that delegate to the owned type.
557
558impl std::fmt::Debug for RecordBatch<'_> {
559    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
560        match self.to_owned() {
561            Ok(o) => o.fmt(f),
562            Err(e) => write!(f, "RecordBatch(<decode error: {e}>)"),
563        }
564    }
565}
566
567impl Clone for RecordBatch<'_> {
568    /// Shallow clone: both `header` and `body` share the same underlying data
569    /// as `self`. For a `Borrowed` body this is a reference copy. For an
570    /// `Owned` body, `Bytes::clone` is a cheap reference-count bump.
571    fn clone(&self) -> Self {
572        RecordBatch {
573            header: self.header,
574            body: match &self.body {
575                RecordBody::Borrowed(s) => RecordBody::Borrowed(s),
576                RecordBody::Owned(b) => RecordBody::Owned(b.clone()),
577            },
578        }
579    }
580}
581
582impl PartialEq for RecordBatch<'_> {
583    fn eq(&self, other: &Self) -> bool {
584        match (self.to_owned(), other.to_owned()) {
585            (Ok(a), Ok(b)) => a == b,
586            _ => false,
587        }
588    }
589}
590
591impl Eq for RecordBatch<'_> {}
592
593// ── Encode trait impl ─────────────────────────────────────────────────────────
594
595impl crate::Encode for RecordBatch<'_> {
596    fn encode<B: bytes::BufMut>(
597        &self,
598        buf: &mut B,
599        version: i16,
600    ) -> Result<(), crate::ProtocolError> {
601        let owned = self.to_owned().map_err(crate::ProtocolError::from)?;
602        crate::Encode::encode(&owned, buf, version)
603    }
604
605    fn encoded_len(&self, version: i16) -> usize {
606        match self.to_owned() {
607            Ok(o) => crate::Encode::encoded_len(&o, version),
608            Err(_) => 0,
609        }
610    }
611}
612
613// ── Tests ─────────────────────────────────────────────────────────────────────
614
615#[cfg(test)]
616mod tests {
617
618    use bytes::BytesMut;
619    use krabka_compression::{CompressionError, CompressionType, RecordDecompressionPolicy};
620    use krabka_units::{bytes, fraction};
621
622    use super::*;
623    use crate::DecodeBorrow;
624
625    fn encode_owned_then_borrow(b: &super::super::owned::RecordBatch) -> Vec<u8> {
626        let mut buf = BytesMut::new();
627        b.encode(&mut buf).unwrap();
628        buf.to_vec()
629    }
630
631    fn refresh_batch_crc(encoded: &mut [u8]) {
632        let coverage_start = crate::records::CRC_COVERAGE_START;
633        let crc = crc32c(&encoded[coverage_start..]);
634        encoded[coverage_start - 4..coverage_start].copy_from_slice(&crc.to_be_bytes());
635    }
636
637    #[test]
638    fn borrowed_roundtrip_cases() {
639        for (_case, codec) in [
640            ("none", CompressionType::None),
641            ("gzip", CompressionType::Gzip),
642            ("snappy", CompressionType::Snappy),
643            ("lz4", CompressionType::Lz4),
644            ("zstd", CompressionType::Zstd),
645        ] {
646            let mut owned = super::super::owned::RecordBatch::default();
647            owned.attributes = owned.attributes.with_compression(codec);
648            owned.records.push(super::super::owned::Record {
649                key: Some(Bytes::from_static(b"key")),
650                value: Some(Bytes::from_static(b"value")),
651                ..Default::default()
652            });
653
654            let encoded = encode_owned_then_borrow(&owned);
655            let mut cur: &[u8] = &encoded[..];
656            let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
657            assert2::assert!(cur.is_empty());
658            assert2::assert!(borrowed.attributes() == owned.attributes);
659
660            let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
661            let expected_records = vec![Record {
662                attributes: 0,
663                timestamp_delta: 0,
664                offset_delta: 0,
665                key: Some(b"key".as_slice()),
666                value: Some(b"value".as_slice()),
667                headers: vec![],
668            }];
669            assert2::assert!(records == expected_records);
670
671            let back_owned = borrowed.to_owned().unwrap();
672            assert2::assert!(back_owned == owned);
673        }
674    }
675
676    /// The borrowed record reports the body size the owned record does, for the
677    /// null and the present forms of every field and on both sides of the
678    /// length-prefix width boundary (a body of 63 bytes takes a one-byte
679    /// prefix, one of 64 takes two).
680    #[test]
681    fn borrowed_body_len_matches_the_owned_record() {
682        use super::super::owned::{Record as OwnedRecord, RecordHeader as OwnedHeader};
683
684        let owned = super::super::owned::RecordBatch {
685            records: vec![
686                OwnedRecord::default(),
687                OwnedRecord {
688                    timestamp_delta: 17,
689                    offset_delta: 2,
690                    key: Some(Bytes::from_static(b"the-key")),
691                    value: Some(Bytes::from_static(b"hello kafka")),
692                    headers: vec![
693                        OwnedHeader {
694                            key: "trace-id".to_string(),
695                            value: Some(Bytes::from_static(b"abc")),
696                        },
697                        OwnedHeader {
698                            key: "null-val".to_string(),
699                            value: None,
700                        },
701                    ],
702                    ..Default::default()
703                },
704                OwnedRecord {
705                    value: Some(Bytes::from(vec![7_u8; 57])),
706                    ..Default::default()
707                },
708                OwnedRecord {
709                    value: Some(Bytes::from(vec![7_u8; 58])),
710                    ..Default::default()
711                },
712            ],
713            last_offset_delta: 3,
714            ..Default::default()
715        };
716        let encoded = encode_owned_then_borrow(&owned);
717        let mut cur: &[u8] = &encoded[..];
718        let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
719        let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
720
721        let lens: Vec<usize> = records.iter().map(Record::body_len).collect();
722        assert2::assert!(lens == vec![6, 47, 63, 64]);
723        let owned_lens: Vec<usize> = owned.records.iter().map(OwnedRecord::body_len).collect();
724        assert2::assert!(lens == owned_lens);
725    }
726
727    #[test]
728    fn decompression_policy_limits_borrowed_decode() {
729        let mut owned = super::super::owned::RecordBatch {
730            attributes: Attributes::default().with_compression(CompressionType::Lz4),
731            records: vec![super::super::owned::Record {
732                value: Some(Bytes::from(vec![b'x'; 4096])),
733                ..Default::default()
734            }],
735            ..Default::default()
736        };
737        owned.last_offset_delta = 0;
738        let encoded = encode_owned_then_borrow(&owned);
739
740        let mut default_cur = &encoded[..];
741        RecordBatch::decode_borrow(&mut default_cur, 0).unwrap();
742
743        let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
744        let mut limited_cur = &encoded[..];
745        assert2::assert!(matches!(
746            RecordBatch::decode_borrow_with_policy(&mut limited_cur, policy),
747            Err(RecordsError::Compression(CompressionError::TooLarge {
748                limit: 32
749            }))
750        ));
751    }
752
753    #[test]
754    fn zero_copy_for_uncompressed() {
755        // Pointer-identity: record key/value slices must point into the
756        // input buffer for uncompressed batches.
757        let mut owned = super::super::owned::RecordBatch::default();
758        owned.records.push(super::super::owned::Record {
759            key: Some(Bytes::from_static(b"k")),
760            value: Some(Bytes::from_static(b"v")),
761            ..Default::default()
762        });
763        let encoded = encode_owned_then_borrow(&owned);
764        let encoded_start = encoded.as_ptr() as usize;
765        let encoded_end = encoded_start + encoded.len();
766
767        let mut cur: &[u8] = &encoded[..];
768        let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
769        let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
770
771        let v_ptr = records[0].value.unwrap().as_ptr() as usize;
772        assert2::assert!(v_ptr >= encoded_start && v_ptr < encoded_end);
773    }
774
775    #[test]
776    fn validate_one_v2_batch_reads_header_and_len() {
777        let mut owned = super::super::owned::RecordBatch {
778            base_offset: 7,
779            partition_leader_epoch: 3,
780            last_offset_delta: 0,
781            producer_id: 99,
782            producer_epoch: 1,
783            base_sequence: 5,
784            max_timestamp: 1_234,
785            ..Default::default()
786        };
787        owned.records.push(super::super::owned::Record {
788            value: Some(Bytes::from_static(b"payload")),
789            ..Default::default()
790        });
791        let encoded = encode_owned_then_borrow(&owned);
792
793        let v = validate_one_v2_batch(&encoded).unwrap();
794        assert2::assert!(
795            (
796                v.total_len,
797                v.header.base_offset.get(),
798                v.header.partition_leader_epoch.get(),
799                v.header.producer_id.get(),
800                v.header.producer_epoch.get(),
801                v.header.base_sequence.get(),
802                v.header.max_timestamp.get(),
803                v.header.magic,
804            ) == (encoded.len(), 7, 3, 99, 1, 5, 1_234, 2)
805        );
806        assert2::assert!(
807            v.validate_records(RecordDecompressionPolicy::default())
808                .is_ok()
809        );
810    }
811
812    #[test]
813    fn borrowed_crc_stops_at_the_declared_batch_boundary() {
814        for codec in [CompressionType::None, CompressionType::Lz4] {
815            for value in [Bytes::new(), Bytes::from_static(b"payload")] {
816                let mut owned = super::super::owned::RecordBatch {
817                    records: vec![super::super::owned::Record {
818                        value: Some(value),
819                        ..Default::default()
820                    }],
821                    ..Default::default()
822                };
823                owned.attributes = owned.attributes.with_compression(codec);
824                let mut encoded = encode_owned_then_borrow(&owned);
825                let first_len = encoded.len();
826                encoded.extend_from_within(..);
827                // A corrupt following batch must not affect the first CRC.
828                encoded[first_len + 21] ^= 1;
829
830                let validated = validate_one_v2_batch(&encoded).unwrap();
831                assert2::assert!(validated.total_len == first_len);
832                let mut input = encoded.as_slice();
833                let borrowed = RecordBatch::decode_borrow(&mut input, 0).unwrap();
834                assert2::assert!(borrowed.to_owned().unwrap() == owned);
835                assert2::assert!(input == &encoded[first_len..]);
836                assert2::assert!(matches!(
837                    validate_one_v2_batch(input),
838                    Err(RecordsError::CrcMismatch { .. })
839                ));
840                assert2::assert!(matches!(
841                    RecordBatch::decode_borrow_with_policy(
842                        &mut input,
843                        RecordDecompressionPolicy::default(),
844                    ),
845                    Err(RecordsError::CrcMismatch { .. })
846                ));
847            }
848        }
849    }
850
851    #[test]
852    fn validated_batch_rejects_crc_valid_malformed_record_body() {
853        let owned = super::super::owned::RecordBatch {
854            records: vec![super::super::owned::Record {
855                value: Some(Bytes::from_static(b"value")),
856                ..Default::default()
857            }],
858            ..Default::default()
859        };
860        let mut encoded = encode_owned_then_borrow(&owned);
861        encoded[HEADER_LEN] = 0; // zero-length first record body
862        refresh_batch_crc(&mut encoded);
863
864        let validated = validate_one_v2_batch(&encoded).expect("CRC-valid header");
865        let error = validated
866            .validate_records(RecordDecompressionPolicy::default())
867            .unwrap_err();
868        assert2::assert!(matches!(error, RecordsError::RecordParse(_)));
869    }
870
871    #[test]
872    fn validated_batch_visits_records_during_validation() {
873        let owned = super::super::owned::RecordBatch {
874            attributes: Attributes::default().with_compression(CompressionType::Lz4),
875            records: vec![
876                super::super::owned::Record {
877                    timestamp_delta: 3,
878                    value: Some(Bytes::from_static(b"first")),
879                    ..Default::default()
880                },
881                super::super::owned::Record {
882                    timestamp_delta: 8,
883                    value: Some(Bytes::from_static(b"second")),
884                    ..Default::default()
885                },
886            ],
887            ..Default::default()
888        };
889        let encoded = encode_owned_then_borrow(&owned);
890        let validated = validate_one_v2_batch(&encoded).unwrap();
891        let mut timestamps = Vec::new();
892
893        validated
894            .validate_records_with(RecordDecompressionPolicy::default(), |record| {
895                timestamps.push(record.timestamp_delta);
896            })
897            .unwrap();
898
899        assert2::assert!(timestamps == vec![3, 8]);
900    }
901
902    #[test]
903    fn validated_batch_rejects_trailing_body_bytes() {
904        let owned = super::super::owned::RecordBatch {
905            records: vec![super::super::owned::Record::default()],
906            ..Default::default()
907        };
908        let mut encoded = encode_owned_then_borrow(&owned);
909        let batch_length = i32::from_be_bytes(encoded[8..12].try_into().unwrap());
910        encoded[8..12].copy_from_slice(&(batch_length + 1).to_be_bytes());
911        encoded.push(0);
912        refresh_batch_crc(&mut encoded);
913
914        let validated = validate_one_v2_batch(&encoded).expect("CRC-valid header");
915        let error = validated
916            .validate_records(RecordDecompressionPolicy::default())
917            .unwrap_err();
918        assert2::assert!(
919            matches!(error, RecordsError::RecordParse(message) if message.contains("trailing bytes"))
920        );
921    }
922
923    #[test]
924    fn validated_batch_applies_decompression_policy() {
925        let owned = super::super::owned::RecordBatch {
926            attributes: Attributes::default().with_compression(CompressionType::Lz4),
927            records: vec![super::super::owned::Record {
928                value: Some(Bytes::from(vec![b'x'; 4096])),
929                ..Default::default()
930            }],
931            ..Default::default()
932        };
933        let encoded = encode_owned_then_borrow(&owned);
934        let validated = validate_one_v2_batch(&encoded).unwrap();
935        let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
936
937        assert2::assert!(matches!(
938            validated.validate_records(policy),
939            Err(RecordsError::Compression(CompressionError::TooLarge {
940                limit: 32
941            }))
942        ));
943    }
944
945    #[test]
946    fn validated_batch_distinguishes_zero_and_negative_record_counts() {
947        let mut encoded = encode_owned_then_borrow(&super::super::owned::RecordBatch::default());
948        let zero = validate_one_v2_batch(&encoded).unwrap();
949        assert2::assert!(
950            zero.validate_records(RecordDecompressionPolicy::default())
951                .is_ok()
952        );
953
954        encoded[57..61].copy_from_slice(&(-1i32).to_be_bytes());
955        refresh_batch_crc(&mut encoded);
956        let negative = validate_one_v2_batch(&encoded).unwrap();
957        assert2::assert!(matches!(
958            negative.validate_records(RecordDecompressionPolicy::default()),
959            Err(RecordsError::RecordParse(message)) if message.contains("negative records_count")
960        ));
961    }
962
963    #[test]
964    fn validate_one_v2_batch_rejects_corrupt_crc() {
965        let owned = super::super::owned::RecordBatch {
966            records: vec![super::super::owned::Record {
967                value: Some(Bytes::from_static(b"x")),
968                ..Default::default()
969            }],
970            ..Default::default()
971        };
972        let mut encoded = encode_owned_then_borrow(&owned);
973        // Flip a body byte (after the 61-byte header) → CRC mismatch.
974        encoded[HEADER_LEN] ^= 0xFF;
975        let err = validate_one_v2_batch(&encoded).unwrap_err();
976        assert2::assert!(matches!(err, RecordsError::CrcMismatch { .. }));
977    }
978
979    #[test]
980    fn validate_one_v2_batch_rejects_truncated() {
981        let owned = super::super::owned::RecordBatch {
982            records: vec![super::super::owned::Record {
983                value: Some(Bytes::from_static(b"value")),
984                ..Default::default()
985            }],
986            ..Default::default()
987        };
988        let encoded = encode_owned_then_borrow(&owned);
989        let err = validate_one_v2_batch(&encoded[..encoded.len() - 2]).unwrap_err();
990        assert2::assert!(matches!(err, RecordsError::BodyTooShort { .. }));
991    }
992
993    #[test]
994    fn count_records_in_v2_batches_counts_single_and_concatenated_batches() {
995        let one = encode_owned_then_borrow(&super::super::owned::RecordBatch {
996            records: vec![super::super::owned::Record::default()],
997            ..Default::default()
998        });
999        let three = encode_owned_then_borrow(&super::super::owned::RecordBatch {
1000            last_offset_delta: 2,
1001            records: vec![
1002                super::super::owned::Record::default(),
1003                super::super::owned::Record::default(),
1004                super::super::owned::Record::default(),
1005            ],
1006            ..Default::default()
1007        });
1008        let mut both = Vec::with_capacity(one.len() + three.len());
1009        both.extend_from_slice(&one);
1010        both.extend_from_slice(&three);
1011
1012        for (_case, input, expected) in [("single", &one[..], 1), ("concatenated", &both[..], 4)] {
1013            assert2::assert!(count_records_in_v2_batches(input) == expected);
1014        }
1015    }
1016
1017    #[test]
1018    fn count_records_in_v2_batches_stops_at_bad_or_non_v2_input() {
1019        let encoded = encode_owned_then_borrow(&super::super::owned::RecordBatch {
1020            records: vec![super::super::owned::Record::default()],
1021            ..Default::default()
1022        });
1023        let mut with_truncated_tail = encoded.clone();
1024        with_truncated_tail.extend_from_slice(&encoded[..HEADER_LEN - 1]);
1025
1026        for (_case, input, expected) in [
1027            ("empty", &[][..], 0),
1028            ("non-v2", &[0u8; HEADER_LEN][..], 0),
1029            ("truncated tail", &with_truncated_tail[..], 1),
1030        ] {
1031            assert2::assert!(count_records_in_v2_batches(input) == expected);
1032        }
1033    }
1034
1035    /// Builds a bare 61-byte v2 batch header with `magic == 2` and all other
1036    /// fields zero, plus the given big-endian `batch_length` and
1037    /// `records_count`. The helper drives the `total_len` boundary checks of
1038    /// `count_records_in_v2_batches` precisely, because `records_count` is read
1039    /// straight from the header and the body needs no real records.
1040    fn v2_header_only(batch_length: i32, records_count: i32) -> Vec<u8> {
1041        let mut buf = vec![0u8; HEADER_LEN];
1042        buf[16] = 2; // magic == v2
1043        buf[8..12].copy_from_slice(&batch_length.to_be_bytes());
1044        buf[57..61].copy_from_slice(&records_count.to_be_bytes());
1045        buf
1046    }
1047
1048    #[test]
1049    fn count_records_in_v2_batches_accepts_batch_that_exactly_fills_buffer() {
1050        // total_len = 12 + 49 = 61 = HEADER_LEN, exactly filling the slice, so
1051        // the batch is accepted and its header count is summed. Guards the
1052        // `total_len < HEADER_LEN` lower bound: a `<`→`==`/`<=` mutation would
1053        // wrongly break here and report 0.
1054        let buf = v2_header_only(49, 7);
1055        assert2::assert!(buf.len() == HEADER_LEN);
1056        assert2::assert!(count_records_in_v2_batches(&buf) == 7);
1057    }
1058
1059    #[test]
1060    fn count_records_in_v2_batches_rejects_batch_overrunning_buffer() {
1061        // total_len = 12 + 100 = 112 > 61 bytes available, so the batch is
1062        // rejected and contributes nothing. Guards the
1063        // `total_len > remaining.len()` arm: an `||`→`&&` mutation would skip
1064        // the break and index past the buffer.
1065        let buf = v2_header_only(100, 9);
1066        assert2::assert!(count_records_in_v2_batches(&buf) == 0);
1067    }
1068
1069    #[test]
1070    fn borrowed_encode_via_trait_roundtrips() {
1071        use crate::Encode as _;
1072        let owned_in = super::super::owned::RecordBatch {
1073            records: vec![super::super::owned::Record {
1074                key: Some(Bytes::from_static(b"x")),
1075                value: Some(Bytes::from_static(b"y")),
1076                ..Default::default()
1077            }],
1078            ..Default::default()
1079        };
1080        let bytes_in = encode_owned_then_borrow(&owned_in);
1081        let mut cur: &[u8] = &bytes_in[..];
1082        let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
1083
1084        let mut out = BytesMut::new();
1085        borrowed.encode(&mut out, 0).unwrap();
1086        assert2::assert!(&out[..] == &bytes_in[..]);
1087    }
1088}