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