diff --git a/moq-transport/src/message/mod.rs b/moq-transport/src/message/mod.rs index 405ba21..4dad338 100644 --- a/moq-transport/src/message/mod.rs +++ b/moq-transport/src/message/mod.rs @@ -86,6 +86,11 @@ macro_rules! message_types { impl Decode for Message { fn decode(r: &mut R) -> Result { let t = u64::decode(r)?; + let len = u64::decode(r)?; + + if r.remaining() < len as usize { + return Err(DecodeError::InvalidLength(len as usize, r.remaining())); + } match t { $($val => { @@ -102,7 +107,18 @@ macro_rules! message_types { match self { $(Self::$name(ref m) => { self.id().encode(w)?; - m.encode(w) + + // Find out the length of the message + // by encoding it into a buffer and then encoding the length. + // This is a bit wasteful, but it's the only way to know the length. + let mut buf = Vec::new(); + m.encode(&mut buf).unwrap(); + (buf.len() as u64).encode(w)?; + + // At least don't encode the message twice. + // Instead, write the buffer directly to the writer. + w.put_slice(&buf); + Ok(()) },)* } }