use crate::version::MAX_PAYLOAD_HARD_CAP;
pub const BATCH_VERSION: u8 = 1;
pub const ACK_BATCH_VERSION: u8 = 1;
pub const BATCH_HEADER_LEN: usize = 1 + 2;
pub const ACK_BATCH_HEADER_LEN: usize = 1 + 2;
pub const MAX_BATCH_RECORDS_HARD_CAP: usize = 2048;
pub const MAX_ACK_BATCH_PAYLOAD_LEN: usize =
ACK_BATCH_HEADER_LEN + MAX_BATCH_RECORDS_HARD_CAP.div_ceil(8);
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum BatchError {
Truncated {
len: usize,
need: usize,
},
UnsupportedVersion {
version: u8,
},
EmptyBatch,
TooManyRecords {
count: usize,
cap: usize,
},
EmptyRecord {
index: usize,
},
RecordTooLarge {
index: usize,
len: usize,
cap: usize,
},
TruncatedRecord {
index: usize,
},
LengthMismatch {
declared: usize,
actual: usize,
},
PaddingNotZero {
byte: u8,
},
}
impl core::fmt::Display for BatchError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Truncated { len, need } => {
write!(f, "batch body truncated: {len} bytes, need at least {need}")
}
Self::UnsupportedVersion { version } => {
write!(f, "unsupported batch version {version:#04x}")
}
Self::EmptyBatch => write!(f, "batch declared zero records"),
Self::TooManyRecords { count, cap } => {
write!(f, "batch declared {count} records, cap is {cap}")
}
Self::EmptyRecord { index } => {
write!(f, "record {index} is zero-length")
}
Self::RecordTooLarge { index, len, cap } => {
write!(f, "record {index} is {len} bytes, cap is {cap}")
}
Self::TruncatedRecord { index } => {
write!(f, "record {index} runs past the end of the body")
}
Self::LengthMismatch { declared, actual } => {
write!(f, "batch declared {declared}, found {actual}")
}
Self::PaddingNotZero { byte } => {
write!(f, "padding bits set in final bitmap byte {byte:#04x}")
}
}
}
}
impl core::error::Error for BatchError {}
pub fn encode_batch_body(records: &[&[u8]]) -> Vec<u8> {
let total: usize = BATCH_HEADER_LEN + records.iter().map(|r| 4 + r.len()).sum::<usize>();
let mut out = Vec::with_capacity(total);
out.push(BATCH_VERSION);
out.extend_from_slice(&(records.len() as u16).to_le_bytes());
for r in records {
out.extend_from_slice(&(r.len() as u32).to_le_bytes());
out.extend_from_slice(r);
}
out
}
pub fn decode_batch_body(
body: &[u8],
max_records: usize,
max_record_len: usize,
) -> Result<Vec<&[u8]>, BatchError> {
if body.len() < BATCH_HEADER_LEN {
return Err(BatchError::Truncated {
len: body.len(),
need: BATCH_HEADER_LEN,
});
}
let version = body[0];
if version != BATCH_VERSION {
return Err(BatchError::UnsupportedVersion { version });
}
let declared = u16::from_le_bytes([body[1], body[2]]) as usize;
if declared == 0 {
return Err(BatchError::EmptyBatch);
}
let cap = max_records.min(MAX_BATCH_RECORDS_HARD_CAP);
if declared > cap {
return Err(BatchError::TooManyRecords {
count: declared,
cap,
});
}
let mut out: Vec<&[u8]> = Vec::with_capacity(declared);
let mut cursor = BATCH_HEADER_LEN;
let record_cap = max_record_len.min(MAX_PAYLOAD_HARD_CAP);
while cursor < body.len() {
let index = out.len();
if cursor + 4 > body.len() {
return Err(BatchError::TruncatedRecord { index });
}
let len = u32::from_le_bytes([
body[cursor],
body[cursor + 1],
body[cursor + 2],
body[cursor + 3],
]) as usize;
cursor += 4;
if len == 0 {
return Err(BatchError::EmptyRecord { index });
}
if len > record_cap {
return Err(BatchError::RecordTooLarge {
index,
len,
cap: record_cap,
});
}
if cursor + len > body.len() {
return Err(BatchError::TruncatedRecord { index });
}
if index == declared {
return Err(BatchError::LengthMismatch {
declared,
actual: declared + 1,
});
}
out.push(&body[cursor..cursor + len]);
cursor += len;
}
if out.len() != declared {
return Err(BatchError::LengthMismatch {
declared,
actual: out.len(),
});
}
if cursor != body.len() {
return Err(BatchError::LengthMismatch {
declared: cursor,
actual: body.len(),
});
}
Ok(out)
}
#[inline]
#[must_use]
pub const fn bitmap_bit(index: usize) -> (usize, u8) {
(index / 8, 1u8 << (index % 8))
}
pub fn encode_ack_batch(accepted: &[bool]) -> Vec<u8> {
let n = accepted.len();
let mut out = Vec::with_capacity(ACK_BATCH_HEADER_LEN + n.div_ceil(8));
out.push(ACK_BATCH_VERSION);
out.extend_from_slice(&(n as u16).to_le_bytes());
out.resize(ACK_BATCH_HEADER_LEN + n.div_ceil(8), 0);
for (i, &ok) in accepted.iter().enumerate() {
if ok {
let (byte, mask) = bitmap_bit(i);
out[ACK_BATCH_HEADER_LEN + byte] |= mask;
}
}
out
}
pub fn decode_ack_batch(payload: &[u8], expected: usize) -> Result<Vec<bool>, BatchError> {
if payload.len() < ACK_BATCH_HEADER_LEN {
return Err(BatchError::Truncated {
len: payload.len(),
need: ACK_BATCH_HEADER_LEN,
});
}
let version = payload[0];
if version != ACK_BATCH_VERSION {
return Err(BatchError::UnsupportedVersion { version });
}
let declared = u16::from_le_bytes([payload[1], payload[2]]) as usize;
if declared == 0 {
return Err(BatchError::EmptyBatch);
}
if declared != expected {
return Err(BatchError::LengthMismatch {
declared,
actual: expected,
});
}
let want = ACK_BATCH_HEADER_LEN + declared.div_ceil(8);
if payload.len() != want {
return Err(BatchError::LengthMismatch {
declared: want,
actual: payload.len(),
});
}
let used_in_last = declared % 8;
if used_in_last != 0 {
let last = payload[payload.len() - 1];
let padding_mask = !((1u8 << used_in_last) - 1);
if last & padding_mask != 0 {
return Err(BatchError::PaddingNotZero { byte: last });
}
}
let bits = &payload[ACK_BATCH_HEADER_LEN..];
Ok((0..declared)
.map(|i| {
let (byte, mask) = bitmap_bit(i);
bits[byte] & mask != 0
})
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
const CAP: usize = 1024;
const REC: usize = 16 * 1024;
#[test]
fn a_batch_body_round_trips() {
let records: Vec<&[u8]> = vec![b"alpha", b"beta", b"gamma"];
let body = encode_batch_body(&records);
assert_eq!(decode_batch_body(&body, CAP, REC).unwrap(), records);
}
#[test]
fn a_single_record_batch_is_legal() {
let body = encode_batch_body(&[b"only"]);
assert_eq!(
decode_batch_body(&body, CAP, REC).unwrap(),
vec![&b"only"[..]]
);
}
#[test]
fn an_unknown_body_version_is_rejected_not_prefix_parsed() {
let mut body = encode_batch_body(&[b"x"]);
body[0] = 0x02;
assert_eq!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::UnsupportedVersion { version: 0x02 })
);
}
#[test]
fn a_zero_record_batch_is_rejected() {
let body = encode_batch_body(&[]);
assert_eq!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::EmptyBatch)
);
}
#[test]
fn a_count_over_the_cap_is_rejected_before_any_reservation() {
let mut body = encode_batch_body(&[b"x"]);
body[1..3].copy_from_slice(&u16::MAX.to_le_bytes());
assert_eq!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::TooManyRecords {
count: 65_535,
cap: CAP
})
);
}
#[test]
fn the_hard_cap_binds_even_when_the_caller_asks_for_more() {
let mut body = encode_batch_body(&[b"x"]);
body[1..3].copy_from_slice(&u16::MAX.to_le_bytes());
match decode_batch_body(&body, usize::MAX, REC) {
Err(BatchError::TooManyRecords { cap, .. }) => {
assert_eq!(cap, MAX_BATCH_RECORDS_HARD_CAP);
}
other => panic!("expected the hard cap to bind, got {other:?}"),
}
}
#[test]
fn an_empty_record_inside_a_batch_is_rejected() {
let body = encode_batch_body(&[b"ok", b"", b"also ok"]);
assert_eq!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::EmptyRecord { index: 1 })
);
}
#[test]
fn a_record_over_the_per_record_cap_is_rejected() {
let big = vec![b'x'; 64];
let body = encode_batch_body(&[b"small", &big]);
assert_eq!(
decode_batch_body(&body, CAP, 32),
Err(BatchError::RecordTooLarge {
index: 1,
len: 64,
cap: 32
})
);
}
#[test]
fn a_truncated_record_is_rejected() {
let mut body = encode_batch_body(&[b"alpha", b"beta"]);
body.truncate(body.len() - 2);
assert!(matches!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::TruncatedRecord { .. })
));
}
#[test]
fn trailing_bytes_after_the_last_record_are_rejected() {
let mut body = encode_batch_body(&[b"alpha"]);
body.push(0xAA);
assert!(matches!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::TruncatedRecord { .. } | BatchError::LengthMismatch { .. })
));
}
#[test]
fn a_count_that_disagrees_with_the_body_is_rejected_either_way() {
let mut body = encode_batch_body(&[b"a", b"b", b"c"]);
body[1..3].copy_from_slice(&2u16.to_le_bytes());
assert!(matches!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::LengthMismatch { .. })
));
let mut body = encode_batch_body(&[b"a", b"b"]);
body[1..3].copy_from_slice(&3u16.to_le_bytes());
assert_eq!(
decode_batch_body(&body, CAP, REC),
Err(BatchError::LengthMismatch {
declared: 3,
actual: 2
})
);
}
#[test]
fn an_ack_bitmap_round_trips() {
for n in [1usize, 7, 8, 9, 63, 64, 65, 255] {
let accepted: Vec<bool> = (0..n).map(|i| i % 3 != 0).collect();
let payload = encode_ack_batch(&accepted);
assert_eq!(
decode_ack_batch(&payload, n).unwrap(),
accepted,
"round trip failed at N={n}"
);
}
}
#[test]
fn the_bitmap_is_lsb_first_and_the_bytes_say_so() {
let mut accepted = vec![false; 9];
accepted[0] = true;
accepted[8] = true;
let payload = encode_ack_batch(&accepted);
assert_eq!(payload[0], ACK_BATCH_VERSION);
assert_eq!(u16::from_le_bytes([payload[1], payload[2]]), 9);
assert_eq!(payload.len(), ACK_BATCH_HEADER_LEN + 2);
assert_eq!(
payload[ACK_BATCH_HEADER_LEN], 0b0000_0001,
"record 0 must be the LEAST significant bit of the first byte"
);
assert_eq!(
payload[ACK_BATCH_HEADER_LEN + 1],
0b0000_0001,
"record 8 must be the least significant bit of the SECOND byte"
);
assert_eq!(decode_ack_batch(&payload, 9).unwrap(), accepted);
}
#[test]
fn a_set_padding_bit_is_rejected_rather_than_ignored() {
let accepted = vec![true; 9];
let mut payload = encode_ack_batch(&accepted);
let last = payload.len() - 1;
payload[last] |= 0b1000_0000;
assert!(matches!(
decode_ack_batch(&payload, 9),
Err(BatchError::PaddingNotZero { .. })
));
}
#[test]
fn an_ack_for_a_different_batch_is_rejected() {
let accepted = vec![true; 1024];
let payload = encode_ack_batch(&accepted);
assert!(decode_ack_batch(&payload, 1024).is_ok());
assert_eq!(
decode_ack_batch(&payload, 1017),
Err(BatchError::LengthMismatch {
declared: 1024,
actual: 1017
})
);
}
#[test]
fn the_largest_ack_batch_fits_the_existing_response_bound() {
let accepted = vec![true; MAX_BATCH_RECORDS_HARD_CAP];
let payload = encode_ack_batch(&accepted);
assert_eq!(payload.len(), MAX_ACK_BATCH_PAYLOAD_LEN);
assert!(
payload.len() <= crate::MAX_TRACKED_ACK_PAYLOAD_LEN,
"the largest AckBatch ({}) exceeds the existing response cap ({}), \
which would force every client to grow its bound",
payload.len(),
crate::MAX_TRACKED_ACK_PAYLOAD_LEN
);
}
}