refactor(proto): centralize packet envelope handling

crates/irtt-proto/src/close.rs

@@ -1,6 +1,8 @@
 use crate::{
-    flags::{FLAG_CLOSE, FLAG_HMAC},
+    envelope::{self, FlagRule},
-    hmac, write_header, Result, HMAC_SIZE,
+    flags::{FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY},
+    layout::PacketLayout,
+    Result,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -9,27 +11,24 @@ pub struct CloseRequest {
 }
 
 pub fn encode_close_request(request: &CloseRequest, hmac_key: Option<&[u8]>) -> Result<Vec<u8>> {
-    let mut flags = FLAG_CLOSE;
+    let mut out = envelope::begin(
-    if hmac_key.is_some() {
+        FLAG_CLOSE,
-        flags |= FLAG_HMAC;
+        hmac_key,
-    }
+        &[
-
+            FlagRule::Require(FLAG_CLOSE),
-    let mut out = Vec::new();
+            FlagRule::Reject(FLAG_OPEN),
-    write_header(&mut out, flags);
+            FlagRule::Reject(FLAG_REPLY),
-    if hmac_key.is_some() {
+        ],
-        out.extend_from_slice(&[0; HMAC_SIZE]);
+        PacketLayout::close_request(hmac_key.is_some()).header_len(),
-    }
+    )?;
     out.extend_from_slice(&request.token.to_le_bytes());
-    if let Some(key) = hmac_key {
+    envelope::finish(out, hmac_key)
-        hmac::compute_hmac_in_place(key, &mut out, hmac::hmac_offset())?;
-    }
-    Ok(out)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{hmac, FLAG_HMAC};
+    use crate::{hmac, FLAG_HMAC, HMAC_SIZE};
 
     #[test]
     fn hmac_close_request_places_token_after_hmac() {

crates/irtt-proto/src/echo.rs

@@ -1,11 +1,9 @@
 use crate::{
-    flags::{has, FLAG_HMAC, FLAG_OPEN, FLAG_REPLY},
+    envelope::{self, FlagRule},
-    hmac,
+    flags::{FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY},
     layout::{echo_packet_len, PacketLayout},
-    open::{check_hmac_presence, reject},
     params::Params,
-    validate_header, write_header, ProtoError, Result, HEADER_SIZE, HMAC_SIZE, RECV_COUNT_SIZE,
+    ProtoError, Result, RECV_COUNT_SIZE, RECV_WINDOW_SIZE, SEQ_SIZE, TIMESTAMP_SIZE, TOKEN_SIZE,
-    RECV_WINDOW_SIZE, SEQ_SIZE, TIMESTAMP_SIZE, TOKEN_SIZE,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -38,10 +36,6 @@ pub struct TimestampFields {
 }
 
 pub fn encode_echo_request(request: &EchoRequest, hmac_key: Option<&[u8]>) -> Result<Vec<u8>> {
-    let mut flags = 0;
-    if hmac_key.is_some() {
-        flags |= FLAG_HMAC;
-    }
     let layout = PacketLayout::echo(hmac_key.is_some(), &request.params);
     let len = echo_packet_len(hmac_key.is_some(), &request.params)?;
     let payload_offset = layout.header_len();
@@ -52,21 +46,23 @@ pub fn encode_echo_request(request: &EchoRequest, hmac_key: Option<&[u8]>) -> Re
             provided: request.payload.len(),
         });
     }
-    let mut out = Vec::with_capacity(len);
+    let mut out = envelope::begin(
-    write_header(&mut out, flags);
+        0,
-    if hmac_key.is_some() {
+        hmac_key,
-        out.extend_from_slice(&[0; HMAC_SIZE]);
+        &[
-    }
+            FlagRule::Reject(FLAG_OPEN),
+            FlagRule::Reject(FLAG_REPLY),
+            FlagRule::Reject(FLAG_CLOSE),
+        ],
+        len,
+    )?;
     out.extend_from_slice(&request.token.to_le_bytes());
     out.extend_from_slice(&request.sequence.to_le_bytes());
     push_zeroed_layout_tail(layout, &mut out);
     out.resize(len, 0);
     out[payload_offset..payload_offset + request.payload.len()].copy_from_slice(&request.payload);
 
-    if let Some(key) = hmac_key {
+    envelope::finish(out, hmac_key)
-        hmac::compute_hmac_in_place(key, &mut out, hmac::hmac_offset())?;
-    }
-    Ok(out)
 }
 
 pub fn decode_echo_reply(
@@ -74,35 +70,17 @@ pub fn decode_echo_reply(
     params: &Params,
     hmac_key: Option<&[u8]>,
 ) -> Result<EchoReply> {
-    let flags = validate_header(packet)?;
+    let envelope = envelope::decode(
-    reject(flags, FLAG_OPEN)?;
+        packet,
-    crate::open::require(flags, FLAG_REPLY)?;
+        hmac_key,
-    check_hmac_presence(flags, hmac_key)?;
+        &[FlagRule::Reject(FLAG_OPEN), FlagRule::Require(FLAG_REPLY)],
+    )?;
+    let layout = PacketLayout::echo(hmac_key.is_some(), params);
+    validate_echo_length(packet, params, layout)?;
+    envelope::verify(packet, hmac_key)?;
 
-    let layout = PacketLayout::echo(has(flags, FLAG_HMAC), params);
     let header_len = layout.header_len();
-    if packet.len() < header_len {
+    let mut pos = envelope.body_offset;
-        return Err(ProtoError::PacketTooShort {
-            needed: header_len,
-            actual: packet.len(),
-        });
-    }
-    let expected_len = echo_packet_len(has(flags, FLAG_HMAC), params)?;
-    if packet.len() != expected_len {
-        return Err(ProtoError::PacketLengthMismatch {
-            expected: expected_len,
-            actual: packet.len(),
-        });
-    }
-
-    if let Some(key) = hmac_key {
-        hmac::verify_hmac(key, packet, hmac::hmac_offset())?;
-    }
-
-    let mut pos = HEADER_SIZE;
-    if layout.hmac {
-        pos += HMAC_SIZE;
-    }
     let token = read_u64(packet, &mut pos);
     let sequence = read_u32(packet, &mut pos);
     let recv_count = layout.recv_count.then(|| read_u32(packet, &mut pos));
@@ -117,7 +95,7 @@ pub fn decode_echo_reply(
     };
 
     Ok(EchoReply {
-        flags,
+        flags: envelope.flags,
         token,
         sequence,
         recv_count,
@@ -127,6 +105,24 @@ pub fn decode_echo_reply(
     })
 }
 
+fn validate_echo_length(packet: &[u8], params: &Params, layout: PacketLayout) -> Result<()> {
+    let header_len = layout.header_len();
+    if packet.len() < header_len {
+        return Err(ProtoError::PacketTooShort {
+            needed: header_len,
+            actual: packet.len(),
+        });
+    }
+    let expected_len = echo_packet_len(layout.hmac, params)?;
+    if packet.len() != expected_len {
+        return Err(ProtoError::PacketLengthMismatch {
+            expected: expected_len,
+            actual: packet.len(),
+        });
+    }
+    Ok(())
+}
+
 fn push_zeroed_layout_tail(layout: PacketLayout, out: &mut Vec<u8>) {
     if layout.recv_count {
         out.extend_from_slice(&[0; RECV_COUNT_SIZE]);
@@ -160,7 +156,11 @@ fn read_i64(packet: &[u8], pos: &mut usize) -> i64 {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::params::{Clock, ReceivedStats, StampAt};
+    use crate::{
+        hmac,
+        params::{Clock, ReceivedStats, StampAt},
+        write_header, FLAG_HMAC, HMAC_SIZE,
+    };
 
     fn default_params() -> Params {
         Params {

crates/irtt-proto/src/envelope.rs

@@ -0,0 +1,180 @@
+use crate::{
+    flags::{self, has, FLAG_HMAC},
+    hmac, validate_header, write_header, ProtoError, Result, HEADER_SIZE, HMAC_SIZE,
+};
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum FlagRule {
+    Require(u8),
+    Reject(u8),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct Envelope {
+    pub(crate) flags: u8,
+    pub(crate) body_offset: usize,
+}
+
+pub(crate) fn decode(
+    packet: &[u8],
+    hmac_key: Option<&[u8]>,
+    rules: &[FlagRule],
+) -> Result<Envelope> {
+    let flags = validate_header(packet)?;
+    validate_packet_type(flags, rules)?;
+    check_hmac_presence(flags, hmac_key)?;
+
+    Ok(Envelope {
+        flags,
+        body_offset: HEADER_SIZE + if has(flags, FLAG_HMAC) { HMAC_SIZE } else { 0 },
+    })
+}
+
+pub(crate) fn begin(
+    flags: u8,
+    hmac_key: Option<&[u8]>,
+    rules: &[FlagRule],
+    capacity: usize,
+) -> Result<Vec<u8>> {
+    // The encoder key is authoritative: authenticated encoders set FLAG_HMAC,
+    // while unauthenticated encoders clear any caller-supplied FLAG_HMAC.
+    let flags = with_hmac_flag(flags, hmac_key.is_some());
+    flags::validate_flags(flags)?;
+    validate_packet_type(flags, rules)?;
+
+    let minimum_capacity = HEADER_SIZE + if hmac_key.is_some() { HMAC_SIZE } else { 0 };
+    let mut out = Vec::with_capacity(capacity.max(minimum_capacity));
+    write_header(&mut out, flags);
+    if hmac_key.is_some() {
+        out.extend_from_slice(&[0; HMAC_SIZE]);
+    }
+    Ok(out)
+}
+
+pub(crate) fn verify(packet: &[u8], hmac_key: Option<&[u8]>) -> Result<()> {
+    if let Some(key) = hmac_key {
+        hmac::verify_hmac(key, packet, hmac::hmac_offset())?;
+    }
+    Ok(())
+}
+
+pub(crate) fn finish(mut packet: Vec<u8>, hmac_key: Option<&[u8]>) -> Result<Vec<u8>> {
+    if let Some(key) = hmac_key {
+        hmac::compute_hmac_in_place(key, &mut packet, hmac::hmac_offset())?;
+    }
+    Ok(packet)
+}
+
+fn validate_packet_type(flags: u8, rules: &[FlagRule]) -> Result<()> {
+    for rule in rules {
+        match *rule {
+            FlagRule::Require(flag) if !has(flags, flag) => {
+                return Err(ProtoError::MissingFlag(flag));
+            }
+            FlagRule::Reject(flag) if has(flags, flag) => {
+                return Err(ProtoError::UnexpectedFlag(flag));
+            }
+            FlagRule::Require(_) | FlagRule::Reject(_) => {}
+        }
+    }
+    Ok(())
+}
+
+fn check_hmac_presence(flags: u8, hmac_key: Option<&[u8]>) -> Result<()> {
+    if has(flags, FLAG_HMAC) == hmac_key.is_some() {
+        Ok(())
+    } else {
+        Err(ProtoError::HmacPresenceMismatch)
+    }
+}
+
+fn with_hmac_flag(flags: u8, authenticated: bool) -> u8 {
+    if authenticated {
+        flags | FLAG_HMAC
+    } else {
+        flags & !FLAG_HMAC
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        flags::{FLAG_OPEN, FLAG_REPLY},
+        MAGIC,
+    };
+
+    const KEY: &[u8] = b"testkey";
+    const OPEN_REQUEST_RULES: &[FlagRule] =
+        &[FlagRule::Require(FLAG_OPEN), FlagRule::Reject(FLAG_REPLY)];
+
+    fn header(flags: u8) -> Vec<u8> {
+        let mut packet = MAGIC.to_vec();
+        packet.push(flags);
+        packet
+    }
+
+    #[test]
+    fn decode_rejects_bad_magic_and_reserved_flags() {
+        assert_eq!(
+            decode(&[0, MAGIC[1], MAGIC[2], FLAG_OPEN], None, &[]),
+            Err(ProtoError::BadMagic)
+        );
+        assert_eq!(
+            decode(&header(FLAG_OPEN | 0x10), None, &[]),
+            Err(ProtoError::ReservedFlags(0x10))
+        );
+    }
+
+    #[test]
+    fn decode_enforces_required_and_rejected_flags() {
+        assert_eq!(
+            decode(&header(0), None, OPEN_REQUEST_RULES),
+            Err(ProtoError::MissingFlag(FLAG_OPEN))
+        );
+        assert_eq!(
+            decode(&header(FLAG_OPEN | FLAG_REPLY), None, OPEN_REQUEST_RULES),
+            Err(ProtoError::UnexpectedFlag(FLAG_REPLY))
+        );
+    }
+
+    #[test]
+    fn decode_rejects_hmac_presence_mismatch() {
+        let mut authenticated = header(FLAG_OPEN | FLAG_HMAC);
+        authenticated.extend_from_slice(&[0; HMAC_SIZE]);
+        assert_eq!(
+            decode(&authenticated, None, OPEN_REQUEST_RULES),
+            Err(ProtoError::HmacPresenceMismatch)
+        );
+        assert_eq!(
+            decode(&header(FLAG_OPEN), Some(KEY), OPEN_REQUEST_RULES),
+            Err(ProtoError::HmacPresenceMismatch)
+        );
+    }
+
+    #[test]
+    fn verify_rejects_bad_hmac() {
+        let mut packet = begin(FLAG_OPEN, Some(KEY), OPEN_REQUEST_RULES, 22).unwrap();
+        packet.extend_from_slice(&[1, 2]);
+        let mut packet = finish(packet, Some(KEY)).unwrap();
+        *packet.last_mut().unwrap() ^= 0x80;
+
+        decode(&packet, Some(KEY), OPEN_REQUEST_RULES).unwrap();
+        assert_eq!(verify(&packet, Some(KEY)), Err(ProtoError::BadHmac));
+    }
+
+    #[test]
+    fn encoder_key_authoritatively_normalizes_hmac_flag() {
+        let plain = begin(FLAG_OPEN | FLAG_HMAC, None, OPEN_REQUEST_RULES, HEADER_SIZE).unwrap();
+        assert_eq!(plain[3], FLAG_OPEN);
+
+        let authenticated = begin(
+            FLAG_OPEN,
+            Some(KEY),
+            OPEN_REQUEST_RULES,
+            HEADER_SIZE + HMAC_SIZE,
+        )
+        .unwrap();
+        assert_eq!(authenticated[3], FLAG_OPEN | FLAG_HMAC);
+    }
+}

crates/irtt-proto/src/lib.rs

@@ -15,11 +15,14 @@
 //! The crate also provides packet layout calculation, open/echo/close packet
 //! encoding and decoding, and optional HMAC placement, computation, and
 //! verification helpers.
+//! Encoder `hmac_key` arguments are authoritative: `Some(key)` adds
+//! `FLAG_HMAC`, while `None` removes it from caller-supplied reply flags.
 //!
 #![forbid(unsafe_code)]
 
 pub mod close;
 pub mod echo;
+mod envelope;
 pub mod error;
 pub mod flags;
 pub mod hmac;

crates/irtt-proto/src/open.rs

@@ -1,8 +1,9 @@
 use crate::{
-    flags::{has, FLAG_CLOSE, FLAG_HMAC, FLAG_OPEN, FLAG_REPLY},
+    envelope::{self, FlagRule},
-    hmac,
+    flags::{has, FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY},
+    layout::PacketLayout,
     params::Params,
-    validate_header, write_header, ProtoError, Result, HEADER_SIZE, HMAC_SIZE, TOKEN_SIZE,
+    ProtoError, Result, TOKEN_SIZE,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -23,36 +24,26 @@ pub fn encode_open_request(request: &OpenRequest, hmac_key: Option<&[u8]>) -> Re
     if request.close {
         flags |= FLAG_CLOSE;
     }
-    if hmac_key.is_some() {
+    let params = request.params.encode();
-        flags |= FLAG_HMAC;
+    let mut out = envelope::begin(
-    }
+        flags,
-
+        hmac_key,
-    let mut out = Vec::new();
+        &[FlagRule::Require(FLAG_OPEN), FlagRule::Reject(FLAG_REPLY)],
-    write_header(&mut out, flags);
+        PacketLayout::open_request(hmac_key.is_some()).header_len() + params.len(),
-    if hmac_key.is_some() {
+    )?;
-        out.extend_from_slice(&[0; HMAC_SIZE]);
+    out.extend_from_slice(&params);
-    }
+    envelope::finish(out, hmac_key)
-    out.extend_from_slice(&request.params.encode());
-
-    if let Some(key) = hmac_key {
-        hmac::compute_hmac_in_place(key, &mut out, hmac::hmac_offset())?;
-    }
-    Ok(out)
 }
 
 pub fn decode_open_reply(packet: &[u8], hmac_key: Option<&[u8]>) -> Result<OpenReply> {
-    let flags = validate_header(packet)?;
+    let envelope = envelope::decode(
-    require(flags, FLAG_OPEN)?;
+        packet,
-    require(flags, FLAG_REPLY)?;
+        hmac_key,
-    check_hmac_presence(flags, hmac_key)?;
+        &[FlagRule::Require(FLAG_OPEN), FlagRule::Require(FLAG_REPLY)],
-    if let Some(key) = hmac_key {
+    )?;
-        hmac::verify_hmac(key, packet, hmac::hmac_offset())?;
+    envelope::verify(packet, hmac_key)?;
-    }
 
-    let mut pos = HEADER_SIZE;
+    let mut pos = envelope.body_offset;
-    if has(flags, FLAG_HMAC) {
-        pos += HMAC_SIZE;
-    }
     let needed = pos + TOKEN_SIZE;
     if packet.len() < needed {
         return Err(ProtoError::PacketTooShort {
@@ -62,45 +53,25 @@ pub fn decode_open_reply(packet: &[u8], hmac_key: Option<&[u8]>) -> Result<OpenR
     }
     let token = u64::from_le_bytes(packet[pos..pos + TOKEN_SIZE].try_into().unwrap());
     pos += TOKEN_SIZE;
-    if token == 0 && !has(flags, FLAG_CLOSE) {
+    if token == 0 && !has(envelope.flags, FLAG_CLOSE) {
         return Err(ProtoError::ZeroToken);
     }
 
     Ok(OpenReply {
-        flags,
+        flags: envelope.flags,
         token,
         params: Params::decode(&packet[pos..])?,
     })
 }
 
-pub(crate) fn require(flags: u8, flag: u8) -> Result<()> {
-    if has(flags, flag) {
-        Ok(())
-    } else {
-        Err(ProtoError::MissingFlag(flag))
-    }
-}
-
-pub(crate) fn reject(flags: u8, flag: u8) -> Result<()> {
-    if has(flags, flag) {
-        Err(ProtoError::UnexpectedFlag(flag))
-    } else {
-        Ok(())
-    }
-}
-
-pub(crate) fn check_hmac_presence(flags: u8, hmac_key: Option<&[u8]>) -> Result<()> {
-    if has(flags, FLAG_HMAC) == hmac_key.is_some() {
-        Ok(())
-    } else {
-        Err(ProtoError::HmacPresenceMismatch)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::params::{Clock, ReceivedStats, StampAt};
+    use crate::{
+        hmac,
+        params::{Clock, ReceivedStats, StampAt},
+        FLAG_HMAC, HMAC_SIZE,
+    };
 
     fn default_params() -> Params {
         Params {