1use 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
17const 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 #[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 fn default() -> Self {
90 use zerocopy::FromZeros as _;
91 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
100impl<'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 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 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 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#[derive(Debug)]
181pub struct ValidatedBatch<'a> {
182 pub header: &'a RecordBatchHeader,
184 pub total_len: usize,
187 raw_body: &'a [u8],
188}
189
190impl ValidatedBatch<'_> {
191 pub fn validate_records(&self, policy: RecordDecompressionPolicy) -> Result<(), RecordsError> {
205 self.validate_records_with(policy, |_| {})
206 }
207
208 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
254pub 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#[must_use]
318pub fn count_records_in_v2_batches(buf: &[u8]) -> u64 {
319 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
348impl RecordBatch<'_> {
351 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 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
511impl RecordBatch<'_> {
514 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
553impl 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 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
593impl 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#[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 #[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 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 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; 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 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 fn v2_header_only(batch_length: i32, records_count: i32) -> Vec<u8> {
1041 let mut buf = vec![0u8; HEADER_LEN];
1042 buf[16] = 2; 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 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 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}