moq_lite/ietf/
message.rs

1use crate::coding::{self, DecodeError, Sizer};
2use crate::ietf::Version;
3use std::fmt::Debug;
4
5use bytes::{Buf, BufMut};
6
7/// A trait for messages that are size-prefixed during encoding/decoding.
8///
9/// This trait wraps the existing Encode/Decode traits and automatically handles:
10/// - Prefixing messages with their encoded size during encoding
11/// - Reading the size prefix and validating exact consumption during decoding
12/// - Ensuring no bytes are left over or missing after decoding
13pub trait Message: Sized + Debug {
14	const ID: u64;
15
16	/// Encode this message with a size prefix.
17	fn encode_msg<W: BufMut>(&self, w: &mut W, version: Version);
18
19	/// Decode a size-prefixed message, ensuring exact size consumption.
20	fn decode_msg<B: Buf>(buf: &mut B, version: Version) -> Result<Self, DecodeError>;
21}
22
23impl<T: Message> coding::Encode<Version> for T {
24	fn encode<W: BufMut>(&self, w: &mut W, version: Version) {
25		// TODO Always encode 2 bytes for the size, then go back and populate it later.
26		// That way we can avoid calculating the size upfront.
27		let mut sizer = Sizer::default();
28		self.encode_msg(&mut sizer, version);
29		let size: u16 = sizer.size.try_into().expect("message too large");
30		size.encode(w, version);
31		self.encode_msg(w, version);
32	}
33}
34
35impl<T: Message> coding::Decode<Version> for T {
36	fn decode<B: Buf>(buf: &mut B, version: Version) -> Result<Self, DecodeError> {
37		let size = u16::decode(buf, version)?;
38		let mut limited = buf.take(size as usize);
39		let result = Self::decode_msg(&mut limited, version)?;
40		if limited.remaining() > 0 {
41			return Err(DecodeError::Long);
42		}
43		Ok(result)
44	}
45}