1use crate::{
6 durability::Durability,
7 error::DecodeError,
8 payload::Payload,
9 version::{MAX_PAYLOAD_HARD_CAP, WIRE_VERSION},
10};
11
12const MAGIC: [u8; 4] = *b"WEIR";
14
15pub const HEADER_LEN: usize = 16;
27
28const HEADER_CRC_COVERAGE: usize = 12;
30
31pub const MIN_FRAME_LEN: usize = HEADER_LEN + 4;
33
34#[repr(u8)]
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum MessageType {
38 Push = 0x01,
40 Ack = 0x02,
42 Nack = 0x03,
44 HealthCheck = 0x04,
46 HealthCheckResponse = 0x05,
48}
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct UnknownMessageType(pub u8);
53
54impl std::fmt::Display for UnknownMessageType {
55 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56 write!(f, "unknown message type byte: {:#04x}", self.0)
57 }
58}
59
60impl std::error::Error for UnknownMessageType {}
61
62impl TryFrom<u8> for MessageType {
63 type Error = UnknownMessageType;
64
65 fn try_from(value: u8) -> Result<Self, Self::Error> {
66 match value {
67 0x01 => Ok(MessageType::Push),
68 0x02 => Ok(MessageType::Ack),
69 0x03 => Ok(MessageType::Nack),
70 0x04 => Ok(MessageType::HealthCheck),
71 0x05 => Ok(MessageType::HealthCheckResponse),
72 v => Err(UnknownMessageType(v)),
73 }
74 }
75}
76
77impl From<MessageType> for u8 {
78 fn from(m: MessageType) -> u8 {
80 m as u8
81 }
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub struct Header {
95 version: u8,
96 message_type: MessageType,
97 durability: Durability,
98 flags: u8,
100 payload_len: u32,
101}
102
103impl Header {
104 #[must_use]
119 pub fn new(message_type: MessageType, durability: Durability, flags: u8) -> Self {
120 Self {
121 version: WIRE_VERSION,
122 message_type,
123 durability,
124 flags,
125 payload_len: 0,
126 }
127 }
128
129 pub fn version(&self) -> u8 {
132 self.version
133 }
134
135 pub fn message_type(&self) -> MessageType {
137 self.message_type
138 }
139
140 pub fn durability(&self) -> Durability {
142 self.durability
143 }
144
145 pub fn flags(&self) -> u8 {
147 self.flags
148 }
149
150 pub fn payload_len(&self) -> u32 {
154 self.payload_len
155 }
156
157 #[must_use]
160 pub fn encode(&self) -> [u8; HEADER_LEN] {
161 let mut buf = [0u8; HEADER_LEN];
162 buf[0..4].copy_from_slice(&MAGIC);
163 buf[4] = self.version;
164 buf[5] = u8::from(self.message_type);
165 buf[6] = u8::from(self.durability);
166 buf[7] = self.flags;
167 buf[8..12].copy_from_slice(&self.payload_len.to_le_bytes());
168 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
169 buf[12..16].copy_from_slice(&crc.to_le_bytes());
170 buf
171 }
172
173 pub fn decode(buf: &[u8; HEADER_LEN]) -> Result<Self, DecodeError> {
183 if buf[0..4] != MAGIC {
184 return Err(DecodeError::BadMagic);
185 }
186
187 let version = buf[4];
188 if version != WIRE_VERSION {
189 return Err(DecodeError::VersionMismatch {
190 supported: WIRE_VERSION,
191 received: version,
192 });
193 }
194
195 let expected_crc = u32::from_le_bytes(buf[12..16].try_into().unwrap());
196 let computed_crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
197 if expected_crc != computed_crc {
198 return Err(DecodeError::HeaderCrcMismatch {
199 expected: expected_crc,
200 computed: computed_crc,
201 });
202 }
203
204 let message_type =
205 MessageType::try_from(buf[5]).map_err(|e| DecodeError::UnknownMessageType(e.0))?;
206 let durability =
207 Durability::try_from(buf[6]).map_err(|e| DecodeError::UnknownDurability(e.0))?;
208 let flags = buf[7];
214 if flags != 0 {
215 return Err(DecodeError::ReservedFlagsSet { flags });
216 }
217 let payload_len = u32::from_le_bytes(buf[8..12].try_into().unwrap());
218
219 Ok(Self {
220 version,
221 message_type,
222 durability,
223 flags,
224 payload_len,
225 })
226 }
227}
228
229#[derive(Debug, Clone, PartialEq, Eq)]
236pub struct Envelope {
237 header: Header,
238 payload: Payload,
239}
240
241impl Envelope {
242 #[must_use]
246 pub fn new(header: Header, payload: impl Into<Payload>) -> Self {
247 let payload = payload.into();
248 let payload_len = u32::try_from(payload.len()).unwrap_or(u32::MAX);
254 let header = Header {
255 payload_len,
256 ..header
257 };
258 Self { header, payload }
259 }
260
261 pub fn header(&self) -> Header {
263 self.header
264 }
265
266 pub fn payload(&self) -> &Payload {
268 &self.payload
269 }
270
271 #[must_use]
273 pub fn encode(&self) -> Vec<u8> {
274 let mut out = Vec::with_capacity(HEADER_LEN + self.payload.len() + 4);
275 out.extend_from_slice(&self.header.encode());
276 out.extend_from_slice(&self.payload);
277 let payload_crc = crc32fast::hash(&self.payload);
278 out.extend_from_slice(&payload_crc.to_le_bytes());
279 out
280 }
281
282 pub fn decode(buf: &[u8]) -> Result<Self, DecodeError> {
296 if buf.len() < HEADER_LEN {
297 return Err(DecodeError::TruncatedFrame);
298 }
299
300 let header_bytes: &[u8; HEADER_LEN] = buf[..HEADER_LEN]
301 .try_into()
302 .expect("slice is exactly HEADER_LEN bytes, checked above");
303 let header = Header::decode(header_bytes)?;
304
305 let payload_len = header.payload_len as usize;
306 if payload_len > MAX_PAYLOAD_HARD_CAP {
307 return Err(DecodeError::PayloadTooLarge {
308 len: payload_len,
309 cap: MAX_PAYLOAD_HARD_CAP,
310 });
311 }
312
313 let frame_len = HEADER_LEN + payload_len + 4;
316 if buf.len() < frame_len {
317 return Err(DecodeError::TruncatedFrame);
318 }
319 if buf.len() > frame_len {
326 return Err(DecodeError::TrailingBytes {
327 extra: buf.len() - frame_len,
328 });
329 }
330
331 let payload = Payload::copy_from_slice(&buf[HEADER_LEN..HEADER_LEN + payload_len]);
332 let expected_crc =
333 u32::from_le_bytes(buf[HEADER_LEN + payload_len..frame_len].try_into().unwrap());
334 let computed_crc = crc32fast::hash(&payload);
335 if expected_crc != computed_crc {
336 return Err(DecodeError::PayloadCrcMismatch {
337 expected: expected_crc,
338 computed: computed_crc,
339 });
340 }
341
342 Ok(Self { header, payload })
343 }
344}
345
346#[cfg(test)]
347mod tests {
348 use super::*;
349 use crate::version::WIRE_VERSION;
350
351 fn push_header() -> Header {
352 Header::new(MessageType::Push, Durability::Durable, 0)
353 }
354
355 #[test]
358 fn header_encode_decode_round_trip() {
359 let h = push_header();
360 let encoded = h.encode();
361 let decoded = Header::decode(&encoded).unwrap();
362 assert_eq!(h, decoded);
363 }
364
365 #[test]
366 fn header_encode_is_16_bytes() {
367 let encoded = push_header().encode();
368 assert_eq!(encoded.len(), HEADER_LEN);
369 }
370
371 #[test]
372 fn header_decode_rejects_bad_magic() {
373 let mut buf = push_header().encode();
374 buf[0] = 0x00;
375 assert_eq!(Header::decode(&buf), Err(DecodeError::BadMagic));
376 }
377
378 #[test]
379 fn header_decode_rejects_future_version() {
380 let mut buf = push_header().encode();
381 buf[4] = WIRE_VERSION + 1;
382 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
384 buf[12..16].copy_from_slice(&crc.to_le_bytes());
385 assert_eq!(
386 Header::decode(&buf),
387 Err(DecodeError::VersionMismatch {
388 supported: WIRE_VERSION,
389 received: WIRE_VERSION + 1,
390 })
391 );
392 }
393
394 #[test]
395 fn header_decode_rejects_older_version() {
396 let mut buf = push_header().encode();
397 buf[4] = WIRE_VERSION - 1;
398 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
399 buf[12..16].copy_from_slice(&crc.to_le_bytes());
400 assert_eq!(
401 Header::decode(&buf),
402 Err(DecodeError::VersionMismatch {
403 supported: WIRE_VERSION,
404 received: WIRE_VERSION - 1,
405 })
406 );
407 }
408
409 #[test]
412 fn header_decode_version_checked_before_crc() {
413 let mut buf = push_header().encode();
414 buf[4] = WIRE_VERSION + 1;
416 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
417 buf[12..16].copy_from_slice(&crc.to_le_bytes());
418 match Header::decode(&buf) {
420 Err(DecodeError::VersionMismatch { .. }) => {}
421 other => panic!("expected VersionMismatch, got {other:?}"),
422 }
423 }
424
425 #[test]
426 fn header_decode_rejects_corrupted_crc() {
427 let mut buf = push_header().encode();
428 buf[12] ^= 0xff;
429 assert!(matches!(
430 Header::decode(&buf),
431 Err(DecodeError::HeaderCrcMismatch { .. })
432 ));
433 }
434
435 #[test]
436 fn header_decode_rejects_unknown_message_type() {
437 let mut buf = push_header().encode();
438 buf[5] = 0xff;
439 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
440 buf[12..16].copy_from_slice(&crc.to_le_bytes());
441 assert_eq!(
442 Header::decode(&buf),
443 Err(DecodeError::UnknownMessageType(0xff))
444 );
445 }
446
447 #[test]
448 fn header_decode_rejects_unknown_durability() {
449 let mut buf = push_header().encode();
450 buf[6] = 0xff;
451 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
452 buf[12..16].copy_from_slice(&crc.to_le_bytes());
453 assert_eq!(
454 Header::decode(&buf),
455 Err(DecodeError::UnknownDurability(0xff))
456 );
457 }
458
459 #[test]
463 fn header_decode_rejects_nonzero_flags() {
464 let buf = Header::new(MessageType::Push, Durability::Durable, 0x01).encode();
465 assert_eq!(
466 Header::decode(&buf),
467 Err(DecodeError::ReservedFlagsSet { flags: 0x01 })
468 );
469 let buf = Header::new(MessageType::Push, Durability::Durable, 0xff).encode();
471 assert_eq!(
472 Header::decode(&buf),
473 Err(DecodeError::ReservedFlagsSet { flags: 0xff })
474 );
475 }
476
477 #[test]
482 fn header_decode_flags_checked_after_durability() {
483 let mut buf = Header::new(MessageType::Push, Durability::Durable, 0x01).encode();
484 buf[6] = 0xff; let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
486 buf[12..16].copy_from_slice(&crc.to_le_bytes());
487 assert_eq!(
488 Header::decode(&buf),
489 Err(DecodeError::UnknownDurability(0xff))
490 );
491 }
492
493 #[test]
496 fn message_type_try_from_accepts_all_known_values() {
497 assert_eq!(MessageType::try_from(0x01).unwrap(), MessageType::Push);
498 assert_eq!(MessageType::try_from(0x02).unwrap(), MessageType::Ack);
499 assert_eq!(MessageType::try_from(0x03).unwrap(), MessageType::Nack);
500 assert_eq!(
501 MessageType::try_from(0x04).unwrap(),
502 MessageType::HealthCheck
503 );
504 assert_eq!(
505 MessageType::try_from(0x05).unwrap(),
506 MessageType::HealthCheckResponse
507 );
508 }
509
510 #[test]
511 fn message_type_try_from_rejects_unknown() {
512 assert!(MessageType::try_from(0x00).is_err());
513 assert!(MessageType::try_from(0x06).is_err());
514 assert!(MessageType::try_from(0xff).is_err());
515 }
516
517 #[test]
518 fn message_type_from_u8_round_trips() {
519 for mt in [
522 MessageType::Push,
523 MessageType::Ack,
524 MessageType::Nack,
525 MessageType::HealthCheck,
526 MessageType::HealthCheckResponse,
527 ] {
528 assert_eq!(MessageType::try_from(u8::from(mt)).unwrap(), mt);
529 }
530 }
531
532 #[test]
533 fn message_type_repr_values_match_wire() {
534 assert_eq!(MessageType::Push as u8, 0x01);
535 assert_eq!(MessageType::Ack as u8, 0x02);
536 assert_eq!(MessageType::Nack as u8, 0x03);
537 assert_eq!(MessageType::HealthCheck as u8, 0x04);
538 assert_eq!(MessageType::HealthCheckResponse as u8, 0x05);
539 }
540
541 fn round_trip(msg_type: MessageType) {
544 let payload = b"hello weir".to_vec();
545 let header = Header::new(msg_type, Durability::Durable, 0);
546 let env = Envelope::new(header, payload);
547 let encoded = env.encode();
548 let decoded = Envelope::decode(&encoded).unwrap();
549 assert_eq!(env, decoded);
550 }
551
552 #[test]
553 fn envelope_round_trip_push() {
554 round_trip(MessageType::Push);
555 }
556
557 #[test]
558 fn envelope_round_trip_ack() {
559 round_trip(MessageType::Ack);
560 }
561
562 #[test]
563 fn envelope_round_trip_nack() {
564 round_trip(MessageType::Nack);
565 }
566
567 #[test]
568 fn envelope_round_trip_health_check() {
569 round_trip(MessageType::HealthCheck);
570 }
571
572 #[test]
573 fn envelope_round_trip_health_check_response() {
574 round_trip(MessageType::HealthCheckResponse);
575 }
576
577 #[test]
578 fn envelope_round_trip_empty_payload() {
579 let header = push_header();
580 let env = Envelope::new(header, vec![]);
581 let encoded = env.encode();
582 let decoded = Envelope::decode(&encoded).unwrap();
583 assert_eq!(env, decoded);
584 }
585
586 #[test]
593 fn envelope_new_payload_len_equals_payload_length() {
594 for len in [0usize, 1, 7, 256, 65_536, 1_000_000] {
595 let env = Envelope::new(push_header(), vec![0xAB; len]);
596 assert_eq!(
597 env.header().payload_len() as usize,
598 len,
599 "payload_len must equal the actual payload length for len={len}"
600 );
601 }
602 }
603
604 #[test]
605 fn envelope_decode_rejects_truncated_header() {
606 let env = Envelope::new(push_header(), b"data".to_vec());
607 let encoded = env.encode();
608 assert_eq!(
609 Envelope::decode(&encoded[..HEADER_LEN - 1]),
610 Err(DecodeError::TruncatedFrame)
611 );
612 }
613
614 #[test]
615 fn envelope_decode_rejects_truncated_payload() {
616 let payload = b"hello".to_vec();
617 let env = Envelope::new(push_header(), payload);
618 let encoded = env.encode();
619 assert_eq!(
621 Envelope::decode(&encoded[..encoded.len() - 1]),
622 Err(DecodeError::TruncatedFrame)
623 );
624 }
625
626 #[test]
627 fn envelope_decode_rejects_header_only() {
628 let env = Envelope::new(push_header(), b"data".to_vec());
629 let encoded = env.encode();
630 assert_eq!(
631 Envelope::decode(&encoded[..HEADER_LEN]),
632 Err(DecodeError::TruncatedFrame)
633 );
634 }
635
636 #[test]
637 fn envelope_decode_rejects_corrupted_payload_crc() {
638 let payload = b"hello".to_vec();
639 let env = Envelope::new(push_header(), payload);
640 let mut encoded = env.encode();
641 let last = encoded.len() - 1;
643 encoded[last] ^= 0xff;
644 assert!(matches!(
645 Envelope::decode(&encoded),
646 Err(DecodeError::PayloadCrcMismatch { .. })
647 ));
648 }
649
650 #[test]
651 fn envelope_decode_rejects_trailing_bytes() {
652 let env = Envelope::new(push_header(), b"hello".to_vec());
656 let mut encoded = env.encode();
657 encoded.extend_from_slice(b"garbage"); assert_eq!(
659 Envelope::decode(&encoded),
660 Err(DecodeError::TrailingBytes { extra: 7 })
661 );
662 }
663
664 #[test]
665 fn envelope_decode_accepts_exact_frame() {
666 let env = Envelope::new(push_header(), b"hello".to_vec());
668 let encoded = env.encode();
669 assert_eq!(Envelope::decode(&encoded).unwrap(), env);
670 }
671
672 #[test]
673 fn envelope_decode_trailing_bytes_checked_after_payload_too_large() {
674 let oversized_len = (MAX_PAYLOAD_HARD_CAP + 1) as u32;
679 let mut header_bytes = push_header().encode();
680 header_bytes[8..12].copy_from_slice(&oversized_len.to_le_bytes());
681 let crc = crc32fast::hash(&header_bytes[..HEADER_CRC_COVERAGE]);
682 header_bytes[12..16].copy_from_slice(&crc.to_le_bytes());
683 assert_eq!(
684 Envelope::decode(&header_bytes),
685 Err(DecodeError::PayloadTooLarge {
686 len: MAX_PAYLOAD_HARD_CAP + 1,
687 cap: MAX_PAYLOAD_HARD_CAP,
688 })
689 );
690 }
691
692 #[test]
693 fn envelope_decode_rejects_oversized_payload_len() {
694 let oversized_len = (MAX_PAYLOAD_HARD_CAP + 1) as u32;
702 let mut header_bytes = Header::new(MessageType::Push, Durability::Durable, 0).encode();
703 header_bytes[8..12].copy_from_slice(&oversized_len.to_le_bytes());
704 let crc = crc32fast::hash(&header_bytes[..HEADER_CRC_COVERAGE]);
705 header_bytes[12..16].copy_from_slice(&crc.to_le_bytes());
706 assert_eq!(
707 Envelope::decode(&header_bytes),
708 Err(DecodeError::PayloadTooLarge {
709 len: MAX_PAYLOAD_HARD_CAP + 1,
710 cap: MAX_PAYLOAD_HARD_CAP,
711 })
712 );
713 }
714}