From cc7f7fc9b2da4f13e4f12dd5a48a3cf6242eba3a Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 00:30:08 +0700 Subject: [PATCH] feat(proto): add directional open and close codecs --- crates/irtt-proto/src/close.rs | 33 ++++- crates/irtt-proto/src/lib.rs | 7 +- crates/irtt-proto/src/open.rs | 31 ++++ crates/irtt-proto/tests/directional_codecs.rs | 135 ++++++++++++++++++ crates/irtt-proto/tests/vectors.rs | 41 +++++- 5 files changed, 240 insertions(+), 7 deletions(-) create mode 100644 crates/irtt-proto/tests/directional_codecs.rs diff --git a/crates/irtt-proto/src/close.rs b/crates/irtt-proto/src/close.rs index f1a8d6b..9e031f0 100644 --- a/crates/irtt-proto/src/close.rs +++ b/crates/irtt-proto/src/close.rs @@ -1,8 +1,8 @@ use crate::{ envelope::{self, FlagRule}, - flags::{FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY}, + flags::{has, FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY}, layout::PacketLayout, - Result, + ProtoError, Result, TOKEN_SIZE, }; #[derive(Debug, Clone, PartialEq, Eq)] @@ -25,6 +25,35 @@ pub fn encode_close_request(request: &CloseRequest, hmac_key: Option<&[u8]>) -> envelope::finish(out, hmac_key) } +pub fn decode_close_request(packet: &[u8], hmac_key: Option<&[u8]>) -> Result { + let envelope = envelope::decode( + packet, + hmac_key, + &[ + FlagRule::Require(FLAG_CLOSE), + FlagRule::Reject(FLAG_OPEN), + FlagRule::Reject(FLAG_REPLY), + ], + )?; + let expected_len = + PacketLayout::close_request(has(envelope.flags, crate::FLAG_HMAC)).header_len(); + if packet.len() != expected_len { + return Err(ProtoError::PacketLengthMismatch { + expected: expected_len, + actual: packet.len(), + }); + } + envelope::verify(packet, hmac_key)?; + + Ok(CloseRequest { + token: u64::from_le_bytes( + packet[envelope.body_offset..envelope.body_offset + TOKEN_SIZE] + .try_into() + .unwrap(), + ), + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/irtt-proto/src/lib.rs b/crates/irtt-proto/src/lib.rs index c72a6c5..a4d007c 100644 --- a/crates/irtt-proto/src/lib.rs +++ b/crates/irtt-proto/src/lib.rs @@ -31,13 +31,16 @@ pub mod open; pub mod params; pub mod varint; -pub use close::{encode_close_request, CloseRequest}; +pub use close::{decode_close_request, encode_close_request, CloseRequest}; pub use echo::{decode_echo_reply, encode_echo_request, EchoReply, EchoRequest, TimestampFields}; pub use error::{ProtoError, Result}; pub use flags::*; pub use hmac::{compute_hmac, compute_hmac_in_place, verify_hmac}; pub use layout::{echo_header_len, echo_packet_len, PacketLayout}; -pub use open::{decode_open_reply, encode_open_request, OpenReply, OpenRequest}; +pub use open::{ + decode_open_reply, decode_open_request, encode_open_reply, encode_open_request, OpenReply, + OpenRequest, +}; pub use params::{Clock, Params, ReceivedStats, ServerFill, StampAt, MAX_SERVER_FILL_BYTES}; pub const MAGIC: [u8; 3] = [0x14, 0xA7, 0x5B]; diff --git a/crates/irtt-proto/src/open.rs b/crates/irtt-proto/src/open.rs index aeb5261..d48f041 100644 --- a/crates/irtt-proto/src/open.rs +++ b/crates/irtt-proto/src/open.rs @@ -35,6 +35,37 @@ pub fn encode_open_request(request: &OpenRequest, hmac_key: Option<&[u8]>) -> Re envelope::finish(out, hmac_key) } +pub fn decode_open_request(packet: &[u8], hmac_key: Option<&[u8]>) -> Result { + let envelope = envelope::decode( + packet, + hmac_key, + &[FlagRule::Require(FLAG_OPEN), FlagRule::Reject(FLAG_REPLY)], + )?; + envelope::verify(packet, hmac_key)?; + + Ok(OpenRequest { + params: Params::decode(&packet[envelope.body_offset..])?, + close: has(envelope.flags, FLAG_CLOSE), + }) +} + +pub fn encode_open_reply(reply: &OpenReply, hmac_key: Option<&[u8]>) -> Result> { + let params = reply.params.encode(); + let mut out = envelope::begin( + reply.flags, + hmac_key, + &[FlagRule::Require(FLAG_OPEN), FlagRule::Require(FLAG_REPLY)], + PacketLayout::open_reply(hmac_key.is_some()).header_len() + params.len(), + )?; + if reply.token == 0 && !has(reply.flags, FLAG_CLOSE) { + return Err(ProtoError::ZeroToken); + } + + out.extend_from_slice(&reply.token.to_le_bytes()); + out.extend_from_slice(¶ms); + envelope::finish(out, hmac_key) +} + pub fn decode_open_reply(packet: &[u8], hmac_key: Option<&[u8]>) -> Result { let envelope = envelope::decode( packet, diff --git a/crates/irtt-proto/tests/directional_codecs.rs b/crates/irtt-proto/tests/directional_codecs.rs new file mode 100644 index 0000000..cf985c6 --- /dev/null +++ b/crates/irtt-proto/tests/directional_codecs.rs @@ -0,0 +1,135 @@ +use irtt_proto::{ + decode_close_request, decode_open_reply, decode_open_request, encode_close_request, + encode_open_reply, encode_open_request, CloseRequest, OpenReply, OpenRequest, Params, + ProtoError, FLAG_CLOSE, FLAG_HMAC, FLAG_OPEN, FLAG_REPLY, MAGIC, +}; + +const KEY: &[u8] = b"testkey"; +const TOKEN: u64 = 0x7896_b6ab_8771_5213; + +#[test] +fn authenticated_open_reply_encoder_sets_hmac_from_key() { + let reply = OpenReply { + flags: FLAG_OPEN | FLAG_REPLY, + token: TOKEN, + params: Params { + protocol_version: 1, + ..Params::default() + }, + }; + let packet = encode_open_reply(&reply, Some(KEY)).unwrap(); + + assert_eq!(packet[3], FLAG_OPEN | FLAG_REPLY | FLAG_HMAC); + assert_eq!( + decode_open_reply(&packet, Some(KEY)), + Ok(OpenReply { + flags: FLAG_OPEN | FLAG_REPLY | FLAG_HMAC, + ..reply + }) + ); +} + +#[test] +fn rejected_open_reply_round_trips_and_zero_token_requires_close() { + let rejected = OpenReply { + flags: FLAG_OPEN | FLAG_REPLY | FLAG_CLOSE, + token: 0, + params: Params { + protocol_version: 1, + ..Params::default() + }, + }; + let packet = encode_open_reply(&rejected, None).unwrap(); + assert_eq!(decode_open_reply(&packet, None), Ok(rejected.clone())); + + assert_eq!( + encode_open_reply( + &OpenReply { + flags: FLAG_OPEN | FLAG_REPLY, + ..rejected + }, + None, + ), + Err(ProtoError::ZeroToken) + ); +} + +#[test] +fn open_and_close_codecs_enforce_directional_flags() { + let params = Params::default(); + let open_reply_packet = encode_open_reply( + &OpenReply { + flags: FLAG_OPEN | FLAG_REPLY, + token: TOKEN, + params: params.clone(), + }, + None, + ) + .unwrap(); + assert_eq!( + decode_open_request(&open_reply_packet, None), + Err(ProtoError::UnexpectedFlag(FLAG_REPLY)) + ); + assert_eq!( + encode_open_reply( + &OpenReply { + flags: FLAG_REPLY, + token: TOKEN, + params: params.clone(), + }, + None, + ), + Err(ProtoError::MissingFlag(FLAG_OPEN)) + ); + + let no_test_packet = encode_open_request( + &OpenRequest { + params, + close: true, + }, + None, + ) + .unwrap(); + assert_eq!( + decode_close_request(&no_test_packet, None), + Err(ProtoError::UnexpectedFlag(FLAG_OPEN)) + ); +} + +#[test] +fn open_and_close_malformed_bodies_and_exact_lengths_are_rejected() { + let mut malformed_open_request = MAGIC.to_vec(); + malformed_open_request.push(FLAG_OPEN); + malformed_open_request.extend_from_slice(&[1, 0x80]); + assert_eq!( + decode_open_request(&malformed_open_request, None), + Err(ProtoError::TruncatedVarint) + ); + + let mut malformed_open_reply = MAGIC.to_vec(); + malformed_open_reply.push(FLAG_OPEN | FLAG_REPLY); + malformed_open_reply.extend_from_slice(&TOKEN.to_le_bytes()); + malformed_open_reply.extend_from_slice(&[1, 0x80]); + assert_eq!( + decode_open_reply(&malformed_open_reply, None), + Err(ProtoError::TruncatedVarint) + ); + + let close_packet = encode_close_request(&CloseRequest { token: TOKEN }, None).unwrap(); + assert_eq!( + decode_close_request(&close_packet[..11], None), + Err(ProtoError::PacketLengthMismatch { + expected: 12, + actual: 11, + }) + ); + let mut long_close_packet = close_packet; + long_close_packet.push(0); + assert_eq!( + decode_close_request(&long_close_packet, None), + Err(ProtoError::PacketLengthMismatch { + expected: 12, + actual: 13, + }) + ); +} diff --git a/crates/irtt-proto/tests/vectors.rs b/crates/irtt-proto/tests/vectors.rs index 93bdf5a..8d5dea9 100644 --- a/crates/irtt-proto/tests/vectors.rs +++ b/crates/irtt-proto/tests/vectors.rs @@ -1,7 +1,7 @@ use irtt_proto::{ - decode_echo_reply, decode_open_reply, encode_close_request, encode_echo_request, - encode_open_request, verify_hmac, Clock, CloseRequest, EchoRequest, OpenRequest, Params, - ReceivedStats, StampAt, + decode_close_request, decode_echo_reply, decode_open_reply, decode_open_request, + encode_close_request, encode_echo_request, encode_open_reply, encode_open_request, verify_hmac, + Clock, CloseRequest, EchoRequest, OpenRequest, Params, ReceivedStats, StampAt, }; fn hex(input: &str) -> Vec { @@ -50,6 +50,13 @@ fn vector_1_open_request_no_hmac() { ) .unwrap(); assert_eq!(packet, expected); + assert_eq!( + decode_open_request(&expected, None), + Ok(OpenRequest { + params: default_params_3s(), + close: false, + }) + ); } #[test] @@ -59,6 +66,7 @@ fn vector_2_open_reply_no_hmac() { let reply = decode_open_reply(&packet, None).unwrap(); assert_eq!(reply.token, 0x7896_b6ab_8771_5213); assert_eq!(reply.params, default_params_3s()); + assert_eq!(encode_open_reply(&reply, None).unwrap(), packet); } #[test] @@ -103,6 +111,12 @@ fn vector_5_close_request_no_hmac() { ) .unwrap(); assert_eq!(packet, expected); + assert_eq!( + decode_close_request(&expected, None), + Ok(CloseRequest { + token: 0x7896_b6ab_8771_5213, + }) + ); } #[test] @@ -120,6 +134,13 @@ fn vector_6_hmac_open_request() { .unwrap(); assert_eq!(packet, expected); verify_hmac(b"testkey", &packet, 4).unwrap(); + assert_eq!( + decode_open_request(&expected, Some(b"testkey")), + Ok(OpenRequest { + params: hmac_params_2s(), + close: false, + }) + ); } #[test] @@ -154,6 +175,12 @@ fn vector_8_hmac_close_request() { ) .unwrap(); assert_eq!(packet, expected); + assert_eq!( + decode_close_request(&expected, Some(b"testkey")), + Ok(CloseRequest { + token: 0x4387_e9eb_5d3f_ca59, + }) + ); } #[test] @@ -169,6 +196,13 @@ fn vector_9_no_test_open_close_request() { ) .unwrap(); assert_eq!(packet, expected); + assert_eq!( + decode_open_request(&expected, None), + Ok(OpenRequest { + params: default_params_60s(), + close: true, + }) + ); } #[test] @@ -179,6 +213,7 @@ fn vector_10_no_test_open_close_reply() { let reply = decode_open_reply(&packet, None).unwrap(); assert_eq!(reply.token, 0); assert_eq!(reply.params, default_params_60s()); + assert_eq!(encode_open_reply(&reply, None).unwrap(), packet); } #[test]