krabka_protocol/records/
payload.rs

1//! `RecordsPayload`: the wire-field type for any Kafka message whose schema
2//! declares a `records` field (`Fetch`, `Produce`, `FetchSnapshot`, `ShareFetch`).
3//!
4//! At the protocol layer the records field is opaque bytes. The contents are
5//! either a v2 `RecordBatch` (current) or a v0/v1 `MessageSet` (legacy, used
6//! by old clients on down-conversion). The discriminator is the magic byte at
7//! offset 16 of the first batch. Both formats put it at that offset by
8//! coincidence of layout: v2 is `base_offset+batch_length+leader_epoch+magic`,
9//! and legacy is `offset+message_size+crc+magic`.
10//!
11//! This type parses the v2 form eagerly, which keeps the broker code paths
12//! unchanged: a caller uses `RecordsPayload` methods such as `encoded_len` in
13//! place of the `RecordBatch` methods. Legacy payloads stay as raw [`Bytes`]
14//! and round-trip verbatim. The
15//! [`krabka-records-legacy`](../../records_legacy/index.html) crate supplies
16//! the codec when an old client appears on the wire.
17
18#[cfg(any(
19    target_os = "linux",
20    target_os = "macos",
21    target_os = "ios",
22    target_os = "tvos",
23    target_os = "watchos",
24    target_os = "freebsd",
25    target_os = "dragonfly",
26))]
27use std::io;
28use std::mem;
29
30use bytes::{Buf, BufMut, Bytes};
31use krabka_compression::{CompressionType, RecordDecompressionPolicy};
32
33#[cfg(any(
34    target_os = "linux",
35    target_os = "macos",
36    target_os = "ios",
37    target_os = "tvos",
38    target_os = "watchos",
39    target_os = "freebsd",
40    target_os = "dragonfly",
41))]
42use crate::records::FileRegion;
43use crate::{
44    primitives::{fixed::put_i32, varint::put_uvarint},
45    records::{RecordsError, borrowed::RecordBatch as RecordBatchBorrowed, owned::RecordBatch},
46};
47
48/// Owned form of a records-field payload.
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum RecordsPayload {
51    /// Zero or more parsed v2 batches. The records field is a sequence.
52    V2(Vec<RecordBatch>),
53    /// Verbatim wire-format v2 bytes, one or more batches, forwarded without
54    /// parsing. The fetch pass-through path produces this variant.
55    Raw(Bytes),
56    /// Opaque pre-v2 bytes, a v0 or v1 `MessageSet`. Decode them with
57    /// `krabka_records_legacy::decode_message_set`.
58    Legacy(Bytes),
59    /// Zero-copy fetch (Increments D + E): the records run stays in segment
60    /// `.log` files. The broker `sendfile(2)`s it straight to a plaintext
61    /// socket and never materializes it in userspace. There is one
62    /// [`crate::records::FileRegion`] per contributing segment. `encode_to`
63    /// falls back to `pread` + `put_slice` on TLS and non-sendfile platforms,
64    /// and for `encoded_len` agreement. This variant is gated on the SENDFILE
65    /// alias (Linux + Apple + FreeBSD/DragonFly) because it only exists to feed
66    /// the `sendfile` drainer. On Windows the fallback `pread` path runs.
67    #[cfg(any(
68        target_os = "linux",
69        target_os = "macos",
70        target_os = "ios",
71        target_os = "tvos",
72        target_os = "watchos",
73        target_os = "freebsd",
74        target_os = "dragonfly",
75    ))]
76    FileRegions(Vec<FileRegion>),
77}
78
79impl RecordsPayload {
80    /// Writes this payload with its Kafka `BYTES` or `COMPACT_BYTES` length
81    /// prefix.
82    ///
83    /// Payloads whose encoded size is exact are written directly to `buf`.
84    /// Compressed v2 batches are first encoded into a temporary buffer because
85    /// their compressed wire size cannot be predicted from
86    /// [`RecordBatch::encoded_len`].
87    ///
88    /// # Errors
89    ///
90    /// Returns an error when the payload cannot be encoded or its length does
91    /// not fit the selected Kafka length prefix.
92    pub fn encode_length_prefixed<B: BufMut>(
93        &self,
94        buf: &mut B,
95        version: i16,
96        flexible: bool,
97    ) -> Result<(), crate::ProtocolError> {
98        let has_compressed_batch = matches!(
99            self,
100            Self::V2(batches)
101                if batches
102                    .iter()
103                    .any(|batch| batch.attributes.compression() != CompressionType::None)
104        );
105
106        if has_compressed_batch {
107            let mut encoded = bytes::BytesMut::new();
108            <Self as crate::Encode>::encode(self, &mut encoded, version)?;
109            put_payload_length(buf, encoded.len(), flexible)?;
110            buf.put_slice(&encoded);
111        } else {
112            put_payload_length(buf, self.payload_len(), flexible)?;
113            <Self as crate::Encode>::encode(self, buf, version)?;
114        }
115        Ok(())
116    }
117
118    /// Constructs a payload from raw records-field bytes.
119    ///
120    /// When the bytes look like v2, this function decodes every batch in the
121    /// field. If they do not, it keeps the bytes as opaque legacy bytes.
122    /// # Errors
123    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
124    pub fn from_bytes(bytes: Bytes) -> Result<Self, RecordsError> {
125        Self::from_bytes_with_policy(bytes, RecordDecompressionPolicy::default())
126    }
127
128    /// Constructs a payload from raw records bytes with explicit decompression
129    /// limits.
130    ///
131    /// # Errors
132    ///
133    /// Returns the underlying records error for malformed, truncated, corrupt,
134    /// or over-limit v2 input.
135    pub fn from_bytes_with_policy(
136        bytes: Bytes,
137        policy: RecordDecompressionPolicy,
138    ) -> Result<Self, RecordsError> {
139        if looks_like_v2(&bytes) {
140            let mut cur: &[u8] = &bytes;
141            let mut batches = Vec::new();
142            while !cur.is_empty() {
143                batches.push(RecordBatch::decode_with_policy(&mut cur, policy)?);
144            }
145            Ok(Self::V2(batches))
146        } else {
147            Ok(Self::Legacy(bytes))
148        }
149    }
150
151    /// Wire size of the records-field bytes (no outer length prefix).
152    #[must_use]
153    pub fn payload_len(&self) -> usize {
154        match self {
155            Self::V2(batches) => batches.iter().map(RecordBatch::encoded_batch_len).sum(),
156            Self::Raw(b) | Self::Legacy(b) => b.len(),
157            #[cfg(any(
158                target_os = "linux",
159                target_os = "macos",
160                target_os = "ios",
161                target_os = "tvos",
162                target_os = "watchos",
163                target_os = "freebsd",
164                target_os = "dragonfly",
165            ))]
166            Self::FileRegions(regions) => regions.iter().map(|r| r.len).sum(),
167        }
168    }
169
170    /// Writes the payload bytes into `buf`.
171    ///
172    /// The caller owns the outer framing.
173    ///
174    /// For `FileRegions` this is the fallback path, used on TLS, on
175    /// non-Linux platforms, and for `encoded_len` agreement. It `pread`s
176    /// each region out of the segment file and copies it into `buf`. The
177    /// broker's zero-copy `sendfile` path never calls this method. That path
178    /// consumes the `FileRegion`s directly.
179    /// # Errors
180    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
181    pub fn encode_to<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
182        match self {
183            Self::V2(batches) => {
184                for b in batches {
185                    b.encode_batch(buf)?;
186                }
187                Ok(())
188            }
189            Self::Raw(b) | Self::Legacy(b) => {
190                buf.put_slice(b);
191                Ok(())
192            }
193            #[cfg(any(
194                target_os = "linux",
195                target_os = "macos",
196                target_os = "ios",
197                target_os = "tvos",
198                target_os = "watchos",
199                target_os = "freebsd",
200                target_os = "dragonfly",
201            ))]
202            Self::FileRegions(regions) => {
203                use std::os::unix::fs::FileExt;
204                let mut scratch = vec![0u8; 0];
205                for region in regions {
206                    scratch.resize(region.len, 0);
207                    let mut filled = 0usize;
208                    let mut offset = region.offset;
209                    while filled < region.len {
210                        let Some(unfilled) = scratch.get_mut(filled..) else {
211                            return Err(RecordsError::RecordParse(
212                                "FileRegion read overran its buffer".into(),
213                            ));
214                        };
215                        match region.file.read_at(unfilled, offset) {
216                            Ok(0) => {
217                                return Err(RecordsError::RecordParse(
218                                    "FileRegion read hit EOF before len bytes".into(),
219                                ));
220                            }
221                            Ok(n) => {
222                                filled += n;
223                                offset += n as u64;
224                            }
225                            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
226                            Err(e) => {
227                                return Err(RecordsError::RecordParse(format!(
228                                    "FileRegion read error: {e}"
229                                )));
230                            }
231                        }
232                    }
233                    buf.put_slice(&scratch);
234                }
235                Ok(())
236            }
237        }
238    }
239
240    /// Borrows the parsed v2 batches, if this is a parsed `V2` payload.
241    ///
242    /// Returns `None` for `Raw`, which is intentionally unparsed, for
243    /// `Legacy`, and for `FileRegions`, which is deliberately never
244    /// materialized.
245    #[must_use]
246    pub fn as_v2(&self) -> Option<&[RecordBatch]> {
247        match self {
248            Self::V2(batches) => Some(batches),
249            #[cfg(any(
250                target_os = "linux",
251                target_os = "macos",
252                target_os = "ios",
253                target_os = "tvos",
254                target_os = "watchos",
255                target_os = "freebsd",
256                target_os = "dragonfly",
257            ))]
258            Self::FileRegions(_) => None,
259            Self::Raw(_) | Self::Legacy(_) => None,
260        }
261    }
262
263    /// Borrows the raw legacy bytes, if this payload is `Legacy`.
264    #[must_use]
265    pub fn as_legacy(&self) -> Option<&Bytes> {
266        match self {
267            Self::Legacy(b) => Some(b),
268            #[cfg(any(
269                target_os = "linux",
270                target_os = "macos",
271                target_os = "ios",
272                target_os = "tvos",
273                target_os = "watchos",
274                target_os = "freebsd",
275                target_os = "dragonfly",
276            ))]
277            Self::FileRegions(_) => None,
278            Self::V2(_) | Self::Raw(_) => None,
279        }
280    }
281
282    /// Decodes a response-side records field and tolerates a truncated
283    /// trailing batch.
284    ///
285    /// Kafka returns a partial final `RecordBatch` when a partition hits its
286    /// fetch byte budget mid-batch. The JVM consumer stops at the first
287    /// incomplete batch and re-fetches it from the next offset. This function
288    /// does the same: it decodes every complete batch, then stops at the first
289    /// `HeaderTooShort` or `BodyTooShort` and drops the remainder. A complete
290    /// batch that is corrupt in its CRC, magic byte, or content still returns
291    /// an error, because the leniency forgives truncation only.
292    /// Produce-request validation keeps the strict
293    /// [`from_bytes`](Self::from_bytes).
294    ///
295    /// This function treats only `HeaderTooShort` and `BodyTooShort` as
296    /// truncation. An invalid `batch_length` gives `RecordParse`, which is
297    /// corruption and still returns an error. Kafka truncation always keeps a
298    /// valid `batch_length` prefix, so truncation can appear only as the two
299    /// too-short variants.
300    /// # Errors
301    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
302    pub fn from_fetch_bytes(bytes: Bytes) -> Result<Self, RecordsError> {
303        if !looks_like_v2(&bytes) {
304            return Ok(Self::Legacy(bytes));
305        }
306        let mut cur: &[u8] = &bytes;
307        let mut batches = Vec::new();
308        while !cur.is_empty() {
309            match RecordBatch::decode_batch(&mut cur) {
310                Ok(rb) => batches.push(rb),
311                Err(RecordsError::HeaderTooShort { .. } | RecordsError::BodyTooShort { .. }) => {
312                    break;
313                }
314                Err(e) => return Err(e),
315            }
316        }
317        Ok(Self::V2(batches))
318    }
319
320    /// Lenient `Decode`-shaped entry point for records fields in response
321    /// messages.
322    ///
323    /// It consumes the whole sliced field buffer, which the caller has
324    /// already framed, and parses the bytes
325    /// leniently through [`from_fetch_bytes`](Self::from_fetch_bytes).
326    /// # Errors
327    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
328    pub fn decode_lenient<B: Buf>(
329        buf: &mut B,
330        _version: i16,
331    ) -> Result<Self, crate::ProtocolError> {
332        let bytes = buf.copy_to_bytes(buf.remaining());
333        Self::from_fetch_bytes(bytes).map_err(Into::into)
334    }
335}
336
337fn put_payload_length<B: BufMut>(
338    buf: &mut B,
339    payload_len: usize,
340    flexible: bool,
341) -> Result<(), crate::ProtocolError> {
342    if flexible {
343        let prefixed_len = payload_len
344            .checked_add(1)
345            .and_then(|len| u32::try_from(len).ok())
346            .ok_or(crate::ProtocolError::InvalidValue(
347                "records payload length exceeds compact bytes range",
348            ))?;
349        put_uvarint(buf, prefixed_len);
350    } else {
351        let payload_len = i32::try_from(payload_len).map_err(|_| {
352            crate::ProtocolError::InvalidValue("records payload length exceeds bytes range")
353        })?;
354        put_i32(buf, payload_len);
355    }
356    Ok(())
357}
358
359impl From<RecordBatch> for RecordsPayload {
360    fn from(rb: RecordBatch) -> Self {
361        Self::V2(vec![rb])
362    }
363}
364
365impl From<Vec<RecordBatch>> for RecordsPayload {
366    fn from(v: Vec<RecordBatch>) -> Self {
367        Self::V2(v)
368    }
369}
370
371impl Default for RecordsPayload {
372    fn default() -> Self {
373        Self::V2(Vec::new())
374    }
375}
376
377impl crate::Encode for RecordsPayload {
378    fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
379        self.encode_to(buf).map_err(Into::into)
380    }
381
382    fn encoded_len(&self, _version: i16) -> usize {
383        self.payload_len()
384    }
385}
386
387impl crate::Decode<'_> for RecordsPayload {
388    fn decode<B: Buf>(buf: &mut B, _version: i16) -> Result<Self, crate::ProtocolError> {
389        // The caller (generated codec) has already sliced the buffer to
390        // exactly the records-field bytes via `get_(nullable_)bytes_owned`,
391        // so we consume everything.
392        let bytes = buf.copy_to_bytes(buf.remaining());
393        Self::from_bytes(bytes).map_err(Into::into)
394    }
395}
396
397/// Borrowed form: zero-copy view into the input buffer.
398#[derive(Debug, Clone, PartialEq, Eq)]
399pub enum RecordsPayloadBorrowed<'a> {
400    V2(Vec<RecordBatchBorrowed<'a>>),
401    Legacy(&'a [u8]),
402}
403
404impl<'a> RecordsPayloadBorrowed<'a> {
405    /// # Errors
406    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
407    pub fn from_slice(bytes: &'a [u8]) -> Result<Self, RecordsError> {
408        if looks_like_v2(bytes) {
409            let mut cur: &'a [u8] = bytes;
410            let mut batches = Vec::new();
411            while !cur.is_empty() {
412                let rb = <RecordBatchBorrowed<'a> as crate::DecodeBorrow<'a>>::decode_borrow(
413                    &mut cur, 0,
414                )
415                .map_err(|e| RecordsError::RecordParse(format!("borrowed v2 decode: {e}")))?;
416                batches.push(rb);
417            }
418            Ok(Self::V2(batches))
419        } else {
420            Ok(Self::Legacy(bytes))
421        }
422    }
423
424    #[must_use]
425    pub fn payload_len(&self) -> usize {
426        match self {
427            Self::V2(batches) => batches
428                .iter()
429                .map(|rb| crate::Encode::encoded_len(rb, 0))
430                .sum(),
431            Self::Legacy(b) => b.len(),
432        }
433    }
434
435    /// # Errors
436    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
437    pub fn encode_to<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
438        match self {
439            Self::V2(batches) => {
440                for rb in batches {
441                    crate::Encode::encode(rb, buf, 0).map_err(|e| {
442                        RecordsError::RecordParse(format!("borrowed v2 encode: {e}"))
443                    })?;
444                }
445                Ok(())
446            }
447            Self::Legacy(b) => {
448                buf.put_slice(b);
449                Ok(())
450            }
451        }
452    }
453
454    /// Converts to the owned form and copies buffers where necessary.
455    /// # Errors
456    /// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
457    pub fn to_owned(&self) -> Result<RecordsPayload, RecordsError> {
458        match self {
459            Self::V2(batches) => {
460                let mut owned = Vec::with_capacity(batches.len());
461                for rb in batches {
462                    owned.push(rb.to_owned()?);
463                }
464                Ok(RecordsPayload::V2(owned))
465            }
466            Self::Legacy(b) => Ok(RecordsPayload::Legacy(Bytes::copy_from_slice(b))),
467        }
468    }
469}
470
471impl Default for RecordsPayloadBorrowed<'_> {
472    fn default() -> Self {
473        Self::V2(Vec::new())
474    }
475}
476
477impl crate::Encode for RecordsPayloadBorrowed<'_> {
478    fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
479        self.encode_to(buf).map_err(Into::into)
480    }
481
482    fn encoded_len(&self, _version: i16) -> usize {
483        self.payload_len()
484    }
485}
486
487impl<'de> crate::DecodeBorrow<'de> for RecordsPayloadBorrowed<'de> {
488    fn decode_borrow(buf: &mut &'de [u8], _version: i16) -> Result<Self, crate::ProtocolError> {
489        // Caller has sliced the buffer to exactly the records-field bytes;
490        // consume everything by swapping in an empty tail.
491        let bytes = mem::take(buf);
492        Self::from_slice(bytes).map_err(Into::into)
493    }
494}
495
496/// True when `bytes` look like a v2 record batch.
497///
498/// A v2 batch has magic byte 2 at the well-known offset. Legacy v0 and v1
499/// `MessageSets` carry magic 0 or 1 at the same offset, so this check separates
500/// the two.
501///
502/// The threshold is `MAGIC_OFFSET + 1 = 17`, not the full `HEADER_LEN`. A
503/// truncated v2 batch must still land in the V2 arm, so that the downstream
504/// `RecordBatch::decode_batch` can give a precise error instead of a silent
505/// misclassification as legacy.
506#[inline]
507fn looks_like_v2(bytes: &[u8]) -> bool {
508    // The magic byte sits at `base_offset(8) + batch_length(4) +
509    // partition_leader_epoch(4) = 16` in v2. Legacy MessageSets place
510    // the first message's magic byte at `offset(8) + message_size(4) +
511    // crc(4) = 16` — same index, different meaning.
512    const MAGIC_OFFSET: usize = 16;
513    bytes.get(MAGIC_OFFSET) == Some(&2)
514}
515
516#[cfg(test)]
517mod tests {
518
519    use bytes::BytesMut;
520    use krabka_compression::{CompressionError, CompressionType, RecordDecompressionPolicy};
521    use krabka_units::{bytes, fraction};
522
523    use super::*;
524    use crate::{
525        primitives::string_bytes::{compact_bytes_len_from_size, put_bytes, put_compact_bytes},
526        records::{Record, RecordBatch},
527    };
528
529    fn sample_v2() -> RecordBatch {
530        RecordBatch {
531            base_offset: 42,
532            records: vec![Record {
533                key: Some(Bytes::from_static(b"k")),
534                value: Some(Bytes::from_static(b"v")),
535                ..Default::default()
536            }],
537            ..RecordBatch::default()
538        }
539    }
540
541    #[test]
542    fn from_bytes_dispatches_v2() {
543        let rb = sample_v2();
544        let mut buf = BytesMut::new();
545        rb.encode_batch(&mut buf).unwrap();
546        let p = RecordsPayload::from_bytes(buf.freeze()).unwrap();
547        match p {
548            RecordsPayload::V2(batches) => assert2::assert!(batches == vec![rb]),
549            _ => panic!("expected V2"),
550        }
551    }
552
553    #[test]
554    fn decompression_policy_limits_payload_decode() {
555        let mut rb = sample_v2();
556        rb.records[0].value = Some(Bytes::from(vec![b'x'; 4096]));
557        rb.attributes = rb.attributes.with_compression(CompressionType::Lz4);
558        let mut wire = BytesMut::new();
559        rb.encode_batch(&mut wire).unwrap();
560        let wire = wire.freeze();
561
562        RecordsPayload::from_bytes(wire.clone()).unwrap();
563
564        let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
565        assert2::assert!(matches!(
566            RecordsPayload::from_bytes_with_policy(wire, policy),
567            Err(RecordsError::Compression(CompressionError::TooLarge {
568                limit: 32
569            }))
570        ));
571    }
572
573    #[test]
574    fn from_bytes_parses_all_batches() {
575        // Two v2 batches concatenated must both decode.
576        let mut b0 = sample_v2();
577        b0.base_offset = 0;
578        let mut b1 = sample_v2();
579        b1.base_offset = 1;
580        let mut buf = BytesMut::new();
581        b0.encode_batch(&mut buf).unwrap();
582        b1.encode_batch(&mut buf).unwrap();
583        let p = RecordsPayload::from_bytes(buf.freeze()).unwrap();
584        let batches = p.as_v2().expect("v2");
585        assert2::assert!(batches == &[b0, b1][..]);
586    }
587
588    #[test]
589    fn raw_passthrough_roundtrips() {
590        let mut b = sample_v2();
591        b.base_offset = 7;
592        let mut wire = BytesMut::new();
593        b.encode_batch(&mut wire).unwrap();
594        let wire = wire.freeze();
595        let p = RecordsPayload::Raw(wire.clone());
596        assert2::assert!(p.payload_len() == wire.len());
597        let mut out = BytesMut::new();
598        p.encode_to(&mut out).unwrap();
599        assert2::assert!(&out[..] == &wire[..]); // verbatim
600        assert2::assert!(p.as_v2().is_none()); // Raw is unparsed
601    }
602
603    #[test]
604    fn length_prefixed_raw_encodes_into_fixed_buffer() {
605        let payload = Bytes::from(vec![0x5a; 1024 * 1024]);
606        let records = RecordsPayload::Raw(payload.clone());
607        let mut output = vec![0_u8; 4 + payload.len()];
608        let mut destination = output.as_mut_slice();
609
610        records
611            .encode_length_prefixed(&mut destination, 0, false)
612            .unwrap();
613
614        assert2::assert!(destination.is_empty());
615        let payload_len = i32::try_from(payload.len()).unwrap();
616        assert2::assert!(&output[..4] == &payload_len.to_be_bytes());
617        assert2::assert!(&output[4..] == payload.as_ref());
618    }
619
620    #[test]
621    fn length_prefixed_legacy_encodes_compact_bytes_directly() {
622        let payload = legacy_bytes();
623        let records = RecordsPayload::Legacy(payload.clone());
624        let expected_len = compact_bytes_len_from_size(payload.len());
625        let mut output = vec![0_u8; expected_len];
626        let mut destination = output.as_mut_slice();
627
628        records
629            .encode_length_prefixed(&mut destination, 0, true)
630            .unwrap();
631
632        assert2::assert!(destination.is_empty());
633        let mut expected = BytesMut::new();
634        put_compact_bytes(&mut expected, &payload).unwrap();
635        assert2::assert!(output.as_slice() == expected.as_ref());
636    }
637
638    #[test]
639    fn length_prefixed_compressed_v2_uses_actual_wire_size() {
640        let mut batch = sample_v2();
641        batch.records[0].value = Some(Bytes::from(vec![b'x'; 4096]));
642        batch.attributes = batch.attributes.with_compression(CompressionType::Lz4);
643        let records = RecordsPayload::V2(vec![batch]);
644
645        let mut payload_wire = BytesMut::new();
646        <RecordsPayload as crate::Encode>::encode(&records, &mut payload_wire, 0).unwrap();
647        let mut expected = BytesMut::new();
648        put_bytes(&mut expected, &payload_wire).unwrap();
649        let mut actual = BytesMut::new();
650        records
651            .encode_length_prefixed(&mut actual, 0, false)
652            .unwrap();
653
654        assert2::assert!(actual == expected);
655        let prefix = i32::from_be_bytes(actual[..4].try_into().unwrap());
656        assert2::assert!(usize::try_from(prefix).unwrap() == payload_wire.len());
657        assert2::assert!(payload_wire.len() != records.payload_len());
658    }
659
660    #[test]
661    fn from_bytes_dispatches_legacy() {
662        // Build a fake legacy MessageSet entry: offset(8) + size(4) + crc(4) +
663        // magic(1)=1 + … . Just need byte 16 to be 1.
664        let mut buf = vec![0u8; 17];
665        buf[16] = 1;
666        let p = RecordsPayload::from_bytes(Bytes::from(buf.clone())).unwrap();
667        match p {
668            RecordsPayload::Legacy(b) => assert2::assert!(&b[..] == &buf[..]),
669            _ => panic!("expected Legacy"),
670        }
671    }
672
673    #[test]
674    fn roundtrip_v2() {
675        let p: RecordsPayload = sample_v2().into();
676        let mut buf = BytesMut::new();
677        p.encode_to(&mut buf).unwrap();
678        let back = RecordsPayload::from_bytes(buf.freeze()).unwrap();
679        assert2::assert!(p == back);
680        assert2::assert!(p.payload_len() == back.payload_len());
681    }
682
683    #[test]
684    fn encode_decode_via_traits() {
685        let p: RecordsPayload = sample_v2().into();
686        let mut buf = BytesMut::new();
687        <RecordsPayload as crate::Encode>::encode(&p, &mut buf, 0).unwrap();
688        let mut cur: &[u8] = &buf;
689        let back = <RecordsPayload as crate::Decode>::decode(&mut cur, 0).unwrap();
690        assert2::assert!(p == back);
691    }
692
693    #[test]
694    fn borrowed_dispatches() {
695        let rb = sample_v2();
696        let mut buf = BytesMut::new();
697        rb.encode_batch(&mut buf).unwrap();
698        let frozen = buf.freeze();
699        let p = RecordsPayloadBorrowed::from_slice(&frozen).unwrap();
700        assert2::assert!(matches!(p, RecordsPayloadBorrowed::V2(_)));
701        let owned = p.to_owned().unwrap();
702        match owned {
703            RecordsPayload::V2(batches) => assert2::assert!(batches[0].base_offset == 42),
704            _ => panic!("expected V2"),
705        }
706    }
707
708    #[test]
709    fn from_record_batch() {
710        let rb = sample_v2();
711        let p: RecordsPayload = rb.clone().into();
712        assert2::assert!(p.as_v2() == Some(&[rb][..]));
713        assert2::assert!(p.as_legacy().is_none());
714    }
715
716    fn legacy_bytes() -> Bytes {
717        let mut buf = vec![0u8; 24];
718        buf[16] = 1;
719        for (i, b) in (b'a'..=b'h').enumerate() {
720            buf[17 + i % 7] = b;
721        }
722        Bytes::from(buf)
723    }
724
725    #[test]
726    fn legacy_payload_len_and_encode_owned() {
727        let bytes = legacy_bytes();
728        let p = RecordsPayload::from_bytes(bytes.clone()).unwrap();
729        assert2::assert!(p == RecordsPayload::Legacy(bytes.clone()));
730        assert2::assert!(p.payload_len() == bytes.len());
731
732        let mut out = BytesMut::new();
733        p.encode_to(&mut out).unwrap();
734        assert2::assert!(&out[..] == &bytes[..]);
735    }
736
737    #[test]
738    fn legacy_roundtrip_via_traits() {
739        let bytes = legacy_bytes();
740        let p = RecordsPayload::from_bytes(bytes.clone()).unwrap();
741        let mut buf = BytesMut::new();
742        <RecordsPayload as crate::Encode>::encode(&p, &mut buf, 0).unwrap();
743        assert2::assert!(<RecordsPayload as crate::Encode>::encoded_len(&p, 0) == bytes.len());
744        let mut cur: &[u8] = &buf;
745        let back = <RecordsPayload as crate::Decode>::decode(&mut cur, 0).unwrap();
746        assert2::assert!(matches!(back, RecordsPayload::Legacy(_)));
747        assert2::assert!(back.as_legacy().unwrap() == &bytes);
748    }
749
750    #[test]
751    fn owned_default_is_empty_v2() {
752        let p = RecordsPayload::default();
753        assert2::assert!(matches!(p, RecordsPayload::V2(ref v) if v.is_empty()));
754    }
755
756    #[test]
757    fn looks_like_v2_rejects_short_buffer() {
758        // Too short to peek the magic byte at offset 16; must fall through to Legacy.
759        let short = Bytes::from_static(&[0u8; 10]);
760        let p = RecordsPayload::from_bytes(short.clone()).unwrap();
761        assert2::assert!(p.as_legacy() == Some(&short));
762    }
763
764    #[test]
765    fn borrowed_legacy_roundtrip() {
766        let bytes = legacy_bytes();
767        let p = RecordsPayloadBorrowed::from_slice(&bytes).unwrap();
768        assert2::assert!(matches!(p, RecordsPayloadBorrowed::Legacy(_)));
769        assert2::assert!(p.payload_len() == bytes.len());
770
771        let mut out = BytesMut::new();
772        p.encode_to(&mut out).unwrap();
773        assert2::assert!(&out[..] == &bytes[..]);
774
775        let owned = p.to_owned().unwrap();
776        match owned {
777            RecordsPayload::Legacy(b) => assert2::assert!(&b[..] == &bytes[..]),
778            _ => panic!("expected Legacy"),
779        }
780    }
781
782    #[test]
783    fn borrowed_v2_payload_len_and_encode() {
784        let rb = sample_v2();
785        let mut buf = BytesMut::new();
786        rb.encode_batch(&mut buf).unwrap();
787        let frozen = buf.freeze();
788        let p = RecordsPayloadBorrowed::from_slice(&frozen).unwrap();
789        assert2::assert!(p.payload_len() == frozen.len());
790
791        let mut out = BytesMut::new();
792        p.encode_to(&mut out).unwrap();
793        assert2::assert!(&out[..] == &frozen[..]);
794    }
795
796    #[test]
797    fn borrowed_encode_decode_via_traits() {
798        let rb = sample_v2();
799        let mut buf = BytesMut::new();
800        rb.encode_batch(&mut buf).unwrap();
801        let frozen = buf.freeze();
802
803        let p = RecordsPayloadBorrowed::from_slice(&frozen).unwrap();
804        let mut out = BytesMut::new();
805        <RecordsPayloadBorrowed as crate::Encode>::encode(&p, &mut out, 0).unwrap();
806        assert2::assert!(
807            <RecordsPayloadBorrowed as crate::Encode>::encoded_len(&p, 0) == frozen.len()
808        );
809
810        let mut cur: &[u8] = &out;
811        let back =
812            <RecordsPayloadBorrowed as crate::DecodeBorrow>::decode_borrow(&mut cur, 0).unwrap();
813        assert2::assert!(matches!(back, RecordsPayloadBorrowed::V2(_)));
814    }
815
816    #[test]
817    fn borrowed_default_is_empty_v2() {
818        let p = RecordsPayloadBorrowed::default();
819        assert2::assert!(matches!(p, RecordsPayloadBorrowed::V2(ref v) if v.is_empty()));
820    }
821
822    #[test]
823    fn from_fetch_bytes_drops_incomplete_trailing_batch() {
824        // Two complete batches followed by a truncated third (only a few
825        // bytes of its header). Kafka sends this when the partition byte
826        // budget cuts the final batch; the consumer must keep the two
827        // complete batches and drop the fragment.
828        let mut b0 = sample_v2();
829        b0.base_offset = 0;
830        let mut b1 = sample_v2();
831        b1.base_offset = 1;
832        let mut buf = BytesMut::new();
833        b0.encode_batch(&mut buf).unwrap();
834        b1.encode_batch(&mut buf).unwrap();
835        buf.extend_from_slice(&[0u8; 7]); // partial trailing batch header
836        let p = RecordsPayload::from_fetch_bytes(buf.freeze()).unwrap();
837        let batches = p.as_v2().expect("v2");
838        assert2::assert!(batches == &[b0, b1][..]);
839    }
840
841    #[test]
842    fn from_fetch_bytes_still_errors_on_corrupt_batch() {
843        // A complete-looking batch whose CRC is wrong must still error, even
844        // leniently — leniency only forgives truncation, not corruption.
845        let rb = sample_v2();
846        let mut buf = BytesMut::new();
847        rb.encode_batch(&mut buf).unwrap();
848        let mut bytes = buf.to_vec();
849        // Corrupt a body byte after the header (HEADER_LEN = 61) to break CRC.
850        bytes[61] ^= 0xFF;
851        let err = RecordsPayload::from_fetch_bytes(Bytes::from(bytes)).unwrap_err();
852        assert2::assert!(matches!(err, RecordsError::CrcMismatch { .. }));
853    }
854
855    #[test]
856    fn from_fetch_bytes_legacy_passes_through() {
857        let bytes = legacy_bytes();
858        let p = RecordsPayload::from_fetch_bytes(bytes.clone()).unwrap();
859        assert2::assert!(p.as_legacy() == Some(&bytes));
860    }
861
862    #[test]
863    fn from_fetch_bytes_empty_is_empty_v2() {
864        // Empty bytes do not carry a magic byte, so looks_like_v2 returns false
865        // and from_fetch_bytes yields an empty Legacy payload (no panic, no error).
866        let p = RecordsPayload::from_fetch_bytes(Bytes::new()).unwrap();
867        assert2::assert!(matches!(p, RecordsPayload::Legacy(ref b) if b.is_empty()));
868    }
869}
krabka_protocol/records/payload.rs

krabka_protocol/records/
payload.rs
