From e908c4ce4f7ba5fe207acb1fe841349813bbe8f1 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 00:27:45 +0700 Subject: [PATCH] refactor(proto): centralize packet envelope handling --- crates/irtt-proto/src/close.rs | 33 +++--- crates/irtt-proto/src/echo.rs | 94 ++++++++-------- crates/irtt-proto/src/envelope.rs | 180 ++++++++++++++++++++++++++++++ crates/irtt-proto/src/lib.rs | 3 + crates/irtt-proto/src/open.rs | 83 +++++--------- 5 files changed, 273 insertions(+), 120 deletions(-) create mode 100644 crates/irtt-proto/src/envelope.rs diff --git a/crates/irtt-proto/src/close.rs b/crates/irtt-proto/src/close.rs index 4f876d9..f1a8d6b 100644 --- a/crates/irtt-proto/src/close.rs +++ b/crates/irtt-proto/src/close.rs @@ -1,6 +1,8 @@ use crate::{ - flags::{FLAG_CLOSE, FLAG_HMAC}, - hmac, write_header, Result, HMAC_SIZE, + envelope::{self, FlagRule}, + 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> { - let mut flags = FLAG_CLOSE; - if hmac_key.is_some() { - flags |= FLAG_HMAC; - } - - let mut out = Vec::new(); - write_header(&mut out, flags); - if hmac_key.is_some() { - out.extend_from_slice(&[0; HMAC_SIZE]); - } + let mut out = envelope::begin( + FLAG_CLOSE, + hmac_key, + &[ + FlagRule::Require(FLAG_CLOSE), + FlagRule::Reject(FLAG_OPEN), + FlagRule::Reject(FLAG_REPLY), + ], + PacketLayout::close_request(hmac_key.is_some()).header_len(), + )?; out.extend_from_slice(&request.token.to_le_bytes()); - if let Some(key) = hmac_key { - hmac::compute_hmac_in_place(key, &mut out, hmac::hmac_offset())?; - } - Ok(out) + envelope::finish(out, hmac_key) } #[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() { diff --git a/crates/irtt-proto/src/echo.rs b/crates/irtt-proto/src/echo.rs index 04bf2ed..56983e3 100644 --- a/crates/irtt-proto/src/echo.rs +++ b/crates/irtt-proto/src/echo.rs @@ -1,11 +1,9 @@ use crate::{ - flags::{has, FLAG_HMAC, FLAG_OPEN, FLAG_REPLY}, - hmac, + envelope::{self, FlagRule}, + 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, - RECV_WINDOW_SIZE, SEQ_SIZE, TIMESTAMP_SIZE, TOKEN_SIZE, + ProtoError, Result, RECV_COUNT_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> { - 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); - write_header(&mut out, flags); - if hmac_key.is_some() { - out.extend_from_slice(&[0; HMAC_SIZE]); - } + let mut out = envelope::begin( + 0, + hmac_key, + &[ + 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 { - hmac::compute_hmac_in_place(key, &mut out, hmac::hmac_offset())?; - } - Ok(out) + envelope::finish(out, hmac_key) } pub fn decode_echo_reply( @@ -74,35 +70,17 @@ pub fn decode_echo_reply( params: &Params, hmac_key: Option<&[u8]>, ) -> Result { - let flags = validate_header(packet)?; - reject(flags, FLAG_OPEN)?; - crate::open::require(flags, FLAG_REPLY)?; - check_hmac_presence(flags, hmac_key)?; + let envelope = envelope::decode( + packet, + 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 { - 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 mut pos = envelope.body_offset; 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) { 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 { diff --git a/crates/irtt-proto/src/envelope.rs b/crates/irtt-proto/src/envelope.rs new file mode 100644 index 0000000..cce9920 --- /dev/null +++ b/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 { + 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> { + // 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, hmac_key: Option<&[u8]>) -> Result> { + 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 { + 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); + } +} diff --git a/crates/irtt-proto/src/lib.rs b/crates/irtt-proto/src/lib.rs index e185347..c72a6c5 100644 --- a/crates/irtt-proto/src/lib.rs +++ b/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; diff --git a/crates/irtt-proto/src/open.rs b/crates/irtt-proto/src/open.rs index 23c1940..aeb5261 100644 --- a/crates/irtt-proto/src/open.rs +++ b/crates/irtt-proto/src/open.rs @@ -1,8 +1,9 @@ use crate::{ - flags::{has, FLAG_CLOSE, FLAG_HMAC, FLAG_OPEN, FLAG_REPLY}, - hmac, + envelope::{self, FlagRule}, + 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() { - flags |= FLAG_HMAC; - } - - let mut out = Vec::new(); - write_header(&mut out, flags); - if hmac_key.is_some() { - out.extend_from_slice(&[0; HMAC_SIZE]); - } - 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) + let params = request.params.encode(); + let mut out = envelope::begin( + flags, + hmac_key, + &[FlagRule::Require(FLAG_OPEN), FlagRule::Reject(FLAG_REPLY)], + PacketLayout::open_request(hmac_key.is_some()).header_len() + params.len(), + )?; + out.extend_from_slice(¶ms); + envelope::finish(out, hmac_key) } pub fn decode_open_reply(packet: &[u8], hmac_key: Option<&[u8]>) -> Result { - let flags = validate_header(packet)?; - require(flags, FLAG_OPEN)?; - require(flags, FLAG_REPLY)?; - check_hmac_presence(flags, hmac_key)?; - if let Some(key) = hmac_key { - hmac::verify_hmac(key, packet, hmac::hmac_offset())?; - } + let envelope = envelope::decode( + packet, + hmac_key, + &[FlagRule::Require(FLAG_OPEN), FlagRule::Require(FLAG_REPLY)], + )?; + envelope::verify(packet, hmac_key)?; - let mut pos = HEADER_SIZE; - if has(flags, FLAG_HMAC) { - pos += HMAC_SIZE; - } + let mut pos = envelope.body_offset; 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 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 {