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 PushTracked = 0x06,
60 AckTracked = 0x07,
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub struct UnknownMessageType(pub u8);
70
71impl std::fmt::Display for UnknownMessageType {
72 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 write!(f, "unknown message type byte: {:#04x}", self.0)
74 }
75}
76
77impl std::error::Error for UnknownMessageType {}
78
79impl TryFrom<u8> for MessageType {
80 type Error = UnknownMessageType;
81
82 fn try_from(value: u8) -> Result<Self, Self::Error> {
83 match value {
84 0x01 => Ok(MessageType::Push),
85 0x02 => Ok(MessageType::Ack),
86 0x03 => Ok(MessageType::Nack),
87 0x04 => Ok(MessageType::HealthCheck),
88 0x05 => Ok(MessageType::HealthCheckResponse),
89 0x06 => Ok(MessageType::PushTracked),
90 0x07 => Ok(MessageType::AckTracked),
91 v => Err(UnknownMessageType(v)),
92 }
93 }
94}
95
96impl From<MessageType> for u8 {
97 fn from(m: MessageType) -> u8 {
99 m as u8
100 }
101}
102
103#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub struct Header {
114 version: u8,
115 message_type: MessageType,
116 durability: Durability,
117 flags: u8,
119 payload_len: u32,
120}
121
122impl Header {
123 #[must_use]
138 pub fn new(message_type: MessageType, durability: Durability, flags: u8) -> Self {
139 Self {
140 version: WIRE_VERSION,
141 message_type,
142 durability,
143 flags,
144 payload_len: 0,
145 }
146 }
147
148 pub fn version(&self) -> u8 {
151 self.version
152 }
153
154 pub fn message_type(&self) -> MessageType {
156 self.message_type
157 }
158
159 pub fn durability(&self) -> Durability {
161 self.durability
162 }
163
164 pub fn flags(&self) -> u8 {
166 self.flags
167 }
168
169 pub fn payload_len(&self) -> u32 {
173 self.payload_len
174 }
175
176 #[must_use]
179 pub fn encode(&self) -> [u8; HEADER_LEN] {
180 let mut buf = [0u8; HEADER_LEN];
181 buf[0..4].copy_from_slice(&MAGIC);
182 buf[4] = self.version;
183 buf[5] = u8::from(self.message_type);
184 buf[6] = u8::from(self.durability);
185 buf[7] = self.flags;
186 buf[8..12].copy_from_slice(&self.payload_len.to_le_bytes());
187 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
188 buf[12..16].copy_from_slice(&crc.to_le_bytes());
189 buf
190 }
191
192 pub fn decode(buf: &[u8; HEADER_LEN]) -> Result<Self, DecodeError> {
202 if buf[0..4] != MAGIC {
203 return Err(DecodeError::BadMagic);
204 }
205
206 let version = buf[4];
207 if version != WIRE_VERSION {
208 return Err(DecodeError::VersionMismatch {
209 supported: WIRE_VERSION,
210 received: version,
211 });
212 }
213
214 let expected_crc = u32::from_le_bytes(buf[12..16].try_into().unwrap());
215 let computed_crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
216 if expected_crc != computed_crc {
217 return Err(DecodeError::HeaderCrcMismatch {
218 expected: expected_crc,
219 computed: computed_crc,
220 });
221 }
222
223 let message_type =
224 MessageType::try_from(buf[5]).map_err(|e| DecodeError::UnknownMessageType(e.0))?;
225 let durability =
226 Durability::try_from(buf[6]).map_err(|e| DecodeError::UnknownDurability(e.0))?;
227 let flags = buf[7];
233 if flags != 0 {
234 return Err(DecodeError::ReservedFlagsSet { flags });
235 }
236 let payload_len = u32::from_le_bytes(buf[8..12].try_into().unwrap());
237
238 Ok(Self {
239 version,
240 message_type,
241 durability,
242 flags,
243 payload_len,
244 })
245 }
246}
247
248#[derive(Debug, Clone, PartialEq, Eq)]
255pub struct Envelope {
256 header: Header,
257 payload: Payload,
258}
259
260impl Envelope {
261 #[must_use]
265 pub fn new(header: Header, payload: impl Into<Payload>) -> Self {
266 let payload = payload.into();
267 let payload_len = u32::try_from(payload.len()).unwrap_or(u32::MAX);
273 let header = Header {
274 payload_len,
275 ..header
276 };
277 Self { header, payload }
278 }
279
280 pub fn header(&self) -> Header {
282 self.header
283 }
284
285 pub fn payload(&self) -> &Payload {
287 &self.payload
288 }
289
290 #[must_use]
292 pub fn encode(&self) -> Vec<u8> {
293 let mut out = Vec::with_capacity(HEADER_LEN + self.payload.len() + 4);
294 out.extend_from_slice(&self.header.encode());
295 out.extend_from_slice(&self.payload);
296 let payload_crc = crc32fast::hash(&self.payload);
297 out.extend_from_slice(&payload_crc.to_le_bytes());
298 out
299 }
300
301 pub fn decode(buf: &[u8]) -> Result<Self, DecodeError> {
315 if buf.len() < HEADER_LEN {
316 return Err(DecodeError::TruncatedFrame);
317 }
318
319 let header_bytes: &[u8; HEADER_LEN] = buf[..HEADER_LEN]
320 .try_into()
321 .expect("slice is exactly HEADER_LEN bytes, checked above");
322 let header = Header::decode(header_bytes)?;
323
324 let payload_len = header.payload_len as usize;
325 if payload_len > MAX_PAYLOAD_HARD_CAP {
326 return Err(DecodeError::PayloadTooLarge {
327 len: payload_len,
328 cap: MAX_PAYLOAD_HARD_CAP,
329 });
330 }
331
332 let frame_len = HEADER_LEN + payload_len + 4;
335 if buf.len() < frame_len {
336 return Err(DecodeError::TruncatedFrame);
337 }
338 if buf.len() > frame_len {
345 return Err(DecodeError::TrailingBytes {
346 extra: buf.len() - frame_len,
347 });
348 }
349
350 let payload = Payload::copy_from_slice(&buf[HEADER_LEN..HEADER_LEN + payload_len]);
351 let expected_crc =
352 u32::from_le_bytes(buf[HEADER_LEN + payload_len..frame_len].try_into().unwrap());
353 let computed_crc = crc32fast::hash(&payload);
354 if expected_crc != computed_crc {
355 return Err(DecodeError::PayloadCrcMismatch {
356 expected: expected_crc,
357 computed: computed_crc,
358 });
359 }
360
361 Ok(Self { header, payload })
362 }
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368 use crate::version::WIRE_VERSION;
369
370 fn push_header() -> Header {
371 Header::new(MessageType::Push, Durability::Durable, 0)
372 }
373
374 #[test]
377 fn header_encode_decode_round_trip() {
378 let h = push_header();
379 let encoded = h.encode();
380 let decoded = Header::decode(&encoded).unwrap();
381 assert_eq!(h, decoded);
382 }
383
384 #[test]
385 fn header_encode_is_16_bytes() {
386 let encoded = push_header().encode();
387 assert_eq!(encoded.len(), HEADER_LEN);
388 }
389
390 #[test]
391 fn header_decode_rejects_bad_magic() {
392 let mut buf = push_header().encode();
393 buf[0] = 0x00;
394 assert_eq!(Header::decode(&buf), Err(DecodeError::BadMagic));
395 }
396
397 #[test]
398 fn header_decode_rejects_future_version() {
399 let mut buf = push_header().encode();
400 buf[4] = WIRE_VERSION + 1;
401 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
403 buf[12..16].copy_from_slice(&crc.to_le_bytes());
404 assert_eq!(
405 Header::decode(&buf),
406 Err(DecodeError::VersionMismatch {
407 supported: WIRE_VERSION,
408 received: WIRE_VERSION + 1,
409 })
410 );
411 }
412
413 #[test]
414 fn header_decode_rejects_older_version() {
415 let mut buf = push_header().encode();
416 buf[4] = WIRE_VERSION - 1;
417 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
418 buf[12..16].copy_from_slice(&crc.to_le_bytes());
419 assert_eq!(
420 Header::decode(&buf),
421 Err(DecodeError::VersionMismatch {
422 supported: WIRE_VERSION,
423 received: WIRE_VERSION - 1,
424 })
425 );
426 }
427
428 #[test]
431 fn header_decode_version_checked_before_crc() {
432 let mut buf = push_header().encode();
433 buf[4] = WIRE_VERSION + 1;
435 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
436 buf[12..16].copy_from_slice(&crc.to_le_bytes());
437 match Header::decode(&buf) {
439 Err(DecodeError::VersionMismatch { .. }) => {}
440 other => panic!("expected VersionMismatch, got {other:?}"),
441 }
442 }
443
444 #[test]
445 fn header_decode_rejects_corrupted_crc() {
446 let mut buf = push_header().encode();
447 buf[12] ^= 0xff;
448 assert!(matches!(
449 Header::decode(&buf),
450 Err(DecodeError::HeaderCrcMismatch { .. })
451 ));
452 }
453
454 #[test]
455 fn header_decode_rejects_unknown_message_type() {
456 let mut buf = push_header().encode();
457 buf[5] = 0xff;
458 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
459 buf[12..16].copy_from_slice(&crc.to_le_bytes());
460 assert_eq!(
461 Header::decode(&buf),
462 Err(DecodeError::UnknownMessageType(0xff))
463 );
464 }
465
466 #[test]
467 fn header_decode_rejects_unknown_durability() {
468 let mut buf = push_header().encode();
469 buf[6] = 0xff;
470 let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
471 buf[12..16].copy_from_slice(&crc.to_le_bytes());
472 assert_eq!(
473 Header::decode(&buf),
474 Err(DecodeError::UnknownDurability(0xff))
475 );
476 }
477
478 #[test]
482 fn header_decode_rejects_nonzero_flags() {
483 let buf = Header::new(MessageType::Push, Durability::Durable, 0x01).encode();
484 assert_eq!(
485 Header::decode(&buf),
486 Err(DecodeError::ReservedFlagsSet { flags: 0x01 })
487 );
488 let buf = Header::new(MessageType::Push, Durability::Durable, 0xff).encode();
490 assert_eq!(
491 Header::decode(&buf),
492 Err(DecodeError::ReservedFlagsSet { flags: 0xff })
493 );
494 }
495
496 #[test]
501 fn header_decode_flags_checked_after_durability() {
502 let mut buf = Header::new(MessageType::Push, Durability::Durable, 0x01).encode();
503 buf[6] = 0xff; let crc = crc32fast::hash(&buf[..HEADER_CRC_COVERAGE]);
505 buf[12..16].copy_from_slice(&crc.to_le_bytes());
506 assert_eq!(
507 Header::decode(&buf),
508 Err(DecodeError::UnknownDurability(0xff))
509 );
510 }
511
512 #[test]
515 fn message_type_try_from_accepts_all_known_values() {
516 assert_eq!(MessageType::try_from(0x01).unwrap(), MessageType::Push);
517 assert_eq!(MessageType::try_from(0x02).unwrap(), MessageType::Ack);
518 assert_eq!(MessageType::try_from(0x03).unwrap(), MessageType::Nack);
519 assert_eq!(
520 MessageType::try_from(0x04).unwrap(),
521 MessageType::HealthCheck
522 );
523 assert_eq!(
524 MessageType::try_from(0x05).unwrap(),
525 MessageType::HealthCheckResponse
526 );
527 assert_eq!(
528 MessageType::try_from(0x06).unwrap(),
529 MessageType::PushTracked
530 );
531 assert_eq!(
532 MessageType::try_from(0x07).unwrap(),
533 MessageType::AckTracked
534 );
535 }
536
537 #[test]
538 fn message_type_try_from_rejects_unknown() {
539 assert!(MessageType::try_from(0x00).is_err());
540 assert!(MessageType::try_from(0x08).is_err());
541 assert!(MessageType::try_from(0xff).is_err());
542 }
543
544 #[test]
545 fn message_type_from_u8_round_trips() {
546 for mt in [
549 MessageType::Push,
550 MessageType::Ack,
551 MessageType::Nack,
552 MessageType::HealthCheck,
553 MessageType::HealthCheckResponse,
554 MessageType::PushTracked,
555 MessageType::AckTracked,
556 ] {
557 assert_eq!(MessageType::try_from(u8::from(mt)).unwrap(), mt);
558 }
559 }
560
561 #[test]
562 fn message_type_repr_values_match_wire() {
563 assert_eq!(MessageType::Push as u8, 0x01);
564 assert_eq!(MessageType::Ack as u8, 0x02);
565 assert_eq!(MessageType::Nack as u8, 0x03);
566 assert_eq!(MessageType::HealthCheck as u8, 0x04);
567 assert_eq!(MessageType::HealthCheckResponse as u8, 0x05);
568 assert_eq!(MessageType::PushTracked as u8, 0x06);
569 assert_eq!(MessageType::AckTracked as u8, 0x07);
570 }
571
572 fn round_trip(msg_type: MessageType) {
575 let payload = b"hello weir".to_vec();
576 let header = Header::new(msg_type, Durability::Durable, 0);
577 let env = Envelope::new(header, payload);
578 let encoded = env.encode();
579 let decoded = Envelope::decode(&encoded).unwrap();
580 assert_eq!(env, decoded);
581 }
582
583 #[test]
584 fn envelope_round_trip_push() {
585 round_trip(MessageType::Push);
586 }
587
588 #[test]
589 fn envelope_round_trip_ack() {
590 round_trip(MessageType::Ack);
591 }
592
593 #[test]
594 fn envelope_round_trip_nack() {
595 round_trip(MessageType::Nack);
596 }
597
598 #[test]
599 fn envelope_round_trip_health_check() {
600 round_trip(MessageType::HealthCheck);
601 }
602
603 #[test]
604 fn envelope_round_trip_health_check_response() {
605 round_trip(MessageType::HealthCheckResponse);
606 }
607
608 #[test]
609 fn envelope_round_trip_push_tracked() {
610 round_trip(MessageType::PushTracked);
611 }
612
613 #[test]
614 fn envelope_round_trip_ack_tracked() {
615 round_trip(MessageType::AckTracked);
616 }
617
618 #[test]
619 fn envelope_round_trip_empty_payload() {
620 let header = push_header();
621 let env = Envelope::new(header, vec![]);
622 let encoded = env.encode();
623 let decoded = Envelope::decode(&encoded).unwrap();
624 assert_eq!(env, decoded);
625 }
626
627 #[test]
634 fn envelope_new_payload_len_equals_payload_length() {
635 for len in [0usize, 1, 7, 256, 65_536, 1_000_000] {
636 let env = Envelope::new(push_header(), vec![0xAB; len]);
637 assert_eq!(
638 env.header().payload_len() as usize,
639 len,
640 "payload_len must equal the actual payload length for len={len}"
641 );
642 }
643 }
644
645 #[test]
646 fn envelope_decode_rejects_truncated_header() {
647 let env = Envelope::new(push_header(), b"data".to_vec());
648 let encoded = env.encode();
649 assert_eq!(
650 Envelope::decode(&encoded[..HEADER_LEN - 1]),
651 Err(DecodeError::TruncatedFrame)
652 );
653 }
654
655 #[test]
656 fn envelope_decode_rejects_truncated_payload() {
657 let payload = b"hello".to_vec();
658 let env = Envelope::new(push_header(), payload);
659 let encoded = env.encode();
660 assert_eq!(
662 Envelope::decode(&encoded[..encoded.len() - 1]),
663 Err(DecodeError::TruncatedFrame)
664 );
665 }
666
667 #[test]
668 fn envelope_decode_rejects_header_only() {
669 let env = Envelope::new(push_header(), b"data".to_vec());
670 let encoded = env.encode();
671 assert_eq!(
672 Envelope::decode(&encoded[..HEADER_LEN]),
673 Err(DecodeError::TruncatedFrame)
674 );
675 }
676
677 #[test]
678 fn envelope_decode_rejects_corrupted_payload_crc() {
679 let payload = b"hello".to_vec();
680 let env = Envelope::new(push_header(), payload);
681 let mut encoded = env.encode();
682 let last = encoded.len() - 1;
684 encoded[last] ^= 0xff;
685 assert!(matches!(
686 Envelope::decode(&encoded),
687 Err(DecodeError::PayloadCrcMismatch { .. })
688 ));
689 }
690
691 #[test]
692 fn envelope_decode_rejects_trailing_bytes() {
693 let env = Envelope::new(push_header(), b"hello".to_vec());
697 let mut encoded = env.encode();
698 encoded.extend_from_slice(b"garbage"); assert_eq!(
700 Envelope::decode(&encoded),
701 Err(DecodeError::TrailingBytes { extra: 7 })
702 );
703 }
704
705 #[test]
706 fn envelope_decode_accepts_exact_frame() {
707 let env = Envelope::new(push_header(), b"hello".to_vec());
709 let encoded = env.encode();
710 assert_eq!(Envelope::decode(&encoded).unwrap(), env);
711 }
712
713 #[test]
714 fn envelope_decode_trailing_bytes_checked_after_payload_too_large() {
715 let oversized_len = (MAX_PAYLOAD_HARD_CAP + 1) as u32;
720 let mut header_bytes = push_header().encode();
721 header_bytes[8..12].copy_from_slice(&oversized_len.to_le_bytes());
722 let crc = crc32fast::hash(&header_bytes[..HEADER_CRC_COVERAGE]);
723 header_bytes[12..16].copy_from_slice(&crc.to_le_bytes());
724 assert_eq!(
725 Envelope::decode(&header_bytes),
726 Err(DecodeError::PayloadTooLarge {
727 len: MAX_PAYLOAD_HARD_CAP + 1,
728 cap: MAX_PAYLOAD_HARD_CAP,
729 })
730 );
731 }
732
733 #[test]
734 fn envelope_decode_rejects_oversized_payload_len() {
735 let oversized_len = (MAX_PAYLOAD_HARD_CAP + 1) as u32;
743 let mut header_bytes = Header::new(MessageType::Push, Durability::Durable, 0).encode();
744 header_bytes[8..12].copy_from_slice(&oversized_len.to_le_bytes());
745 let crc = crc32fast::hash(&header_bytes[..HEADER_CRC_COVERAGE]);
746 header_bytes[12..16].copy_from_slice(&crc.to_le_bytes());
747 assert_eq!(
748 Envelope::decode(&header_bytes),
749 Err(DecodeError::PayloadTooLarge {
750 len: MAX_PAYLOAD_HARD_CAP + 1,
751 cap: MAX_PAYLOAD_HARD_CAP,
752 })
753 );
754 }
755}