1use bytes::Bytes;
4use zerocopy::FromBytes as _;
5
6use crate::primitives::varint::{get_varint, get_varlong};
7use crate::records::RecordsError;
8use crate::records::crc::{crc32c, crc32c_append};
9use crate::records::header::{Attributes, HEADER_LEN, RecordBatchHeader};
10
11const HEADER_TAIL_LEN: i32 = 49;
17
18pub struct RecordBatch<'a> {
19 pub(crate) header: &'a RecordBatchHeader,
20 pub(crate) body: RecordBody<'a>,
21}
22
23pub(crate) enum RecordBody<'a> {
24 Borrowed(&'a [u8]),
25 Owned(Bytes),
26}
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct Record<'a> {
30 pub attributes: i8,
31 pub timestamp_delta: i64,
32 pub offset_delta: i32,
33 pub key: Option<&'a [u8]>,
34 pub value: Option<&'a [u8]>,
35 pub headers: Vec<RecordHeader<'a>>,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct RecordHeader<'a> {
40 pub key: &'a str,
41 pub value: Option<&'a [u8]>,
42}
43
44impl RecordBatch<'_> {
45 #[must_use]
46 pub fn header(&self) -> &RecordBatchHeader {
47 self.header
48 }
49
50 #[must_use]
51 pub fn attributes(&self) -> Attributes {
52 Attributes(self.header.attributes.get())
53 }
54}
55
56impl<'a> Default for RecordBatch<'a> {
57 fn default() -> Self {
61 use zerocopy::FromZeros as _;
62 let header: &'a RecordBatchHeader = Box::leak(Box::new(RecordBatchHeader::new_zeroed()));
64 Self {
65 header,
66 body: RecordBody::Owned(bytes::Bytes::new()),
67 }
68 }
69}
70
71impl<'de> crate::DecodeBorrow<'de> for RecordBatch<'de> {
74 fn decode_borrow(buf: &mut &'de [u8], _version: i16) -> Result<Self, crate::ProtocolError> {
75 decode_borrow_impl(buf).map_err(Into::into)
76 }
77}
78
79fn decode_borrow_impl<'de>(buf: &mut &'de [u8]) -> Result<RecordBatch<'de>, RecordsError> {
80 if buf.len() < HEADER_LEN {
81 return Err(RecordsError::HeaderTooShort {
82 needed: HEADER_LEN - buf.len(),
83 });
84 }
85 let (hdr_slice, rest) = buf.split_at(HEADER_LEN);
87 let hdr: &'de RecordBatchHeader =
88 RecordBatchHeader::ref_from_bytes(hdr_slice).map_err(|_| RecordsError::ZerocopyFailure)?;
89 if hdr.magic != 2 {
90 return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
91 }
92
93 let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
94 .and_then(|n| usize::try_from(n).ok())
95 .ok_or_else(|| RecordsError::RecordParse("negative or oversized batch_length".into()))?;
96
97 if rest.len() < body_len {
98 return Err(RecordsError::BodyTooShort {
99 needed: body_len - rest.len(),
100 });
101 }
102 let (raw_body, after) = rest.split_at(body_len);
103 *buf = after;
104
105 let expected = hdr.crc.get();
108 let mut computed = crc32c(&hdr_slice[21..HEADER_LEN]);
109 computed = crc32c_append(computed, raw_body);
110 if computed != expected {
111 return Err(RecordsError::CrcMismatch { expected, computed });
112 }
113
114 let attributes = Attributes(hdr.attributes.get());
115 let codec = attributes.compression();
116 let body = if codec == crabka_compression::CompressionType::None {
117 RecordBody::Borrowed(raw_body)
118 } else {
119 const DECOMPRESS_MIN_CAP: usize = 16 * 1024 * 1024; const DECOMPRESS_MAX_RATIO: usize = 100; const DECOMPRESS_ABSOLUTE_CEILING: usize = 1024 * 1024 * 1024; let max_output = raw_body
126 .len()
127 .saturating_mul(DECOMPRESS_MAX_RATIO)
128 .clamp(DECOMPRESS_MIN_CAP, DECOMPRESS_ABSOLUTE_CEILING);
129 let decompressed = crabka_compression::decompress(codec, raw_body, max_output)?;
130 RecordBody::Owned(decompressed)
131 };
132
133 Ok(RecordBatch { header: hdr, body })
134}
135
136#[derive(Debug)]
140pub struct ValidatedBatch<'a> {
141 pub header: &'a RecordBatchHeader,
143 pub total_len: usize,
146}
147
148pub fn validate_one_v2_batch(buf: &[u8]) -> Result<ValidatedBatch<'_>, RecordsError> {
169 if buf.len() < HEADER_LEN {
170 return Err(RecordsError::HeaderTooShort {
171 needed: HEADER_LEN - buf.len(),
172 });
173 }
174 let (hdr_slice, rest) = buf.split_at(HEADER_LEN);
175 let hdr: &RecordBatchHeader =
176 RecordBatchHeader::ref_from_bytes(hdr_slice).map_err(|_| RecordsError::ZerocopyFailure)?;
177 if hdr.magic != 2 {
178 return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
179 }
180 let body_len = i32::checked_sub(hdr.batch_length.get(), HEADER_TAIL_LEN)
181 .and_then(|n| usize::try_from(n).ok())
182 .ok_or_else(|| RecordsError::RecordParse("negative or oversized batch_length".into()))?;
183 if rest.len() < body_len {
184 return Err(RecordsError::BodyTooShort {
185 needed: body_len - rest.len(),
186 });
187 }
188 let raw_body = &rest[..body_len];
189
190 let expected = hdr.crc.get();
191 let mut computed = crc32c(&hdr_slice[21..HEADER_LEN]);
192 computed = crc32c_append(computed, raw_body);
193 if computed != expected {
194 return Err(RecordsError::CrcMismatch { expected, computed });
195 }
196
197 Ok(ValidatedBatch {
198 header: hdr,
199 total_len: HEADER_LEN + body_len,
200 })
201}
202
203#[must_use]
211pub fn count_records_in_v2_batches(buf: &[u8]) -> u64 {
212 const MAGIC_OFFSET: usize = 16;
215 if buf.len() <= MAGIC_OFFSET || buf[MAGIC_OFFSET] != 2 {
216 return 0;
217 }
218
219 let mut total = 0u64;
220 let mut remaining = buf;
221 while remaining.len() >= HEADER_LEN {
222 if remaining[MAGIC_OFFSET] != 2 {
223 break;
224 }
225 let Ok(hdr) = RecordBatchHeader::ref_from_bytes(&remaining[..HEADER_LEN]) else {
226 break;
227 };
228 let Ok(after_len) = usize::try_from(hdr.batch_length.get()) else {
229 break;
230 };
231 let total_len = 12 + after_len;
232 if total_len < HEADER_LEN || total_len > remaining.len() {
233 break;
234 }
235 total += u64::try_from(hdr.records_count.get().max(0)).unwrap_or(0);
236 remaining = &remaining[total_len..];
237 }
238 total
239}
240
241impl RecordBatch<'_> {
244 pub fn iter(&self) -> RecordIter<'_> {
251 let body: &[u8] = match &self.body {
252 RecordBody::Borrowed(s) => s,
253 RecordBody::Owned(b) => b.as_ref(),
254 };
255 #[allow(clippy::cast_sign_loss)] let count = self.header.records_count.get().max(0) as usize;
257 RecordIter {
258 remaining: body,
259 count,
260 index: 0,
261 }
262 }
263}
264
265impl<'a> IntoIterator for &'a RecordBatch<'_> {
266 type Item = Result<Record<'a>, RecordsError>;
267 type IntoIter = RecordIter<'a>;
268
269 fn into_iter(self) -> Self::IntoIter {
270 self.iter()
271 }
272}
273
274pub struct RecordIter<'a> {
275 remaining: &'a [u8],
276 count: usize,
277 index: usize,
278}
279
280impl<'a> Iterator for RecordIter<'a> {
281 type Item = Result<Record<'a>, RecordsError>;
282
283 fn next(&mut self) -> Option<Self::Item> {
284 if self.index >= self.count {
285 return None;
286 }
287 self.index += 1;
288 Some(parse_one_record(&mut self.remaining))
289 }
290}
291
292#[inline]
293pub(crate) fn parse_one_record<'a>(buf: &mut &'a [u8]) -> Result<Record<'a>, RecordsError> {
294 let body_len =
295 get_varlong(buf).map_err(|e| RecordsError::RecordParse(format!("record length: {e}")))?;
296 let body_len = usize::try_from(body_len).map_err(|_| {
297 RecordsError::RecordParse(format!("record length negative or too large: {body_len}"))
298 })?;
299 if buf.len() < body_len {
300 return Err(RecordsError::BodyTooShort {
301 needed: body_len - buf.len(),
302 });
303 }
304 let (body, rest) = buf.split_at(body_len);
305 *buf = rest;
306 let mut body_cur = body;
307 let r = parse_body(&mut body_cur)?;
308 if !body_cur.is_empty() {
309 return Err(RecordsError::RecordParse(format!(
310 "trailing bytes inside record (left={})",
311 body_cur.len()
312 )));
313 }
314 Ok(r)
315}
316
317fn parse_body<'a>(buf: &mut &'a [u8]) -> Result<Record<'a>, RecordsError> {
318 if buf.is_empty() {
319 return Err(RecordsError::RecordParse("record body empty".into()));
320 }
321 #[allow(clippy::cast_possible_wrap)] let attributes = buf[0] as i8;
323 *buf = &buf[1..];
324 let timestamp_delta =
325 get_varlong(buf).map_err(|e| RecordsError::RecordParse(format!("timestamp_delta: {e}")))?;
326 let offset_delta =
327 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("offset_delta: {e}")))?;
328
329 let key = read_nullable_slice(buf, "key")?;
330 let value = read_nullable_slice(buf, "value")?;
331
332 let header_count =
333 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("header_count: {e}")))?;
334 if header_count < 0 {
335 return Err(RecordsError::RecordParse(format!(
336 "negative header count {header_count}"
337 )));
338 }
339 #[allow(clippy::cast_sign_loss)] let mut headers = Vec::with_capacity((header_count as usize).min(buf.len()));
344 for i in 0..header_count {
345 let key_len = get_varint(buf)
346 .map_err(|e| RecordsError::RecordParse(format!("header[{i}] key length: {e}")))?;
347 if key_len < 0 {
348 return Err(RecordsError::RecordParse(format!(
349 "header[{i}] negative key length"
350 )));
351 }
352 #[allow(clippy::cast_sign_loss)] let n = key_len as usize;
354 if buf.len() < n {
355 return Err(RecordsError::BodyTooShort {
356 needed: n - buf.len(),
357 });
358 }
359 let (key_bytes, rest) = buf.split_at(n);
360 *buf = rest;
361 let key_str = std::str::from_utf8(key_bytes)
362 .map_err(|e| RecordsError::RecordParse(format!("header[{i}] key utf-8: {e}")))?;
363
364 let value = read_nullable_slice(buf, &format!("header[{i}] value"))?;
365 headers.push(RecordHeader {
366 key: key_str,
367 value,
368 });
369 }
370
371 Ok(Record {
372 attributes,
373 timestamp_delta,
374 offset_delta,
375 key,
376 value,
377 headers,
378 })
379}
380
381fn read_nullable_slice<'a>(
382 buf: &mut &'a [u8],
383 label: &str,
384) -> Result<Option<&'a [u8]>, RecordsError> {
385 let len =
386 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("{label} length: {e}")))?;
387 if len < 0 {
388 Ok(None)
389 } else {
390 #[allow(clippy::cast_sign_loss)] let n = len as usize;
392 if buf.len() < n {
393 return Err(RecordsError::BodyTooShort {
394 needed: n - buf.len(),
395 });
396 }
397 let (head, rest) = buf.split_at(n);
398 *buf = rest;
399 Ok(Some(head))
400 }
401}
402
403impl RecordBatch<'_> {
406 pub fn to_owned(&self) -> Result<super::owned::RecordBatch, RecordsError> {
409 let mut records = Vec::new();
410 for r in self {
411 let r = r?;
412 records.push(super::owned::Record {
413 attributes: r.attributes,
414 timestamp_delta: r.timestamp_delta,
415 offset_delta: r.offset_delta,
416 key: r.key.map(Bytes::copy_from_slice),
417 value: r.value.map(Bytes::copy_from_slice),
418 headers: r
419 .headers
420 .into_iter()
421 .map(|h| super::owned::RecordHeader {
422 key: h.key.to_string(),
423 value: h.value.map(Bytes::copy_from_slice),
424 })
425 .collect(),
426 });
427 }
428 Ok(super::owned::RecordBatch {
429 base_offset: self.header.base_offset.get(),
430 partition_leader_epoch: self.header.partition_leader_epoch.get(),
431 attributes: self.attributes(),
432 last_offset_delta: self.header.last_offset_delta.get(),
433 base_timestamp: self.header.base_timestamp.get(),
434 max_timestamp: self.header.max_timestamp.get(),
435 producer_id: self.header.producer_id.get(),
436 producer_epoch: self.header.producer_epoch.get(),
437 base_sequence: self.header.base_sequence.get(),
438 records,
439 })
440 }
441}
442
443impl std::fmt::Debug for RecordBatch<'_> {
449 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
450 match self.to_owned() {
451 Ok(o) => o.fmt(f),
452 Err(e) => write!(f, "RecordBatch(<decode error: {e}>)"),
453 }
454 }
455}
456
457impl Clone for RecordBatch<'_> {
458 fn clone(&self) -> Self {
462 RecordBatch {
463 header: self.header,
464 body: match &self.body {
465 RecordBody::Borrowed(s) => RecordBody::Borrowed(s),
466 RecordBody::Owned(b) => RecordBody::Owned(b.clone()),
467 },
468 }
469 }
470}
471
472impl PartialEq for RecordBatch<'_> {
473 fn eq(&self, other: &Self) -> bool {
474 match (self.to_owned(), other.to_owned()) {
475 (Ok(a), Ok(b)) => a == b,
476 _ => false,
477 }
478 }
479}
480
481impl Eq for RecordBatch<'_> {}
482
483impl crate::Encode for RecordBatch<'_> {
486 fn encode<B: bytes::BufMut>(
487 &self,
488 buf: &mut B,
489 version: i16,
490 ) -> Result<(), crate::ProtocolError> {
491 let owned = self.to_owned().map_err(crate::ProtocolError::from)?;
492 crate::Encode::encode(&owned, buf, version)
493 }
494
495 fn encoded_len(&self, version: i16) -> usize {
496 match self.to_owned() {
497 Ok(o) => crate::Encode::encoded_len(&o, version),
498 Err(_) => 0,
499 }
500 }
501}
502
503#[cfg(test)]
506mod tests {
507 use super::*;
508 use crate::DecodeBorrow;
509 use assert2::assert;
510 use bytes::BytesMut;
511 use crabka_compression::CompressionType;
512
513 fn encode_owned_then_borrow(b: &super::super::owned::RecordBatch) -> Vec<u8> {
514 let mut buf = BytesMut::new();
515 b.encode(&mut buf).unwrap();
516 buf.to_vec()
517 }
518
519 macro_rules! borrowed_roundtrip {
520 ($name:ident, $codec:expr) => {
521 #[test]
522 fn $name() {
523 let mut owned = super::super::owned::RecordBatch::default();
524 owned.attributes = owned.attributes.with_compression($codec);
525 owned.records.push(super::super::owned::Record {
526 key: Some(Bytes::from_static(b"key")),
527 value: Some(Bytes::from_static(b"value")),
528 ..Default::default()
529 });
530
531 let encoded = encode_owned_then_borrow(&owned);
532 let mut cur: &[u8] = &encoded[..];
533 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
534 assert!(cur.is_empty());
535 assert!(borrowed.attributes() == owned.attributes);
536
537 let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
538 assert!(records.len() == 1);
539 assert!(records[0].key == Some(b"key".as_slice()));
540 assert!(records[0].value == Some(b"value".as_slice()));
541
542 let back_owned = borrowed.to_owned().unwrap();
543 assert!(back_owned == owned);
544 }
545 };
546 }
547
548 borrowed_roundtrip!(roundtrip_none, CompressionType::None);
549 borrowed_roundtrip!(roundtrip_gzip, CompressionType::Gzip);
550 borrowed_roundtrip!(roundtrip_snappy, CompressionType::Snappy);
551 borrowed_roundtrip!(roundtrip_lz4, CompressionType::Lz4);
552 borrowed_roundtrip!(roundtrip_zstd, CompressionType::Zstd);
553
554 #[test]
555 fn zero_copy_for_uncompressed() {
556 let mut owned = super::super::owned::RecordBatch::default();
559 owned.records.push(super::super::owned::Record {
560 key: Some(Bytes::from_static(b"k")),
561 value: Some(Bytes::from_static(b"v")),
562 ..Default::default()
563 });
564 let encoded = encode_owned_then_borrow(&owned);
565 let encoded_start = encoded.as_ptr() as usize;
566 let encoded_end = encoded_start + encoded.len();
567
568 let mut cur: &[u8] = &encoded[..];
569 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
570 let records: Vec<_> = borrowed.iter().collect::<Result<_, _>>().unwrap();
571
572 let v_ptr = records[0].value.unwrap().as_ptr() as usize;
573 assert!(
574 v_ptr >= encoded_start && v_ptr < encoded_end,
575 "value slice does not point into the input buffer: \
576 input range [{encoded_start:#x}, {encoded_end:#x}), value ptr {v_ptr:#x}",
577 );
578 }
579
580 #[test]
581 fn validate_one_v2_batch_reads_header_and_len() {
582 let mut owned = super::super::owned::RecordBatch {
583 base_offset: 7,
584 partition_leader_epoch: 3,
585 last_offset_delta: 0,
586 producer_id: 99,
587 producer_epoch: 1,
588 base_sequence: 5,
589 max_timestamp: 1_234,
590 ..Default::default()
591 };
592 owned.records.push(super::super::owned::Record {
593 value: Some(Bytes::from_static(b"payload")),
594 ..Default::default()
595 });
596 let encoded = encode_owned_then_borrow(&owned);
597
598 let v = validate_one_v2_batch(&encoded).unwrap();
599 assert!(v.total_len == encoded.len());
600 assert!(v.header.base_offset.get() == 7);
601 assert!(v.header.partition_leader_epoch.get() == 3);
602 assert!(v.header.producer_id.get() == 99);
603 assert!(v.header.producer_epoch.get() == 1);
604 assert!(v.header.base_sequence.get() == 5);
605 assert!(v.header.max_timestamp.get() == 1_234);
606 assert!(v.header.magic == 2);
607 }
608
609 #[test]
610 fn validate_one_v2_batch_rejects_corrupt_crc() {
611 let owned = super::super::owned::RecordBatch {
612 records: vec![super::super::owned::Record {
613 value: Some(Bytes::from_static(b"x")),
614 ..Default::default()
615 }],
616 ..Default::default()
617 };
618 let mut encoded = encode_owned_then_borrow(&owned);
619 encoded[HEADER_LEN] ^= 0xFF;
621 let err = validate_one_v2_batch(&encoded).unwrap_err();
622 assert!(matches!(err, RecordsError::CrcMismatch { .. }));
623 }
624
625 #[test]
626 fn validate_one_v2_batch_rejects_truncated() {
627 let owned = super::super::owned::RecordBatch {
628 records: vec![super::super::owned::Record {
629 value: Some(Bytes::from_static(b"value")),
630 ..Default::default()
631 }],
632 ..Default::default()
633 };
634 let encoded = encode_owned_then_borrow(&owned);
635 let err = validate_one_v2_batch(&encoded[..encoded.len() - 2]).unwrap_err();
636 assert!(matches!(err, RecordsError::BodyTooShort { .. }));
637 }
638
639 #[test]
640 fn count_records_in_v2_batches_counts_single_and_concatenated_batches() {
641 let one = encode_owned_then_borrow(&super::super::owned::RecordBatch {
642 records: vec![super::super::owned::Record::default()],
643 ..Default::default()
644 });
645 let three = encode_owned_then_borrow(&super::super::owned::RecordBatch {
646 last_offset_delta: 2,
647 records: vec![
648 super::super::owned::Record::default(),
649 super::super::owned::Record::default(),
650 super::super::owned::Record::default(),
651 ],
652 ..Default::default()
653 });
654 let mut both = Vec::with_capacity(one.len() + three.len());
655 both.extend_from_slice(&one);
656 both.extend_from_slice(&three);
657
658 assert!(count_records_in_v2_batches(&one) == 1);
659 assert!(count_records_in_v2_batches(&both) == 4);
660 }
661
662 #[test]
663 fn count_records_in_v2_batches_stops_at_bad_or_non_v2_input() {
664 assert!(count_records_in_v2_batches(&[]) == 0);
665 assert!(count_records_in_v2_batches(&[0u8; HEADER_LEN]) == 0);
666
667 let encoded = encode_owned_then_borrow(&super::super::owned::RecordBatch {
668 records: vec![super::super::owned::Record::default()],
669 ..Default::default()
670 });
671 let mut with_truncated_tail = encoded.clone();
672 with_truncated_tail.extend_from_slice(&encoded[..HEADER_LEN - 1]);
673
674 assert!(count_records_in_v2_batches(&with_truncated_tail) == 1);
675 }
676
677 fn v2_header_only(batch_length: i32, records_count: i32) -> Vec<u8> {
683 let mut buf = vec![0u8; HEADER_LEN];
684 buf[16] = 2; buf[8..12].copy_from_slice(&batch_length.to_be_bytes());
686 buf[57..61].copy_from_slice(&records_count.to_be_bytes());
687 buf
688 }
689
690 #[test]
691 fn count_records_in_v2_batches_accepts_batch_that_exactly_fills_buffer() {
692 let buf = v2_header_only(49, 7);
697 assert!(buf.len() == HEADER_LEN);
698 assert!(count_records_in_v2_batches(&buf) == 7);
699 }
700
701 #[test]
702 fn count_records_in_v2_batches_rejects_batch_overrunning_buffer() {
703 let buf = v2_header_only(100, 9);
708 assert!(count_records_in_v2_batches(&buf) == 0);
709 }
710
711 #[test]
712 fn borrowed_encode_via_trait_roundtrips() {
713 use crate::Encode as _;
714 let owned_in = super::super::owned::RecordBatch {
715 records: vec![super::super::owned::Record {
716 key: Some(Bytes::from_static(b"x")),
717 value: Some(Bytes::from_static(b"y")),
718 ..Default::default()
719 }],
720 ..Default::default()
721 };
722 let bytes_in = encode_owned_then_borrow(&owned_in);
723 let mut cur: &[u8] = &bytes_in[..];
724 let borrowed = RecordBatch::decode_borrow(&mut cur, 0).unwrap();
725
726 let mut out = BytesMut::new();
727 borrowed.encode(&mut out, 0).unwrap();
728 assert!(&out[..] == &bytes_in[..]);
729 }
730}