use super::error::MessageError;
pub const DOC_FLAG_TOMBSTONE: u8 = 0x01;
pub const DOC_FLAG_ENCRYPTED: u8 = 0x02;
pub const DOC_FLAGS_MASK: u8 = DOC_FLAG_TOMBSTONE | DOC_FLAG_ENCRYPTED;
pub const MAX_COLLECTION_LEN: usize = u8::MAX as usize;
pub const MAX_DOC_ID_LEN: usize = u16::MAX as usize;
pub const MAX_BODY_LEN: usize = u16::MAX as usize;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct DocumentRef<'a> {
pub flags: u8,
pub collection: &'a str,
pub doc_id: &'a str,
pub timestamp_ms: i64,
pub body: &'a [u8],
}
impl<'a> DocumentRef<'a> {
#[inline]
pub fn is_tombstone(&self) -> bool {
self.flags & DOC_FLAG_TOMBSTONE != 0
}
#[inline]
pub fn is_encrypted(&self) -> bool {
self.flags & DOC_FLAG_ENCRYPTED != 0
}
}
#[inline]
pub const fn encoded_len(
collection_len: usize,
doc_id_len: usize,
body_len: usize,
) -> Option<usize> {
let fixed = 1 + 1 + 2 + 8 + 2;
let Some(t) = collection_len.checked_add(doc_id_len) else {
return None;
};
let Some(t) = t.checked_add(body_len) else {
return None;
};
let Some(t) = t.checked_add(fixed) else {
return None;
};
Some(t)
}
pub fn encode(
flags: u8,
collection: &str,
doc_id: &str,
timestamp_ms: i64,
body: &[u8],
buf: &mut [u8],
) -> Result<usize, MessageError> {
if flags & !DOC_FLAGS_MASK != 0 || flags & DOC_FLAG_ENCRYPTED != 0 {
return Err(MessageError::InvalidFlags);
}
let coll_bytes = collection.as_bytes();
let id_bytes = doc_id.as_bytes();
if coll_bytes.is_empty() {
return Err(MessageError::EmptyCollection);
}
if coll_bytes.contains(&0) || id_bytes.contains(&0) {
return Err(MessageError::InvalidUtf8);
}
if flags & DOC_FLAG_TOMBSTONE != 0 && !body.is_empty() {
return Err(MessageError::TombstoneWithBody);
}
if coll_bytes.len() > MAX_COLLECTION_LEN
|| id_bytes.len() > MAX_DOC_ID_LEN
|| body.len() > MAX_BODY_LEN
{
return Err(MessageError::FieldTooLarge);
}
let needed = encoded_len(coll_bytes.len(), id_bytes.len(), body.len())
.ok_or(MessageError::FieldTooLarge)?;
if buf.len() < needed {
return Err(MessageError::BufferTooSmall);
}
let mut o = 0;
buf[o] = flags;
o += 1;
buf[o] = coll_bytes.len() as u8;
o += 1;
buf[o..o + coll_bytes.len()].copy_from_slice(coll_bytes);
o += coll_bytes.len();
buf[o..o + 2].copy_from_slice(&(id_bytes.len() as u16).to_le_bytes());
o += 2;
buf[o..o + id_bytes.len()].copy_from_slice(id_bytes);
o += id_bytes.len();
buf[o..o + 8].copy_from_slice(×tamp_ms.to_le_bytes());
o += 8;
buf[o..o + 2].copy_from_slice(&(body.len() as u16).to_le_bytes());
o += 2;
buf[o..o + body.len()].copy_from_slice(body);
o += body.len();
assert_eq!(o, needed);
Ok(o)
}
pub fn decode(buf: &[u8]) -> Result<DocumentRef<'_>, MessageError> {
let mut o = 0;
if buf.len() < 1 + 1 {
return Err(MessageError::TooShort);
}
let flags = buf[o];
o += 1;
let coll_len = buf[o] as usize;
o += 1;
if buf.len() < o + coll_len {
return Err(MessageError::TruncatedField);
}
let collection =
core::str::from_utf8(&buf[o..o + coll_len]).map_err(|_| MessageError::InvalidUtf8)?;
if collection.is_empty() {
return Err(MessageError::EmptyCollection);
}
if collection.as_bytes().contains(&0) {
return Err(MessageError::InvalidUtf8);
}
o += coll_len;
if buf.len() < o + 2 {
return Err(MessageError::TruncatedField);
}
let id_len = u16::from_le_bytes([buf[o], buf[o + 1]]) as usize;
o += 2;
if buf.len() < o + id_len {
return Err(MessageError::TruncatedField);
}
let doc_id =
core::str::from_utf8(&buf[o..o + id_len]).map_err(|_| MessageError::InvalidUtf8)?;
if doc_id.as_bytes().contains(&0) {
return Err(MessageError::InvalidUtf8);
}
o += id_len;
if buf.len() < o + 8 {
return Err(MessageError::TruncatedField);
}
let mut ts = [0u8; 8];
ts.copy_from_slice(&buf[o..o + 8]);
let timestamp_ms = i64::from_le_bytes(ts);
o += 8;
if buf.len() < o + 2 {
return Err(MessageError::TruncatedField);
}
let body_len = u16::from_le_bytes([buf[o], buf[o + 1]]) as usize;
o += 2;
if buf.len() < o + body_len {
return Err(MessageError::TruncatedField);
}
let body = &buf[o..o + body_len];
Ok(DocumentRef {
flags,
collection,
doc_id,
timestamp_ms,
body,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn roundtrip_full() {
let mut buf = [0u8; 512];
let body = b"opaque-body-bytes";
let n = encode(
0,
"markers",
"uuid-abc-123",
1_700_000_000_000,
body,
&mut buf,
)
.expect("encode");
let view = decode(&buf[..n]).expect("decode");
assert_eq!(view.flags, 0);
assert_eq!(view.collection, "markers");
assert_eq!(view.doc_id, "uuid-abc-123");
assert_eq!(view.timestamp_ms, 1_700_000_000_000);
assert_eq!(view.body, body);
assert!(!view.is_tombstone());
assert!(!view.is_encrypted());
}
#[test]
fn roundtrip_empty_doc_id() {
let mut buf = [0u8; 256];
let n = encode(0, "markers", "", 0, b"x", &mut buf).expect("encode");
let view = decode(&buf[..n]).expect("decode");
assert_eq!(view.doc_id, "");
assert_eq!(view.body, b"x");
}
#[test]
fn roundtrip_tombstone_empty_body() {
let mut buf = [0u8; 64];
let n = encode(
DOC_FLAG_TOMBSTONE,
"platforms",
"ANDROID-x",
42,
&[],
&mut buf,
)
.expect("encode");
let view = decode(&buf[..n]).expect("decode");
assert!(view.is_tombstone());
assert!(!view.is_encrypted());
assert_eq!(view.body.len(), 0);
}
#[test]
fn empty_collection_is_rejected_with_distinct_variant() {
let mut buf = [0u8; 64];
assert_eq!(
encode(0, "", "id", 0, b"x", &mut buf),
Err(MessageError::EmptyCollection),
"encode-side empty-collection must surface EmptyCollection"
);
let mut wire = [0u8; 16];
wire[0] = 0; wire[1] = 0; assert_eq!(
decode(&wire[..4]),
Err(MessageError::EmptyCollection),
"decode-side empty-collection must surface EmptyCollection"
);
}
#[test]
fn truncated_input_returns_error_not_panic() {
let mut buf = [0u8; 256];
let n = encode(0, "markers", "id", 1, b"body", &mut buf).expect("encode");
for trunc in 0..n {
let res = decode(&buf[..trunc]);
assert!(
matches!(
res,
Err(MessageError::TooShort) | Err(MessageError::TruncatedField),
),
"trunc {} should be TooShort/TruncatedField, got {:?}",
trunc,
res
);
}
}
#[test]
fn non_utf8_collection_rejected() {
let mut wire = [0u8; 32];
wire[0] = 0;
wire[1] = 3;
wire[2] = 0xFF;
wire[3] = 0xFE;
wire[4] = 0xFD;
wire[5] = 0; wire[6] = 0; assert_eq!(decode(&wire[..17]), Err(MessageError::InvalidUtf8));
}
#[test]
fn encoded_len_matches_encode() {
let mut buf = [0u8; 1024];
let body = vec![0xAB; 300];
let n = encode(0, "tracks", "id-1", 7, &body, &mut buf).expect("encode");
assert_eq!(
n,
encoded_len("tracks".len(), "id-1".len(), body.len()).expect("no overflow"),
"encoded_len drift",
);
}
#[test]
fn encoded_len_returns_none_on_overflow() {
assert_eq!(encoded_len(usize::MAX, 1, 0), None);
assert_eq!(encoded_len(0, usize::MAX, 1), None);
assert_eq!(encoded_len(usize::MAX - 13, 0, 0), None);
assert_eq!(encoded_len(0, 0, usize::MAX - 13), None);
assert_eq!(
encoded_len(usize::MAX - 14, 0, 0),
Some(usize::MAX - 14 + 14)
);
}
#[test]
fn oversize_fields_are_rejected() {
let mut buf = [0u8; 512];
let big_coll = "x".repeat(MAX_COLLECTION_LEN + 1);
assert_eq!(
encode(0, &big_coll, "id", 0, b"x", &mut buf),
Err(MessageError::FieldTooLarge),
);
let body_too_large = vec![0u8; MAX_BODY_LEN + 1];
let needed = encoded_len("c".len(), "id".len(), body_too_large.len()).expect("no overflow");
let mut huge = vec![0u8; needed];
assert_eq!(
encode(0, "c", "id", 0, &body_too_large, &mut huge),
Err(MessageError::FieldTooLarge),
);
}
#[test]
fn encrypted_flag_is_rejected_today() {
let mut buf = [0u8; 64];
assert_eq!(
encode(DOC_FLAG_ENCRYPTED, "markers", "id", 0, b"x", &mut buf),
Err(MessageError::InvalidFlags),
);
}
#[test]
fn reserved_flag_bits_are_rejected() {
let mut buf = [0u8; 64];
for bit in 2..=7 {
let flags = 1u8 << bit;
assert_eq!(
encode(flags, "markers", "id", 0, b"x", &mut buf),
Err(MessageError::InvalidFlags),
"bit {} should be rejected as reserved",
bit
);
}
}
#[test]
fn tombstone_with_non_empty_body_is_rejected() {
let mut buf = [0u8; 64];
assert_eq!(
encode(
DOC_FLAG_TOMBSTONE,
"tracks",
"id",
0,
b"unwanted body",
&mut buf
),
Err(MessageError::TombstoneWithBody),
);
}
#[test]
fn nul_bytes_in_collection_are_rejected() {
let mut buf = [0u8; 64];
assert_eq!(
encode(0, "mark\0ers", "id", 0, b"x", &mut buf),
Err(MessageError::InvalidUtf8),
);
let mut wire = [0u8; 32];
wire[0] = 0; wire[1] = 8; wire[2..10].copy_from_slice(b"mark\0ers");
wire[10] = 0; wire[11] = 0; assert_eq!(
decode(&wire[..22]),
Err(MessageError::InvalidUtf8),
"decode must reject NUL-bearing collection for symmetry with encode"
);
}
#[test]
fn nul_bytes_in_doc_id_are_rejected() {
let mut buf = [0u8; 64];
assert_eq!(
encode(0, "markers", "id\0bad", 0, b"x", &mut buf),
Err(MessageError::InvalidUtf8),
);
let mut wire = [0u8; 64];
wire[0] = 0; wire[1] = 7; wire[2..9].copy_from_slice(b"markers");
wire[9..11].copy_from_slice(&6u16.to_le_bytes()); wire[11..17].copy_from_slice(b"id\0bad");
assert_eq!(
decode(&wire[..27]),
Err(MessageError::InvalidUtf8),
"decode must reject NUL-bearing doc_id for symmetry with encode"
);
}
#[test]
fn tombstone_with_oversize_body_surfaces_tombstone_error() {
let body_too_large = vec![0u8; MAX_BODY_LEN + 1];
let needed = encoded_len("c".len(), "id".len(), body_too_large.len()).expect("no overflow");
let mut buf = vec![0u8; needed];
assert_eq!(
encode(DOC_FLAG_TOMBSTONE, "c", "id", 0, &body_too_large, &mut buf),
Err(MessageError::TombstoneWithBody),
"tombstone-vs-body invariant must win against size-cap"
);
}
#[test]
fn tombstone_combined_with_encrypted_is_rejected() {
let mut buf = [0u8; 64];
assert_eq!(
encode(
DOC_FLAG_TOMBSTONE | DOC_FLAG_ENCRYPTED,
"markers",
"id",
0,
&[],
&mut buf
),
Err(MessageError::InvalidFlags),
"tombstone+encrypted must surface InvalidFlags (encrypted clause wins)"
);
}
#[test]
fn roundtrip_at_collection_max_length_boundary() {
let max_coll = "x".repeat(MAX_COLLECTION_LEN); let mut buf = vec![0u8; 1024];
let n = encode(0, &max_coll, "id", 7, b"body", &mut buf).expect("encode at boundary");
let view = decode(&buf[..n]).expect("decode at boundary");
assert_eq!(view.collection.len(), MAX_COLLECTION_LEN);
assert_eq!(view.collection, max_coll);
assert_eq!(view.body, b"body");
let too_long = "x".repeat(MAX_COLLECTION_LEN + 1);
let mut buf2 = vec![0u8; 2048];
assert_eq!(
encode(0, &too_long, "id", 0, b"x", &mut buf2),
Err(MessageError::FieldTooLarge),
);
}
#[test]
fn fuzzy_roundtrip_random_inputs() {
let mut state: u32 = 0xCAFEBABE;
let mut next = || {
state = state.wrapping_mul(1664525).wrapping_add(1013904223);
state
};
for _ in 0..200 {
let coll_choices = ["m", "markers", "platforms", "alerts", "company_summaries"];
let coll = coll_choices[(next() as usize) % coll_choices.len()];
let id_len = (next() as usize) % 257;
let mut id = String::with_capacity(id_len);
for _ in 0..id_len {
id.push(char::from(0x21 + ((next() & 0x3F) as u8))); }
let body_len = (next() as usize) % 1024;
let mut body = vec![0u8; body_len];
for byte in body.iter_mut() {
*byte = next() as u8;
}
let ts = next() as i64;
let needed = encoded_len(coll.len(), id.len(), body.len()).expect("no overflow");
let mut buf = vec![0u8; needed];
let n = encode(0, coll, &id, ts, &body, &mut buf).expect("encode random input");
assert_eq!(n, needed);
let view = decode(&buf).expect("decode random input");
assert_eq!(view.collection, coll);
assert_eq!(view.doc_id, id);
assert_eq!(view.timestamp_ms, ts);
assert_eq!(view.body, body.as_slice());
}
}
}