use tsoracle_core::LeaseRecord;
pub const MAGIC: &[u8; 4] = b"TSOO";
pub const VERSION: u8 = 1;
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum LeaseRecordError {
#[error("lease record too short: {0} bytes")]
TooShort(usize),
#[error("bad magic: {0:?}")]
BadMagic([u8; 4]),
#[error("unsupported lease record version {0}")]
UnsupportedVersion(u8),
#[error("checksum mismatch: stored {stored:#010x}, computed {computed:#010x}")]
ChecksumMismatch { stored: u32, computed: u32 },
#[error("malformed lease record body")]
Malformed,
}
pub fn encode(records: &[LeaseRecord]) -> Vec<u8> {
debug_assert!(
records.len() <= u32::MAX as usize,
"record_count exceeds u32"
);
let mut records = records.to_vec();
records.sort_by_key(|r| r.lease_id);
let holder_bytes: usize = records.iter().map(|r| r.holder.len()).sum();
let mut buf = Vec::with_capacity(4 + 1 + 4 + records.len() * 43 + holder_bytes + 4);
buf.extend_from_slice(MAGIC);
buf.push(VERSION);
buf.extend_from_slice(&(records.len() as u32).to_le_bytes());
for record in &records {
debug_assert!(
record.holder.len() <= u16::MAX as usize,
"holder length exceeds u16"
);
buf.extend_from_slice(&record.lease_id.to_le_bytes());
buf.extend_from_slice(&(record.holder.len() as u16).to_le_bytes());
buf.extend_from_slice(&record.holder);
buf.extend_from_slice(&record.holder_epoch.to_le_bytes());
buf.extend_from_slice(&record.ttl_ms.to_le_bytes());
buf.extend_from_slice(&record.ts_upper_bound.to_le_bytes());
buf.extend_from_slice(&record.expires_at_ms.to_le_bytes());
buf.push(u8::from(record.superseded));
}
let crc = crc32c::crc32c(&buf);
buf.extend_from_slice(&crc.to_le_bytes());
buf
}
pub fn decode(bytes: &[u8]) -> Result<Vec<LeaseRecord>, LeaseRecordError> {
if bytes.len() < 13 {
return Err(LeaseRecordError::TooShort(bytes.len()));
}
if &bytes[0..4] != MAGIC {
return Err(LeaseRecordError::BadMagic([
bytes[0], bytes[1], bytes[2], bytes[3],
]));
}
if bytes[4] != VERSION {
return Err(LeaseRecordError::UnsupportedVersion(bytes[4]));
}
let body = &bytes[..bytes.len() - 4];
let stored = u32::from_le_bytes(
bytes[bytes.len() - 4..]
.try_into()
.map_err(|_| LeaseRecordError::Malformed)?,
);
let computed = crc32c::crc32c(body);
if stored != computed {
return Err(LeaseRecordError::ChecksumMismatch { stored, computed });
}
let record_count = u32::from_le_bytes(
body[5..9]
.try_into()
.map_err(|_| LeaseRecordError::Malformed)?,
);
let mut pos = 9;
let mut records = Vec::with_capacity(record_count as usize);
let mut prev_id = None;
for _ in 0..record_count {
if pos + 10 > body.len() {
return Err(LeaseRecordError::Malformed);
}
let lease_id = u64::from_le_bytes(
body[pos..pos + 8]
.try_into()
.map_err(|_| LeaseRecordError::Malformed)?,
);
pos += 8;
if let Some(prev) = prev_id
&& lease_id <= prev
{
return Err(LeaseRecordError::Malformed);
}
prev_id = Some(lease_id);
let holder_len = u16::from_le_bytes(
body[pos..pos + 2]
.try_into()
.map_err(|_| LeaseRecordError::Malformed)?,
) as usize;
pos += 2;
if pos + holder_len + 33 > body.len() {
return Err(LeaseRecordError::Malformed);
}
let holder = body[pos..pos + holder_len].to_vec();
pos += holder_len;
let holder_epoch = read_u64(body, &mut pos)?;
let ttl_ms = read_u64(body, &mut pos)?;
let ts_upper_bound = read_u64(body, &mut pos)?;
let expires_at_ms = read_u64(body, &mut pos)?;
let superseded = match body.get(pos) {
Some(0) => false,
Some(1) => true,
Some(_) | None => return Err(LeaseRecordError::Malformed),
};
pos += 1;
records.push(LeaseRecord {
lease_id,
holder,
holder_epoch,
ttl_ms,
ts_upper_bound,
expires_at_ms,
superseded,
});
}
if pos != body.len() {
return Err(LeaseRecordError::Malformed);
}
Ok(records)
}
fn read_u64(body: &[u8], pos: &mut usize) -> Result<u64, LeaseRecordError> {
if *pos + 8 > body.len() {
return Err(LeaseRecordError::Malformed);
}
let value = u64::from_le_bytes(
body[*pos..*pos + 8]
.try_into()
.map_err(|_| LeaseRecordError::Malformed)?,
);
*pos += 8;
Ok(value)
}
#[cfg(test)]
mod tests {
use super::*;
fn rec(
lease_id: u64,
holder: &[u8],
holder_epoch: u64,
ttl_ms: u64,
ts_upper_bound: u64,
expires_at_ms: u64,
superseded: bool,
) -> LeaseRecord {
LeaseRecord {
lease_id,
holder: holder.to_vec(),
holder_epoch,
ttl_ms,
ts_upper_bound,
expires_at_ms,
superseded,
}
}
#[test]
fn empty_set_round_trips() {
let records = Vec::new();
assert_eq!(decode(&encode(&records)).unwrap(), records);
}
#[test]
fn two_records_round_trip_preserving_fields() {
let records = vec![
rec(2, b"holder-b", 7, 10_000, 2_000, 12_000, true),
rec(1, b"holder-a", 3, 5_000, 1_000, 6_000, false),
];
let mut expected = records.clone();
expected.sort_by_key(|r| r.lease_id);
assert_eq!(decode(&encode(&records)).unwrap(), expected);
}
#[test]
fn rejects_bad_magic() {
let mut bytes = encode(&[rec(1, b"h", 1, 5_000, 10, 5_010, false)]);
bytes[0] = b'X';
assert!(matches!(decode(&bytes), Err(LeaseRecordError::BadMagic(_))));
}
#[test]
fn rejects_wrong_version() {
let mut bytes = encode(&[rec(1, b"h", 1, 5_000, 10, 5_010, false)]);
bytes[4] = 0;
assert_eq!(decode(&bytes), Err(LeaseRecordError::UnsupportedVersion(0)));
}
#[test]
fn rejects_flipped_crc_byte() {
let mut bytes = encode(&[rec(1, b"h", 1, 5_000, 10, 5_010, false)]);
let last = bytes.len() - 1;
bytes[last] ^= 0xff;
assert!(matches!(
decode(&bytes),
Err(LeaseRecordError::ChecksumMismatch { .. })
));
}
#[test]
fn rejects_truncated_body() {
let mut bytes = encode(&[rec(1, b"h", 1, 5_000, 10, 5_010, false)]);
bytes.truncate(bytes.len() - 2);
assert!(matches!(
decode(&bytes),
Err(LeaseRecordError::ChecksumMismatch { .. }) | Err(LeaseRecordError::Malformed)
));
}
}