use crate::{
durability::Durability,
error::DecodeError,
payload::Payload,
version::{MAX_PAYLOAD_HARD_CAP, WIRE_VERSION},
};
const MAGIC: [u8; 4] = *b"WEIR";
pub const HEADER_LEN: usize = 16;
const HEADER_CRC_COVERAGE: usize = 12;
pub const MIN_FRAME_LEN: usize = HEADER_LEN + 4;
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageType {
Push = 0x01,
Ack = 0x02,
Nack = 0x03,
HealthCheck = 0x04,
HealthCheckResponse = 0x05,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UnknownMessageType(pub u8);
impl std::fmt::Display for UnknownMessageType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "unknown message type byte: {:#04x}", self.0)
}
}
impl std::error::Error for UnknownMessageType {}
impl TryFrom<u8> for MessageType {
type Error = UnknownMessageType;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
0x01 => Ok(MessageType::Push),
0x02 => Ok(MessageType::Ack),
0x03 => Ok(MessageType::Nack),
0x04 => Ok(MessageType::HealthCheck),
0x05 => Ok(MessageType::HealthCheckResponse),
v => Err(UnknownMessageType(v)),
}
}
}
impl From<MessageType> for u8 {
fn from(m: MessageType) -> u8 {
m as u8
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
version: u8,
message_type: MessageType,
durability: Durability,
flags: u8,
payload_len: u32,
}
impl Header {
#[must_use]
pub fn new(message_type: MessageType, durability: Durability, flags: u8) -> Self {
Self {
version: WIRE_VERSION,
message_type,
durability,
flags,
payload_len: 0,
}
}
pub fn version(&self) -> u8 {
self.version
}
pub fn message_type(&self) -> MessageType {
self.message_type
}
pub fn durability(&self) -> Durability {
self.durability
}
pub fn flags(&self) -> u8 {
self.flags
}
pub fn payload_len(&self) -> u32 {
self.payload_len
}
#[must_use]
pub fn encode(&self) -> [u8; HEADER_LEN] {
let mut buf = [0u8; HEADER_LEN];
buf[0..4].copy_from_slice(&MAGIC);
buf[4] = self.version;
buf[5] = u8::from(self.message_type);
buf[6] = u8::from(self.durability);
buf[7] = self.flags;
buf[8..12].copy_from_slice(&self.payload_len.to_le_bytes());
let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
buf
}
pub fn decode(buf: &[u8; HEADER_LEN]) -> Result<Self, DecodeError> {
if buf[0..4] != MAGIC {
return Err(DecodeError::BadMagic);
}
let version = buf[4];
if version != WIRE_VERSION {
return Err(DecodeError::VersionMismatch {
supported: WIRE_VERSION,
received: version,
});
}
let expected_crc = u32::from_le_bytes(buf[12..16].try_into().unwrap());
let computed_crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
if expected_crc != computed_crc {
return Err(DecodeError::HeaderCrcMismatch {
expected: expected_crc,
computed: computed_crc,
});
}
let message_type =
MessageType::try_from(buf[5]).map_err(|e| DecodeError::UnknownMessageType(e.0))?;
let durability =
Durability::try_from(buf[6]).map_err(|e| DecodeError::UnknownDurability(e.0))?;
let flags = buf[7];
if flags != 0 {
return Err(DecodeError::ReservedFlagsSet { flags });
}
let payload_len = u32::from_le_bytes(buf[8..12].try_into().unwrap());
Ok(Self {
version,
message_type,
durability,
flags,
payload_len,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Envelope {
header: Header,
payload: Payload,
}
impl Envelope {
#[must_use]
pub fn new(header: Header, payload: impl Into<Payload>) -> Self {
let payload = payload.into();
let payload_len = u32::try_from(payload.len()).unwrap_or(u32::MAX);
let header = Header {
payload_len,
..header
};
Self { header, payload }
}
pub fn header(&self) -> Header {
self.header
}
pub fn payload(&self) -> &Payload {
&self.payload
}
#[must_use]
pub fn encode(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(HEADER_LEN + self.payload.len() + 4);
out.extend_from_slice(&self.header.encode());
out.extend_from_slice(&self.payload);
let payload_crc = crc32fast::hash(&self.payload);
out.extend_from_slice(&payload_crc.to_le_bytes());
out
}
pub fn decode(buf: &[u8]) -> Result<Self, DecodeError> {
if buf.len() < HEADER_LEN {
return Err(DecodeError::TruncatedFrame);
}
let header_bytes: &[u8; HEADER_LEN] = buf[..HEADER_LEN]
.try_into()
.expect("slice is exactly HEADER_LEN bytes, checked above");
let header = Header::decode(header_bytes)?;
let payload_len = header.payload_len as usize;
if payload_len > MAX_PAYLOAD_HARD_CAP {
return Err(DecodeError::PayloadTooLarge {
len: payload_len,
cap: MAX_PAYLOAD_HARD_CAP,
});
}
let frame_len = HEADER_LEN + payload_len + 4;
if buf.len() < frame_len {
return Err(DecodeError::TruncatedFrame);
}
if buf.len() > frame_len {
return Err(DecodeError::TrailingBytes {
extra: buf.len() - frame_len,
});
}
let payload = Payload::copy_from_slice(&buf[HEADER_LEN..HEADER_LEN + payload_len]);
let expected_crc =
u32::from_le_bytes(buf[HEADER_LEN + payload_len..frame_len].try_into().unwrap());
let computed_crc = crc32fast::hash(&payload);
if expected_crc != computed_crc {
return Err(DecodeError::PayloadCrcMismatch {
expected: expected_crc,
computed: computed_crc,
});
}
Ok(Self { header, payload })
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::version::WIRE_VERSION;
fn push_header() -> Header {
Header::new(MessageType::Push, Durability::Durable, 0)
}
#[test]
fn header_encode_decode_round_trip() {
let h = push_header();
let encoded = h.encode();
let decoded = Header::decode(&encoded).unwrap();
assert_eq!(h, decoded);
}
#[test]
fn header_encode_is_16_bytes() {
let encoded = push_header().encode();
assert_eq!(encoded.len(), HEADER_LEN);
}
#[test]
fn header_decode_rejects_bad_magic() {
let mut buf = push_header().encode();
buf[0] = 0x00;
assert_eq!(Header::decode(&buf), Err(DecodeError::BadMagic));
}
#[test]
fn header_decode_rejects_future_version() {
let mut buf = push_header().encode();
buf[4] = WIRE_VERSION + 1;
let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Header::decode(&buf),
Err(DecodeError::VersionMismatch {
supported: WIRE_VERSION,
received: WIRE_VERSION + 1,
})
);
}
#[test]
fn header_decode_rejects_older_version() {
let mut buf = push_header().encode();
buf[4] = WIRE_VERSION - 1;
let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Header::decode(&buf),
Err(DecodeError::VersionMismatch {
supported: WIRE_VERSION,
received: WIRE_VERSION - 1,
})
);
}
#[test]
fn header_decode_version_checked_before_crc() {
let mut buf = push_header().encode();
buf[4] = WIRE_VERSION + 1;
let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
match Header::decode(&buf) {
Err(DecodeError::VersionMismatch { .. }) => {}
other => panic!("expected VersionMismatch, got {other:?}"),
}
}
#[test]
fn header_decode_rejects_corrupted_crc() {
let mut buf = push_header().encode();
buf[12] ^= 0xff;
assert!(matches!(
Header::decode(&buf),
Err(DecodeError::HeaderCrcMismatch { .. })
));
}
#[test]
fn header_decode_rejects_unknown_message_type() {
let mut buf = push_header().encode();
buf[5] = 0xff;
let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Header::decode(&buf),
Err(DecodeError::UnknownMessageType(0xff))
);
}
#[test]
fn header_decode_rejects_unknown_durability() {
let mut buf = push_header().encode();
buf[6] = 0xff;
let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Header::decode(&buf),
Err(DecodeError::UnknownDurability(0xff))
);
}
#[test]
fn header_decode_rejects_nonzero_flags() {
let buf = Header::new(MessageType::Push, Durability::Durable, 0x01).encode();
assert_eq!(
Header::decode(&buf),
Err(DecodeError::ReservedFlagsSet { flags: 0x01 })
);
let buf = Header::new(MessageType::Push, Durability::Durable, 0xff).encode();
assert_eq!(
Header::decode(&buf),
Err(DecodeError::ReservedFlagsSet { flags: 0xff })
);
}
#[test]
fn header_decode_flags_checked_after_durability() {
let mut buf = Header::new(MessageType::Push, Durability::Durable, 0x01).encode();
buf[6] = 0xff; let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
buf[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Header::decode(&buf),
Err(DecodeError::UnknownDurability(0xff))
);
}
#[test]
fn message_type_try_from_accepts_all_known_values() {
assert_eq!(MessageType::try_from(0x01).unwrap(), MessageType::Push);
assert_eq!(MessageType::try_from(0x02).unwrap(), MessageType::Ack);
assert_eq!(MessageType::try_from(0x03).unwrap(), MessageType::Nack);
assert_eq!(
MessageType::try_from(0x04).unwrap(),
MessageType::HealthCheck
);
assert_eq!(
MessageType::try_from(0x05).unwrap(),
MessageType::HealthCheckResponse
);
}
#[test]
fn message_type_try_from_rejects_unknown() {
assert!(MessageType::try_from(0x00).is_err());
assert!(MessageType::try_from(0x06).is_err());
assert!(MessageType::try_from(0xff).is_err());
}
#[test]
fn message_type_from_u8_round_trips() {
for mt in [
MessageType::Push,
MessageType::Ack,
MessageType::Nack,
MessageType::HealthCheck,
MessageType::HealthCheckResponse,
] {
assert_eq!(MessageType::try_from(u8::from(mt)).unwrap(), mt);
}
}
#[test]
fn message_type_repr_values_match_wire() {
assert_eq!(MessageType::Push as u8, 0x01);
assert_eq!(MessageType::Ack as u8, 0x02);
assert_eq!(MessageType::Nack as u8, 0x03);
assert_eq!(MessageType::HealthCheck as u8, 0x04);
assert_eq!(MessageType::HealthCheckResponse as u8, 0x05);
}
fn round_trip(msg_type: MessageType) {
let payload = b"hello weir".to_vec();
let header = Header::new(msg_type, Durability::Durable, 0);
let env = Envelope::new(header, payload);
let encoded = env.encode();
let decoded = Envelope::decode(&encoded).unwrap();
assert_eq!(env, decoded);
}
#[test]
fn envelope_round_trip_push() {
round_trip(MessageType::Push);
}
#[test]
fn envelope_round_trip_ack() {
round_trip(MessageType::Ack);
}
#[test]
fn envelope_round_trip_nack() {
round_trip(MessageType::Nack);
}
#[test]
fn envelope_round_trip_health_check() {
round_trip(MessageType::HealthCheck);
}
#[test]
fn envelope_round_trip_health_check_response() {
round_trip(MessageType::HealthCheckResponse);
}
#[test]
fn envelope_round_trip_empty_payload() {
let header = push_header();
let env = Envelope::new(header, vec![]);
let encoded = env.encode();
let decoded = Envelope::decode(&encoded).unwrap();
assert_eq!(env, decoded);
}
#[test]
fn envelope_new_payload_len_equals_payload_length() {
for len in [0usize, 1, 7, 256, 65_536, 1_000_000] {
let env = Envelope::new(push_header(), vec![0xAB; len]);
assert_eq!(
env.header().payload_len() as usize,
len,
"payload_len must equal the actual payload length for len={len}"
);
}
}
#[test]
fn envelope_decode_rejects_truncated_header() {
let env = Envelope::new(push_header(), b"data".to_vec());
let encoded = env.encode();
assert_eq!(
Envelope::decode(&encoded[..HEADER_LEN - 1]),
Err(DecodeError::TruncatedFrame)
);
}
#[test]
fn envelope_decode_rejects_truncated_payload() {
let payload = b"hello".to_vec();
let env = Envelope::new(push_header(), payload);
let encoded = env.encode();
assert_eq!(
Envelope::decode(&encoded[..encoded.len() - 1]),
Err(DecodeError::TruncatedFrame)
);
}
#[test]
fn envelope_decode_rejects_header_only() {
let env = Envelope::new(push_header(), b"data".to_vec());
let encoded = env.encode();
assert_eq!(
Envelope::decode(&encoded[..HEADER_LEN]),
Err(DecodeError::TruncatedFrame)
);
}
#[test]
fn envelope_decode_rejects_corrupted_payload_crc() {
let payload = b"hello".to_vec();
let env = Envelope::new(push_header(), payload);
let mut encoded = env.encode();
let last = encoded.len() - 1;
encoded[last] ^= 0xff;
assert!(matches!(
Envelope::decode(&encoded),
Err(DecodeError::PayloadCrcMismatch { .. })
));
}
#[test]
fn envelope_decode_rejects_trailing_bytes() {
let env = Envelope::new(push_header(), b"hello".to_vec());
let mut encoded = env.encode();
encoded.extend_from_slice(b"garbage"); assert_eq!(
Envelope::decode(&encoded),
Err(DecodeError::TrailingBytes { extra: 7 })
);
}
#[test]
fn envelope_decode_accepts_exact_frame() {
let env = Envelope::new(push_header(), b"hello".to_vec());
let encoded = env.encode();
assert_eq!(Envelope::decode(&encoded).unwrap(), env);
}
#[test]
fn envelope_decode_trailing_bytes_checked_after_payload_too_large() {
let oversized_len = (MAX_PAYLOAD_HARD_CAP + 1) as u32;
let mut header_bytes = push_header().encode();
header_bytes[8..12].copy_from_slice(&oversized_len.to_le_bytes());
let crc = crc32fast::hash(&header_bytes[..HEADER_CRC_COVERAGE]);
header_bytes[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Envelope::decode(&header_bytes),
Err(DecodeError::PayloadTooLarge {
len: MAX_PAYLOAD_HARD_CAP + 1,
cap: MAX_PAYLOAD_HARD_CAP,
})
);
}
#[test]
fn envelope_decode_rejects_oversized_payload_len() {
let oversized_len = (MAX_PAYLOAD_HARD_CAP + 1) as u32;
let mut header_bytes = Header::new(MessageType::Push, Durability::Durable, 0).encode();
header_bytes[8..12].copy_from_slice(&oversized_len.to_le_bytes());
let crc = crc32fast::hash(&header_bytes[..HEADER_CRC_COVERAGE]);
header_bytes[12..16].copy_from_slice(&crc.to_le_bytes());
assert_eq!(
Envelope::decode(&header_bytes),
Err(DecodeError::PayloadTooLarge {
len: MAX_PAYLOAD_HARD_CAP + 1,
cap: MAX_PAYLOAD_HARD_CAP,
})
);
}
}