1use std::fmt;
4
5use bytes::Bytes;
6use krabka_compression::RecordDecompressionPolicy;
7use krabka_units::{ByteSize, convert::ByteSizeExt as _};
8use zerocopy::FromBytes as _;
9
10use crate::{
11 primitives::varint::{get_varint, get_varlong},
12 records::{
13 RecordsError,
14 crc::crc32c,
15 header::{Attributes, CRC_COVERAGE_START, HEADER_LEN, RecordBatchHeader},
16 },
17};
18
19const HEADER_TAIL_LEN: i32 = 49;
25
26pub struct RecordBatch<'a> {
27 pub(crate) header: &'a RecordBatchHeader,
28 pub(crate) body: RecordBody<'a>,
29}
30
31pub enum RecordBody<'a> {
32 Borrowed(&'a [u8]),
33 Owned(Bytes),
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct Record<'a> {
38 pub attributes: i8,
39 pub timestamp_delta: i64,
40 pub offset_delta: i32,
41 pub key: Option<&'a [u8]>,
42 pub value: Option<&'a [u8]>,
43 pub headers: Vec<RecordHeader<'a>>,
44}
45
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct RecordHeader<'a> {
48 pub key: &'a str,
49 pub value: Option<&'a [u8]>,
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct RecordSummary<'a> {
59 pub attributes: i8,
60 pub timestamp_delta: i64,
61 pub offset_delta: i32,
62 pub key: Option<&'a [u8]>,
63 pub value: Option<&'a [u8]>,
64 pub header_count: usize,
66 pub body_len: usize,
72}
73
74impl Record<'_> {
75 #[must_use]
84 pub fn body_len(&self) -> usize {
85 super::owned::body_len_of(
86 self.timestamp_delta,
87 self.offset_delta,
88 self.key.map(<[u8]>::len),
89 self.value.map(<[u8]>::len),
90 self.headers
91 .iter()
92 .map(|h| (h.key.len(), h.value.map(<[u8]>::len))),
93 )
94 }
95}
96
97impl RecordBatch<'_> {
98 #[must_use]
99 pub fn header(&self) -> &RecordBatchHeader {
100 self.header
101 }
102
103 #[must_use]
104 pub fn attributes(&self) -> Attributes {
105 Attributes(self.header.attributes.get())
106 }
107}
108
109impl<'a> Default for RecordBatch<'a> {
110 fn default() -> Self {
114 use zerocopy::FromZeros as _;
115 let header: &'a RecordBatchHeader = Box::leak(Box::new(RecordBatchHeader::new_zeroed()));
117 Self {
118 header,
119 body: RecordBody::Owned(bytes::Bytes::new()),
120 }
121 }
122}
123
124impl<'de> crate::DecodeBorrow<'de> for RecordBatch<'de> {
127 fn decode_borrow(buf: &mut &'de [u8], _version: i16) -> Result<Self, crate::ProtocolError> {
128 Self::decode_borrow_with_policy(buf, RecordDecompressionPolicy::default())
129 .map_err(Into::into)
130 }
131}
132
133impl<'de> RecordBatch<'de> {
134 pub fn decode_borrow_with_policy(
141 buf: &mut &'de [u8],
142 policy: RecordDecompressionPolicy,
143 ) -> Result<Self, RecordsError> {
144 decode_borrow_impl(buf, policy)
145 }
146}
147
148fn decode_borrow_impl<'de>(
149 buf: &mut &'de [u8],
150 policy: RecordDecompressionPolicy,
151) -> Result<RecordBatch<'de>, RecordsError> {
152 if buf.len() < HEADER_LEN {
153 return Err(RecordsError::HeaderTooShort {
154 needed: HEADER_LEN - buf.len(),
155 });
156 }
157 let (hdr_slice, rest) = buf.split_at(HEADER_LEN);
159 let hdr: &'de RecordBatchHeader =
160 RecordBatchHeader::ref_from_bytes(hdr_slice).map_err(|_| RecordsError::ZerocopyFailure)?;
161 if hdr.magic != 2 {
162 return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
163 }
164
165 let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
166 .and_then(|n| usize::try_from(n).ok())
167 .ok_or_else(|| RecordsError::RecordParse("negative or oversized batch_length".into()))?;
168
169 let Some((raw_body, after)) = rest.split_at_checked(body_len) else {
170 return Err(RecordsError::BodyTooShort {
171 needed: body_len - rest.len(),
172 });
173 };
174 let computed = crc32c(crc_covered(buf, body_len)?);
176 *buf = after;
177
178 let expected = hdr.crc.get();
179 if computed != expected {
180 return Err(RecordsError::CrcMismatch { expected, computed });
181 }
182
183 let attributes = Attributes(hdr.attributes.get());
184 let codec = attributes.compression();
185 let body = if codec == krabka_compression::CompressionType::None {
186 RecordBody::Borrowed(raw_body)
187 } else {
188 let decompressed = krabka_compression::decompress(
189 codec,
190 raw_body,
191 policy.output_limit(ByteSize::from_bytes(raw_body.len() as u64)),
192 )?;
193 RecordBody::Owned(decompressed)
194 };
195
196 Ok(RecordBatch { header: hdr, body })
197}
198
199#[derive(Debug)]
204pub struct ValidatedBatch<'a> {
205 pub header: &'a RecordBatchHeader,
207 pub total_len: usize,
210 raw_body: &'a [u8],
211}
212
213impl ValidatedBatch<'_> {
214 pub fn validate_records(&self, policy: RecordDecompressionPolicy) -> Result<(), RecordsError> {
228 self.validate_record_summaries_with(policy, |_| {})
229 }
230
231 pub fn validate_records_with(
242 &self,
243 policy: RecordDecompressionPolicy,
244 mut visit: impl FnMut(&Record<'_>),
245 ) -> Result<(), RecordsError> {
246 self.walk_records(policy, |remaining| {
247 visit(&parse_one_record(remaining)?);
248 Ok(())
249 })
250 }
251
252 pub fn validate_record_summaries_with(
263 &self,
264 policy: RecordDecompressionPolicy,
265 mut visit: impl FnMut(&RecordSummary<'_>),
266 ) -> Result<(), RecordsError> {
267 self.walk_records(policy, |remaining| {
268 visit(&parse_one_record_into(remaining, &mut SkipHeaders)?);
269 Ok(())
270 })
271 }
272
273 fn walk_records(
276 &self,
277 policy: RecordDecompressionPolicy,
278 mut parse: impl FnMut(&mut &[u8]) -> Result<(), RecordsError>,
279 ) -> Result<(), RecordsError> {
280 let count = self.header.records_count.get();
281 if count < 0 {
282 return Err(RecordsError::RecordParse(format!(
283 "negative records_count {count}"
284 )));
285 }
286
287 let codec = Attributes(self.header.attributes.get()).compression();
288 let decompressed = if codec == krabka_compression::CompressionType::None {
289 None
290 } else {
291 Some(krabka_compression::decompress(
292 codec,
293 self.raw_body,
294 policy.output_limit(ByteSize::from_bytes(self.raw_body.len() as u64)),
295 )?)
296 };
297 let mut remaining = decompressed.as_deref().unwrap_or(self.raw_body);
298 for i in 0..count {
299 parse(&mut remaining)
300 .map_err(|error| RecordsError::RecordParse(format!("record[{i}]: {error}")))?;
301 }
302 if !remaining.is_empty() {
303 return Err(RecordsError::RecordParse(format!(
304 "trailing bytes after records (left={})",
305 remaining.len()
306 )));
307 }
308 Ok(())
309 }
310}
311
312pub fn validate_one_v2_batch(buf: &[u8]) -> Result<ValidatedBatch<'_>, RecordsError> {
334 if buf.len() < HEADER_LEN {
335 return Err(RecordsError::HeaderTooShort {
336 needed: HEADER_LEN - buf.len(),
337 });
338 }
339 let (hdr_slice, rest) = buf.split_at(HEADER_LEN);
340 let hdr: &RecordBatchHeader =
341 RecordBatchHeader::ref_from_bytes(hdr_slice).map_err(|_| RecordsError::ZerocopyFailure)?;
342 if hdr.magic != 2 {
343 return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
344 }
345 let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
346 .and_then(|n| usize::try_from(n).ok())
347 .ok_or_else(|| RecordsError::RecordParse("negative or oversized batch_length".into()))?;
348 let Some(raw_body) = rest.get(..body_len) else {
349 return Err(RecordsError::BodyTooShort {
350 needed: body_len - rest.len(),
351 });
352 };
353
354 let expected = hdr.crc.get();
355 let computed = crc32c(crc_covered(buf, body_len)?);
356 if computed != expected {
357 return Err(RecordsError::CrcMismatch { expected, computed });
358 }
359
360 Ok(ValidatedBatch {
361 header: hdr,
362 total_len: HEADER_LEN + body_len,
363 raw_body,
364 })
365}
366
367fn crc_covered(batch: &[u8], body_len: usize) -> Result<&[u8], RecordsError> {
373 let end = HEADER_LEN + body_len;
374 batch
375 .get(CRC_COVERAGE_START..end)
376 .ok_or_else(|| RecordsError::BodyTooShort {
377 needed: end.saturating_sub(batch.len()),
378 })
379}
380
381#[must_use]
389pub fn count_records_in_v2_batches(buf: &[u8]) -> u64 {
390 const MAGIC_OFFSET: usize = 16;
393 if buf.get(MAGIC_OFFSET) != Some(&2) {
394 return 0;
395 }
396
397 let mut total = 0u64;
398 let mut remaining = buf;
399 while let Some((hdr_bytes, _)) = remaining.split_first_chunk::<HEADER_LEN>() {
400 if hdr_bytes[MAGIC_OFFSET] != 2 {
401 break;
402 }
403 let Ok(hdr) = RecordBatchHeader::ref_from_bytes(hdr_bytes) else {
404 break;
405 };
406 let Ok(after_len) = usize::try_from(hdr.batch_length.get()) else {
407 break;
408 };
409 let total_len = 12 + after_len;
410 if total_len < HEADER_LEN {
411 break;
412 }
413 let Some(next) = remaining.get(total_len..) else {
414 break;
415 };
416 total += u64::try_from(hdr.records_count.get().max(0)).unwrap_or(0);
417 remaining = next;
418 }
419 total
420}
421
422impl RecordBatch<'_> {
425 pub fn iter(&self) -> RecordIter<'_> {
432 let body: &[u8] = match &self.body {
433 RecordBody::Borrowed(s) => s,
434 RecordBody::Owned(b) => b.as_ref(),
435 };
436 let count = usize::try_from(self.header.records_count.get().max(0)).unwrap_or(usize::MAX);
439 RecordIter {
440 remaining: body,
441 count,
442 index: 0,
443 }
444 }
445}
446
447impl<'a> IntoIterator for &'a RecordBatch<'_> {
448 type Item = Result<Record<'a>, RecordsError>;
449 type IntoIter = RecordIter<'a>;
450
451 fn into_iter(self) -> Self::IntoIter {
452 self.iter()
453 }
454}
455
456pub struct RecordIter<'a> {
457 remaining: &'a [u8],
458 count: usize,
459 index: usize,
460}
461
462impl<'a> Iterator for RecordIter<'a> {
463 type Item = Result<Record<'a>, RecordsError>;
464
465 fn next(&mut self) -> Option<Self::Item> {
466 if self.index >= self.count {
467 return None;
468 }
469 self.index += 1;
470 Some(parse_one_record(&mut self.remaining))
471 }
472}
473
474#[inline]
475pub fn parse_one_record<'a>(buf: &mut &'a [u8]) -> Result<Record<'a>, RecordsError> {
476 let mut headers = Vec::new();
477 let summary = parse_one_record_into(buf, &mut headers)?;
478 Ok(Record {
479 attributes: summary.attributes,
480 timestamp_delta: summary.timestamp_delta,
481 offset_delta: summary.offset_delta,
482 key: summary.key,
483 value: summary.value,
484 headers,
485 })
486}
487
488#[inline]
491fn parse_one_record_into<'a>(
492 buf: &mut &'a [u8],
493 headers: &mut impl HeaderSink<'a>,
494) -> Result<RecordSummary<'a>, RecordsError> {
495 let body_len =
496 get_varlong(buf).map_err(|e| RecordsError::RecordParse(format!("record length: {e}")))?;
497 let body_len = usize::try_from(body_len).map_err(|_| {
498 RecordsError::RecordParse(format!("record length negative or too large: {body_len}"))
499 })?;
500 if buf.len() < body_len {
501 return Err(RecordsError::BodyTooShort {
502 needed: body_len - buf.len(),
503 });
504 }
505 let (body, rest) = buf.split_at(body_len);
506 *buf = rest;
507 let mut body_cur = body;
508 let r = parse_body(&mut body_cur, body_len, headers)?;
509 if !body_cur.is_empty() {
510 return Err(RecordsError::RecordParse(format!(
511 "trailing bytes inside record (left={})",
512 body_cur.len()
513 )));
514 }
515 Ok(r)
516}
517
518trait HeaderSink<'a> {
523 fn reserve(&mut self, count: usize);
526
527 fn push(&mut self, header: RecordHeader<'a>);
528}
529
530impl<'a> HeaderSink<'a> for Vec<RecordHeader<'a>> {
531 fn reserve(&mut self, count: usize) {
532 self.reserve_exact(count);
533 }
534
535 fn push(&mut self, header: RecordHeader<'a>) {
536 Vec::push(self, header);
537 }
538}
539
540struct SkipHeaders;
542
543impl HeaderSink<'_> for SkipHeaders {
544 fn reserve(&mut self, _count: usize) {}
545
546 fn push(&mut self, _header: RecordHeader<'_>) {}
547}
548
549fn parse_body<'a>(
550 buf: &mut &'a [u8],
551 body_len: usize,
552 headers: &mut impl HeaderSink<'a>,
553) -> Result<RecordSummary<'a>, RecordsError> {
554 let Some((&attributes, rest)) = buf.split_first() else {
555 return Err(RecordsError::RecordParse("record body empty".into()));
556 };
557 let attributes = i8::from_ne_bytes([attributes]);
558 *buf = rest;
559 let timestamp_delta =
560 get_varlong(buf).map_err(|e| RecordsError::RecordParse(format!("timestamp_delta: {e}")))?;
561 let offset_delta =
562 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("offset_delta: {e}")))?;
563
564 let key = read_nullable_slice(buf, SliceLabel::Key)?;
565 let value = read_nullable_slice(buf, SliceLabel::Value)?;
566
567 let header_count =
568 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("header_count: {e}")))?;
569 let Ok(header_count_usize) = usize::try_from(header_count) else {
570 return Err(RecordsError::RecordParse(format!(
571 "negative header count {header_count}"
572 )));
573 };
574 headers.reserve(header_count_usize.min(buf.len()));
578 for i in 0..header_count {
579 let key_len = get_varint(buf)
580 .map_err(|e| RecordsError::RecordParse(format!("header[{i}] key length: {e}")))?;
581 let Ok(n) = usize::try_from(key_len) else {
582 return Err(RecordsError::RecordParse(format!(
583 "header[{i}] negative key length"
584 )));
585 };
586 if buf.len() < n {
587 return Err(RecordsError::BodyTooShort {
588 needed: n - buf.len(),
589 });
590 }
591 let (key_bytes, rest) = buf.split_at(n);
592 *buf = rest;
593 let key_str = str::from_utf8(key_bytes)
594 .map_err(|e| RecordsError::RecordParse(format!("header[{i}] key utf-8: {e}")))?;
595
596 let value = read_nullable_slice(buf, SliceLabel::HeaderValue(i))?;
597 headers.push(RecordHeader {
598 key: key_str,
599 value,
600 });
601 }
602
603 Ok(RecordSummary {
604 attributes,
605 timestamp_delta,
606 offset_delta,
607 key,
608 value,
609 header_count: header_count_usize,
610 body_len,
611 })
612}
613
614#[derive(Debug, Clone, Copy, PartialEq, Eq)]
620enum SliceLabel {
621 Key,
622 Value,
623 HeaderValue(i32),
624}
625
626impl fmt::Display for SliceLabel {
627 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
628 match self {
629 Self::Key => f.write_str("key"),
630 Self::Value => f.write_str("value"),
631 Self::HeaderValue(index) => write!(f, "header[{index}] value"),
632 }
633 }
634}
635
636fn read_nullable_slice<'a>(
637 buf: &mut &'a [u8],
638 label: SliceLabel,
639) -> Result<Option<&'a [u8]>, RecordsError> {
640 let len =
641 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("{label} length: {e}")))?;
642 let Ok(n) = usize::try_from(len) else {
644 return Ok(None);
645 };
646 if buf.len() < n {
647 return Err(RecordsError::BodyTooShort {
648 needed: n - buf.len(),
649 });
650 }
651 let (head, rest) = buf.split_at(n);
652 *buf = rest;
653 Ok(Some(head))
654}
655
656impl RecordBatch<'_> {
659 pub fn to_owned(&self) -> Result<super::owned::RecordBatch, RecordsError> {
664 let mut records = Vec::new();
665 for r in self {
666 let r = r?;
667 records.push(super::owned::Record {
668 attributes: r.attributes,
669 timestamp_delta: r.timestamp_delta,
670 offset_delta: r.offset_delta,
671 key: r.key.map(Bytes::copy_from_slice),
672 value: r.value.map(Bytes::copy_from_slice),
673 headers: r
674 .headers
675 .into_iter()
676 .map(|h| super::owned::RecordHeader {
677 key: h.key.to_owned(),
678 value: h.value.map(Bytes::copy_from_slice),
679 })
680 .collect(),
681 });
682 }
683 Ok(super::owned::RecordBatch {
684 base_offset: self.header.base_offset.get(),
685 partition_leader_epoch: self.header.partition_leader_epoch.get(),
686 attributes: self.attributes(),
687 last_offset_delta: self.header.last_offset_delta.get(),
688 base_timestamp: self.header.base_timestamp.get(),
689 max_timestamp: self.header.max_timestamp.get(),
690 producer_id: self.header.producer_id.get(),
691 producer_epoch: self.header.producer_epoch.get(),
692 base_sequence: self.header.base_sequence.get(),
693 records,
694 })
695 }
696}
697
698impl fmt::Debug for RecordBatch<'_> {
704 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
705 match self.to_owned() {
706 Ok(o) => o.fmt(f),
707 Err(e) => write!(f, "RecordBatch(<decode error: {e}>)"),
708 }
709 }
710}
711
712impl Clone for RecordBatch<'_> {
713 fn clone(&self) -> Self {
717 RecordBatch {
718 header: self.header,
719 body: match &self.body {
720 RecordBody::Borrowed(s) => RecordBody::Borrowed(s),
721 RecordBody::Owned(b) => RecordBody::Owned(b.clone()),
722 },
723 }
724 }
725}
726
727impl PartialEq for RecordBatch<'_> {
728 fn eq(&self, other: &Self) -> bool {
729 match (self.to_owned(), other.to_owned()) {
730 (Ok(a), Ok(b)) => a == b,
731 _ => false,
732 }
733 }
734}
735
736impl Eq for RecordBatch<'_> {}
737
738impl crate::Encode for RecordBatch<'_> {
741 fn encode<B: bytes::BufMut>(
742 &self,
743 buf: &mut B,
744 version: i16,
745 ) -> Result<(), crate::ProtocolError> {
746 let owned = self.to_owned().map_err(crate::ProtocolError::from)?;
747 crate::Encode::encode(&owned, buf, version)
748 }
749
750 fn encoded_len(&self, version: i16) -> usize {
751 self.to_owned()
752 .map_or(0, |o| crate::Encode::encoded_len(&o, version))
753 }
754}
755
756#[cfg(test)]
759mod tests {
760
761 use bytes::BytesMut;
762 use krabka_compression::{CompressionError, CompressionType, RecordDecompressionPolicy};
763 use krabka_units::{bytes, fraction};
764
765 use super::*;
766 use crate::{DecodeBorrow, primitives::varint::put_varlong};
767
768 fn encode_owned_then_borrow(b: &super::super::owned::RecordBatch) -> Vec<u8> {
769 let mut buf = BytesMut::new();
770 b.encode_batch(&mut buf).unwrap();
771 buf.to_vec()
772 }
773
774 fn refresh_batch_crc(encoded: &mut [u8]) {
775 let coverage_start = CRC_COVERAGE_START;
776 let crc = crc32c(&encoded[coverage_start..]);
777 encoded[coverage_start - 4..coverage_start].copy_from_slice(&crc.to_be_bytes());
778 }
779
780 #[test]
781 fn borrowed_roundtrip_cases() {
782 for (_case, codec) in [
783 ("none", CompressionType::None),
784 ("gzip", CompressionType::Gzip),
785 ("snappy", CompressionType::Snappy),
786 ("lz4", CompressionType::Lz4),
787 ("zstd", CompressionType::Zstd),
788 ] {
789 let mut owned = super::super::owned::RecordBatch::default();
790 owned.attributes = owned.attributes.with_compression(codec);
791 owned.records.push(super::super::owned::Record {
792 key: Some(Bytes::from_static(b"key")),
793 value: Some(Bytes::from_static(b"value")),
794 ..Default::default()
795 });
796
797 let encoded = encode_owned_then_borrow(&owned);
798 let mut cur: &[u8] = &encoded[..];
799 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
800 assert2::assert!(cur.is_empty());
801 assert2::assert!(borrowed.attributes() == owned.attributes);
802
803 let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
804 let expected_records = vec![Record {
805 attributes: 0,
806 timestamp_delta: 0,
807 offset_delta: 0,
808 key: Some(b"key".as_slice()),
809 value: Some(b"value".as_slice()),
810 headers: vec![],
811 }];
812 assert2::assert!(records == expected_records);
813
814 let back_owned = borrowed.to_owned().unwrap();
815 assert2::assert!(back_owned == owned);
816 }
817 }
818
819 #[test]
824 fn borrowed_body_len_matches_the_owned_record() {
825 use super::super::owned::{Record as OwnedRecord, RecordHeader as OwnedHeader};
826
827 let owned = super::super::owned::RecordBatch {
828 records: vec![
829 OwnedRecord::default(),
830 OwnedRecord {
831 timestamp_delta: 17,
832 offset_delta: 2,
833 key: Some(Bytes::from_static(b"the-key")),
834 value: Some(Bytes::from_static(b"hello kafka")),
835 headers: vec![
836 OwnedHeader {
837 key: "trace-id".to_owned(),
838 value: Some(Bytes::from_static(b"abc")),
839 },
840 OwnedHeader {
841 key: "null-val".to_owned(),
842 value: None,
843 },
844 ],
845 ..Default::default()
846 },
847 OwnedRecord {
848 value: Some(Bytes::from(vec![7_u8; 57])),
849 ..Default::default()
850 },
851 OwnedRecord {
852 value: Some(Bytes::from(vec![7_u8; 58])),
853 ..Default::default()
854 },
855 ],
856 last_offset_delta: 3,
857 ..Default::default()
858 };
859 let encoded = encode_owned_then_borrow(&owned);
860 let mut cur: &[u8] = &encoded[..];
861 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
862 let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
863
864 let lens: Vec<usize> = records.iter().map(Record::body_len).collect();
865 assert2::assert!(lens == vec![6, 47, 63, 64]);
866 let owned_lens: Vec<usize> = owned.records.iter().map(OwnedRecord::body_len).collect();
867 assert2::assert!(lens == owned_lens);
868 }
869
870 #[test]
871 fn decompression_policy_limits_borrowed_decode() {
872 let mut owned = super::super::owned::RecordBatch {
873 attributes: Attributes::default().with_compression(CompressionType::Lz4),
874 records: vec![super::super::owned::Record {
875 value: Some(Bytes::from(vec![b'x'; 4096])),
876 ..Default::default()
877 }],
878 ..Default::default()
879 };
880 owned.last_offset_delta = 0;
881 let encoded = encode_owned_then_borrow(&owned);
882
883 let mut default_cur = &encoded[..];
884 RecordBatch::decode_borrow(&mut default_cur, 0).unwrap();
885
886 let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
887 let mut limited_cur = &encoded[..];
888 assert2::assert!(matches!(
889 RecordBatch::decode_borrow_with_policy(&mut limited_cur, policy),
890 Err(RecordsError::Compression(CompressionError::TooLarge {
891 limit: 32
892 }))
893 ));
894 }
895
896 #[test]
897 fn zero_copy_for_uncompressed() {
898 let mut owned = super::super::owned::RecordBatch::default();
901 owned.records.push(super::super::owned::Record {
902 key: Some(Bytes::from_static(b"k")),
903 value: Some(Bytes::from_static(b"v")),
904 ..Default::default()
905 });
906 let encoded = encode_owned_then_borrow(&owned);
907 let encoded_start = encoded.as_ptr() as usize;
908 let encoded_end = encoded_start + encoded.len();
909
910 let mut cur: &[u8] = &encoded[..];
911 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
912 let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
913
914 let v_ptr = records[0].value.unwrap().as_ptr() as usize;
915 assert2::assert!(v_ptr >= encoded_start && v_ptr < encoded_end);
916 }
917
918 #[test]
919 fn validate_one_v2_batch_reads_header_and_len() {
920 let mut owned = super::super::owned::RecordBatch {
921 base_offset: 7,
922 partition_leader_epoch: 3,
923 last_offset_delta: 0,
924 producer_id: 99,
925 producer_epoch: 1,
926 base_sequence: 5,
927 max_timestamp: 1_234,
928 ..Default::default()
929 };
930 owned.records.push(super::super::owned::Record {
931 value: Some(Bytes::from_static(b"payload")),
932 ..Default::default()
933 });
934 let encoded = encode_owned_then_borrow(&owned);
935
936 let v = validate_one_v2_batch(&encoded).unwrap();
937 assert2::assert!(
938 (
939 v.total_len,
940 v.header.base_offset.get(),
941 v.header.partition_leader_epoch.get(),
942 v.header.producer_id.get(),
943 v.header.producer_epoch.get(),
944 v.header.base_sequence.get(),
945 v.header.max_timestamp.get(),
946 v.header.magic,
947 ) == (encoded.len(), 7, 3, 99, 1, 5, 1_234, 2)
948 );
949 assert2::assert!(
950 v.validate_records(RecordDecompressionPolicy::default())
951 .is_ok()
952 );
953 }
954
955 #[test]
956 fn borrowed_crc_stops_at_the_declared_batch_boundary() {
957 for codec in [CompressionType::None, CompressionType::Lz4] {
958 for value in [Bytes::new(), Bytes::from_static(b"payload")] {
959 let mut owned = super::super::owned::RecordBatch {
960 records: vec![super::super::owned::Record {
961 value: Some(value),
962 ..Default::default()
963 }],
964 ..Default::default()
965 };
966 owned.attributes = owned.attributes.with_compression(codec);
967 let mut encoded = encode_owned_then_borrow(&owned);
968 let first_len = encoded.len();
969 encoded.extend_from_within(..);
970 encoded[first_len + 21] ^= 1;
972
973 let validated = validate_one_v2_batch(&encoded).unwrap();
974 assert2::assert!(validated.total_len == first_len);
975 let mut input = encoded.as_slice();
976 let borrowed = RecordBatch::decode_borrow(&mut input, 0).unwrap();
977 assert2::assert!(borrowed.to_owned().unwrap() == owned);
978 assert2::assert!(input == &encoded[first_len..]);
979 assert2::assert!(matches!(
980 validate_one_v2_batch(input),
981 Err(RecordsError::CrcMismatch { .. })
982 ));
983 assert2::assert!(matches!(
984 RecordBatch::decode_borrow_with_policy(
985 &mut input,
986 RecordDecompressionPolicy::default(),
987 ),
988 Err(RecordsError::CrcMismatch { .. })
989 ));
990 }
991 }
992 }
993
994 #[test]
995 fn validated_batch_rejects_crc_valid_malformed_record_body() {
996 let owned = super::super::owned::RecordBatch {
997 records: vec![super::super::owned::Record {
998 value: Some(Bytes::from_static(b"value")),
999 ..Default::default()
1000 }],
1001 ..Default::default()
1002 };
1003 let mut encoded = encode_owned_then_borrow(&owned);
1004 encoded[HEADER_LEN] = 0; refresh_batch_crc(&mut encoded);
1006
1007 let validated = validate_one_v2_batch(&encoded).expect("CRC-valid header");
1008 let error = validated
1009 .validate_records(RecordDecompressionPolicy::default())
1010 .unwrap_err();
1011 assert2::assert!(matches!(error, RecordsError::RecordParse(_)));
1012 }
1013
1014 #[test]
1015 fn validated_batch_visits_records_during_validation() {
1016 let owned = super::super::owned::RecordBatch {
1017 attributes: Attributes::default().with_compression(CompressionType::Lz4),
1018 records: vec![
1019 super::super::owned::Record {
1020 timestamp_delta: 3,
1021 value: Some(Bytes::from_static(b"first")),
1022 ..Default::default()
1023 },
1024 super::super::owned::Record {
1025 timestamp_delta: 8,
1026 value: Some(Bytes::from_static(b"second")),
1027 ..Default::default()
1028 },
1029 ],
1030 ..Default::default()
1031 };
1032 let encoded = encode_owned_then_borrow(&owned);
1033 let validated = validate_one_v2_batch(&encoded).unwrap();
1034 let mut timestamps = Vec::new();
1035
1036 validated
1037 .validate_records_with(RecordDecompressionPolicy::default(), |record| {
1038 timestamps.push(record.timestamp_delta);
1039 })
1040 .unwrap();
1041
1042 assert2::assert!(timestamps == vec![3, 8]);
1043 }
1044
1045 #[test]
1046 fn validated_batch_visits_record_summaries_during_validation() {
1047 use super::super::owned::{Record as OwnedRecord, RecordHeader as OwnedHeader};
1048
1049 for codec in [CompressionType::None, CompressionType::Lz4] {
1050 let owned = super::super::owned::RecordBatch {
1051 attributes: Attributes::default().with_compression(codec),
1052 last_offset_delta: 1,
1053 records: vec![
1054 OwnedRecord {
1055 timestamp_delta: 3,
1056 value: Some(Bytes::from_static(b"first")),
1057 ..Default::default()
1058 },
1059 OwnedRecord {
1060 attributes: 0,
1061 timestamp_delta: 8,
1062 offset_delta: 1,
1063 key: Some(Bytes::from_static(b"k")),
1064 value: None,
1065 headers: vec![
1066 OwnedHeader {
1067 key: "trace-id".to_owned(),
1068 value: Some(Bytes::from_static(b"abc")),
1069 },
1070 OwnedHeader {
1071 key: "null-val".to_owned(),
1072 value: None,
1073 },
1074 ],
1075 },
1076 ],
1077 ..Default::default()
1078 };
1079 let encoded = encode_owned_then_borrow(&owned);
1080 let validated = validate_one_v2_batch(&encoded).unwrap();
1081 let mut summaries = Vec::new();
1082 validated
1083 .validate_record_summaries_with(RecordDecompressionPolicy::default(), |summary| {
1084 summaries.push((
1085 summary.attributes,
1086 summary.timestamp_delta,
1087 summary.offset_delta,
1088 summary.key.map(<[u8]>::to_vec),
1089 summary.value.map(<[u8]>::to_vec),
1090 summary.header_count,
1091 summary.body_len,
1092 ));
1093 })
1094 .unwrap();
1095
1096 let expected: Vec<_> = owned
1097 .records
1098 .iter()
1099 .map(|record| {
1100 (
1101 record.attributes,
1102 record.timestamp_delta,
1103 record.offset_delta,
1104 record.key.as_deref().map(<[u8]>::to_vec),
1105 record.value.as_deref().map(<[u8]>::to_vec),
1106 record.headers.len(),
1107 record.body_len(),
1108 )
1109 })
1110 .collect();
1111 assert2::assert!(summaries == expected, "{codec:?}");
1112 }
1113 }
1114
1115 #[test]
1118 fn summary_walk_rejects_malformed_headers_like_the_collecting_walk() {
1119 let owned = super::super::owned::RecordBatch {
1120 records: vec![super::super::owned::Record {
1121 headers: vec![super::super::owned::RecordHeader {
1122 key: "ab".to_owned(),
1123 value: Some(Bytes::from_static(b"v")),
1124 }],
1125 ..Default::default()
1126 }],
1127 ..Default::default()
1128 };
1129 let encoded = encode_owned_then_borrow(&owned);
1130 let key_len_at = HEADER_LEN + 7;
1134 for (case, offset, byte) in [
1135 ("invalid utf-8 key", key_len_at + 1, 0xFF),
1136 ("negative key length", key_len_at, 0x01),
1137 ("key overruns record", key_len_at, 0x7E),
1138 ] {
1139 let mut corrupt = encoded.clone();
1140 corrupt[offset] = byte;
1141 refresh_batch_crc(&mut corrupt);
1142 let validated = validate_one_v2_batch(&corrupt).unwrap();
1143 let policy = RecordDecompressionPolicy::default();
1144 let collected = validated
1145 .validate_records_with(policy, |_| {})
1146 .unwrap_err()
1147 .to_string();
1148 let summarized = validated
1149 .validate_record_summaries_with(policy, |_| {})
1150 .unwrap_err()
1151 .to_string();
1152 assert2::assert!(summarized == collected, "{case}");
1153 }
1154 }
1155
1156 #[test]
1157 fn validated_batch_rejects_trailing_body_bytes() {
1158 let owned = super::super::owned::RecordBatch {
1159 records: vec![super::super::owned::Record::default()],
1160 ..Default::default()
1161 };
1162 let mut encoded = encode_owned_then_borrow(&owned);
1163 let batch_length = i32::from_be_bytes(encoded[8..12].try_into().unwrap());
1164 encoded[8..12].copy_from_slice(&(batch_length + 1).to_be_bytes());
1165 encoded.push(0);
1166 refresh_batch_crc(&mut encoded);
1167
1168 let validated = validate_one_v2_batch(&encoded).expect("CRC-valid header");
1169 let error = validated
1170 .validate_records(RecordDecompressionPolicy::default())
1171 .unwrap_err();
1172 assert2::assert!(
1173 matches!(error, RecordsError::RecordParse(message) if message.contains("trailing bytes"))
1174 );
1175 }
1176
1177 #[test]
1178 fn validated_batch_applies_decompression_policy() {
1179 let owned = super::super::owned::RecordBatch {
1180 attributes: Attributes::default().with_compression(CompressionType::Lz4),
1181 records: vec![super::super::owned::Record {
1182 value: Some(Bytes::from(vec![b'x'; 4096])),
1183 ..Default::default()
1184 }],
1185 ..Default::default()
1186 };
1187 let encoded = encode_owned_then_borrow(&owned);
1188 let validated = validate_one_v2_batch(&encoded).unwrap();
1189 let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
1190
1191 assert2::assert!(matches!(
1192 validated.validate_records(policy),
1193 Err(RecordsError::Compression(CompressionError::TooLarge {
1194 limit: 32
1195 }))
1196 ));
1197 }
1198
1199 #[test]
1200 fn validated_batch_distinguishes_zero_and_negative_record_counts() {
1201 let mut encoded = encode_owned_then_borrow(&super::super::owned::RecordBatch::default());
1202 let zero = validate_one_v2_batch(&encoded).unwrap();
1203 assert2::assert!(
1204 zero.validate_records(RecordDecompressionPolicy::default())
1205 .is_ok()
1206 );
1207
1208 encoded[57..61].copy_from_slice(&(-1i32).to_be_bytes());
1209 refresh_batch_crc(&mut encoded);
1210 let negative = validate_one_v2_batch(&encoded).unwrap();
1211 assert2::assert!(matches!(
1212 negative.validate_records(RecordDecompressionPolicy::default()),
1213 Err(RecordsError::RecordParse(message)) if message.contains("negative records_count")
1214 ));
1215 }
1216
1217 #[test]
1218 fn validate_one_v2_batch_rejects_corrupt_crc() {
1219 let owned = super::super::owned::RecordBatch {
1220 records: vec![super::super::owned::Record {
1221 value: Some(Bytes::from_static(b"x")),
1222 ..Default::default()
1223 }],
1224 ..Default::default()
1225 };
1226 let mut encoded = encode_owned_then_borrow(&owned);
1227 encoded[HEADER_LEN] ^= 0xFF;
1229 let err = validate_one_v2_batch(&encoded).unwrap_err();
1230 assert2::assert!(matches!(err, RecordsError::CrcMismatch { .. }));
1231 }
1232
1233 #[test]
1234 fn validate_one_v2_batch_rejects_truncated() {
1235 let owned = super::super::owned::RecordBatch {
1236 records: vec![super::super::owned::Record {
1237 value: Some(Bytes::from_static(b"value")),
1238 ..Default::default()
1239 }],
1240 ..Default::default()
1241 };
1242 let encoded = encode_owned_then_borrow(&owned);
1243 let err = validate_one_v2_batch(&encoded[..encoded.len() - 2]).unwrap_err();
1244 assert2::assert!(matches!(err, RecordsError::BodyTooShort { .. }));
1245 }
1246
1247 #[test]
1248 fn count_records_in_v2_batches_counts_single_and_concatenated_batches() {
1249 let one = encode_owned_then_borrow(&super::super::owned::RecordBatch {
1250 records: vec![super::super::owned::Record::default()],
1251 ..Default::default()
1252 });
1253 let three = encode_owned_then_borrow(&super::super::owned::RecordBatch {
1254 last_offset_delta: 2,
1255 records: vec![
1256 super::super::owned::Record::default(),
1257 super::super::owned::Record::default(),
1258 super::super::owned::Record::default(),
1259 ],
1260 ..Default::default()
1261 });
1262 let mut both = Vec::with_capacity(one.len() + three.len());
1263 both.extend_from_slice(&one);
1264 both.extend_from_slice(&three);
1265
1266 for (_case, input, expected) in [("single", &one[..], 1), ("concatenated", &both[..], 4)] {
1267 assert2::assert!(count_records_in_v2_batches(input) == expected);
1268 }
1269 }
1270
1271 #[test]
1272 fn count_records_in_v2_batches_stops_at_bad_or_non_v2_input() {
1273 let encoded = encode_owned_then_borrow(&super::super::owned::RecordBatch {
1274 records: vec![super::super::owned::Record::default()],
1275 ..Default::default()
1276 });
1277 let mut with_truncated_tail = encoded.clone();
1278 with_truncated_tail.extend_from_slice(&encoded[..HEADER_LEN - 1]);
1279
1280 for (_case, input, expected) in [
1281 ("empty", &[][..], 0),
1282 ("non-v2", &[0u8; HEADER_LEN][..], 0),
1283 ("truncated tail", &with_truncated_tail[..], 1),
1284 ] {
1285 assert2::assert!(count_records_in_v2_batches(input) == expected);
1286 }
1287 }
1288
1289 fn v2_header_only(batch_length: i32, records_count: i32) -> Vec<u8> {
1295 let mut buf = vec![0u8; HEADER_LEN];
1296 buf[16] = 2; buf[8..12].copy_from_slice(&batch_length.to_be_bytes());
1298 buf[57..61].copy_from_slice(&records_count.to_be_bytes());
1299 buf
1300 }
1301
1302 #[test]
1303 fn count_records_in_v2_batches_accepts_batch_that_exactly_fills_buffer() {
1304 let buf = v2_header_only(49, 7);
1309 assert2::assert!(buf.len() == HEADER_LEN);
1310 assert2::assert!(count_records_in_v2_batches(&buf) == 7);
1311 }
1312
1313 #[test]
1314 fn count_records_in_v2_batches_rejects_batch_overrunning_buffer() {
1315 let buf = v2_header_only(100, 9);
1320 assert2::assert!(count_records_in_v2_batches(&buf) == 0);
1321 }
1322
1323 fn framed_record(body: &[u8]) -> Vec<u8> {
1325 let mut framed = Vec::new();
1326 put_varlong(&mut framed, i64::try_from(body.len()).unwrap());
1327 framed.extend_from_slice(body);
1328 framed
1329 }
1330
1331 #[test]
1332 fn slice_labels_name_the_field() {
1333 for (label, expected) in [
1334 (SliceLabel::Key, "key"),
1335 (SliceLabel::Value, "value"),
1336 (SliceLabel::HeaderValue(0), "header[0] value"),
1337 (SliceLabel::HeaderValue(17), "header[17] value"),
1338 ] {
1339 assert2::assert!(label.to_string() == expected);
1340 }
1341 }
1342
1343 #[test]
1346 fn truncated_nullable_length_reports_the_field_label() {
1347 let varint_error = get_varint(&mut &[0x80_u8][..]).unwrap_err();
1348 let prefix: &[u8] = &[0, 0, 0];
1350 for (case, tail, label) in [
1351 ("key", &[0x80_u8][..], "key"),
1352 ("value", &[0x01, 0x80][..], "value"),
1353 (
1356 "second header value",
1357 &[0x01, 0x01, 0x04, 0x02, b'a', 0x01, 0x02, b'b', 0x80][..],
1358 "header[1] value",
1359 ),
1360 ] {
1361 let framed = framed_record(&[prefix, tail].concat());
1362 let error = parse_one_record(&mut framed.as_slice()).unwrap_err();
1363 let expected = format!("{label} length: {varint_error}");
1364 assert2::assert!(
1365 matches!(&error, RecordsError::RecordParse(message) if *message == expected),
1366 "{case}: {error:?}"
1367 );
1368 }
1369 }
1370
1371 #[test]
1372 fn borrowed_encode_via_trait_roundtrips() {
1373 use crate::Encode as _;
1374 let owned_in = super::super::owned::RecordBatch {
1375 records: vec![super::super::owned::Record {
1376 key: Some(Bytes::from_static(b"x")),
1377 value: Some(Bytes::from_static(b"y")),
1378 ..Default::default()
1379 }],
1380 ..Default::default()
1381 };
1382 let bytes_in = encode_owned_then_borrow(&owned_in);
1383 let mut cur: &[u8] = &bytes_in[..];
1384 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
1385
1386 let mut out = BytesMut::new();
1387 borrowed.encode(&mut out, 0).unwrap();
1388 assert2::assert!(&out[..] == &bytes_in[..]);
1389 }
1390}