From dad4d2a29812a696c7e6992894d79da65b8df818 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 01:23:03 +0700 Subject: [PATCH] refactor(session): prepare opening transactions --- crates/irtt-client/src/session/machine.rs | 513 ++++++++++++++++++++-- 1 file changed, 465 insertions(+), 48 deletions(-) diff --git a/crates/irtt-client/src/session/machine.rs b/crates/irtt-client/src/session/machine.rs index 2c53753..d14755c 100644 --- a/crates/irtt-client/src/session/machine.rs +++ b/crates/irtt-client/src/session/machine.rs @@ -61,6 +61,71 @@ struct ActiveSession { completed: CompletedSet, } +#[derive(Debug)] +pub(crate) struct PreparedOpenRequest { + pub(crate) bytes: Box<[u8]>, +} + +#[derive(Debug)] +pub(crate) enum OpenDatagramDisposition { + Ignore, + Trusted(OpenReply), +} + +#[derive(Debug)] +pub(crate) struct PreparedOpenAcceptance { + next_state: MachineState, + outcome: OpenOutcome, + cleanup_close: Option>, +} + +impl PreparedOpenAcceptance { + pub(crate) fn negotiated(&self) -> &NegotiatedParams { + match &self.outcome { + OpenOutcome::Started { negotiated, .. } + | OpenOutcome::NoTestCompleted { negotiated, .. } => negotiated, + } + } + + pub(crate) fn normal_negotiated(&self) -> Option<&NegotiatedParams> { + match &self.next_state { + MachineState::Open(session) => Some(&session.negotiated), + MachineState::NoTestCompleted => None, + MachineState::Connected | MachineState::Closed { .. } => { + unreachable!("open acceptance only prepares open or no-test state") + } + } + } + + pub(crate) fn cleanup_close_packet(&self) -> Option<&[u8]> { + self.cleanup_close.as_deref() + } + + #[cfg(test)] + pub(crate) fn has_prepared_active_session(&self) -> bool { + matches!(self.next_state, MachineState::Open(_)) + } +} + +#[derive(Debug)] +pub(crate) struct OpenAcceptanceFailure { + pub(crate) primary: ClientError, + pub(crate) cleanup_close: Option>, +} + +impl OpenAcceptanceFailure { + fn new(primary: ClientError, cleanup_close: Option>) -> Self { + Self { + primary, + cleanup_close, + } + } + + fn without_cleanup(primary: ClientError) -> Self { + Self::new(primary, None) + } +} + #[derive(Debug)] pub(crate) struct PreparedProbe { pub(crate) bytes: Box<[u8]>, @@ -121,22 +186,106 @@ impl SessionMachine { self.config.hmac_key.is_some() } - pub(crate) fn open_packet(&self) -> Result, ClientError> { - match self.state { - MachineState::Connected => {} - MachineState::Open(_) => return Err(ClientError::AlreadyOpen), - MachineState::Closed { .. } => return Err(ClientError::AlreadyClosed), - MachineState::NoTestCompleted => return Err(ClientError::AlreadyCompleted), - } - + pub(crate) fn prepare_open_request(&self) -> Result { + self.ensure_connected()?; let request = OpenRequest { params: self.requested.clone(), close: self.config.run_mode == RunMode::NoTest, }; - Ok(encode_open_request( - &request, - self.config.hmac_key.as_deref(), - )?) + let bytes = encode_open_request(&request, self.config.hmac_key.as_deref())?; + Ok(PreparedOpenRequest { + bytes: bytes.into_boxed_slice(), + }) + } + + pub(crate) fn inspect_open_datagram( + &self, + packet: &[u8], + ) -> Result { + self.ensure_connected()?; + match irtt_proto::decode_open_reply(packet, self.config.hmac_key.as_deref()) { + Ok(reply) => Ok(OpenDatagramDisposition::Trusted(reply)), + Err(irtt_proto::ProtoError::ZeroToken) => Err(ClientError::ZeroToken), + Err( + error @ (irtt_proto::ProtoError::TruncatedVarint + | irtt_proto::ProtoError::VarintOverflow + | irtt_proto::ProtoError::InvalidUtf8 + | irtt_proto::ProtoError::InvalidEnum { .. } + | irtt_proto::ProtoError::NegativePacketLength { .. } + | irtt_proto::ProtoError::ParameterLengthTooLarge { .. } + | irtt_proto::ProtoError::MalformedParams), + ) => Err(ClientError::Protocol(error)), + Err(_) => Ok(OpenDatagramDisposition::Ignore), + } + } + + pub(crate) fn prepare_open_acceptance( + &self, + reply: OpenReply, + now: ClientTimestamp, + ) -> Result { + self.ensure_connected() + .map_err(OpenAcceptanceFailure::without_cleanup)?; + + let reply_is_close = flags::has(reply.flags, flags::FLAG_CLOSE); + let cleanup_close = + if self.config.run_mode == RunMode::Normal && !reply_is_close && reply.token != 0 { + let bytes = encode_close_request( + &CloseRequest { token: reply.token }, + self.config.hmac_key.as_deref(), + ) + .map_err(ClientError::from) + .map_err(OpenAcceptanceFailure::without_cleanup)?; + Some(bytes.into_boxed_slice()) + } else { + None + }; + + if reply.params.protocol_version != PROTOCOL_VERSION { + return Err(OpenAcceptanceFailure::new( + ClientError::ProtocolVersionMismatch { + requested: PROTOCOL_VERSION, + received: reply.params.protocol_version, + }, + cleanup_close, + )); + } + + match self.config.run_mode { + RunMode::Normal if reply_is_close => Err(OpenAcceptanceFailure::new( + ClientError::ServerRejected, + cleanup_close, + )), + RunMode::Normal if reply.token == 0 => Err(OpenAcceptanceFailure::new( + ClientError::ZeroToken, + cleanup_close, + )), + RunMode::Normal => self.prepare_normal_open(reply, now, cleanup_close), + RunMode::NoTest if !reply_is_close => Err(OpenAcceptanceFailure::new( + ClientError::UnexpectedNoTestReply, + cleanup_close, + )), + RunMode::NoTest if reply.token != 0 => Err(OpenAcceptanceFailure::new( + ClientError::NonZeroNoTestToken { token: reply.token }, + cleanup_close, + )), + RunMode::NoTest => self.prepare_no_test_open(reply, now), + } + } + + pub(crate) fn commit_open(&mut self, prepared: PreparedOpenAcceptance) -> OpenOutcome { + debug_assert!( + matches!(self.state, MachineState::Connected), + "open acceptance commits only from connected state" + ); + self.state = prepared.next_state; + prepared.outcome + } + + // Temporary private bridges keep this first commit independently compiling. + // The blocking and managed migrations remove them before the PR is complete. + pub(crate) fn open_packet(&self) -> Result, ClientError> { + Ok(self.prepare_open_request()?.bytes.into_vec()) } pub(crate) fn decode_open_reply(&self, packet: &[u8]) -> Result { @@ -155,24 +304,13 @@ impl SessionMachine { where F: FnOnce(&NegotiatedParams) -> Result<(), ClientError>, { - if reply.params.protocol_version != PROTOCOL_VERSION { - return Err(ClientError::ProtocolVersionMismatch { - requested: PROTOCOL_VERSION, - received: reply.params.protocol_version, - }); - } - - let reply_is_close = flags::has(reply.flags, flags::FLAG_CLOSE); - match self.config.run_mode { - RunMode::Normal if reply_is_close => Err(ClientError::ServerRejected), - RunMode::Normal if reply.token == 0 => Err(ClientError::ZeroToken), - RunMode::Normal => self.accept_normal_open(reply, now, before_normal_open), - RunMode::NoTest if !reply_is_close => Err(ClientError::UnexpectedNoTestReply), - RunMode::NoTest if reply.token != 0 => { - Err(ClientError::NonZeroNoTestToken { token: reply.token }) - } - RunMode::NoTest => self.accept_no_test_open(reply, now), + let prepared = self + .prepare_open_acceptance(reply, now) + .map_err(|failure| failure.primary)?; + if let Some(negotiated) = prepared.normal_negotiated() { + before_normal_open(negotiated)?; } + Ok(self.commit_open(prepared)) } pub(crate) fn probe_timeout(&self) -> Duration { @@ -374,24 +512,22 @@ impl SessionMachine { self.open_session().map(|_| ()) } - fn accept_normal_open( - &mut self, + fn prepare_normal_open( + &self, reply: OpenReply, now: ClientTimestamp, - before_normal_open: F, - ) -> Result - where - F: FnOnce(&NegotiatedParams) -> Result<(), ClientError>, - { + cleanup_close: Option>, + ) -> Result { let token = reply.token; - let negotiated = negotiate_params( + let negotiated = match negotiate_params( &self.requested, reply.params, self.config.negotiation_policy, - )?; - before_normal_open(&negotiated)?; - - self.state = MachineState::Open(Box::new(ActiveSession { + ) { + Ok(negotiated) => negotiated, + Err(primary) => return Err(OpenAcceptanceFailure::new(primary, cleanup_close)), + }; + let next_state = MachineState::Open(Box::new(ActiveSession { token, negotiated: negotiated.clone(), next_wire_seq: 0, @@ -409,34 +545,44 @@ impl SessionMachine { at: now, }; - Ok(OpenOutcome::Started { + let outcome = OpenOutcome::Started { remote: self.remote, token, negotiated, event, + }; + Ok(PreparedOpenAcceptance { + next_state, + outcome, + cleanup_close, }) } - fn accept_no_test_open( - &mut self, + fn prepare_no_test_open( + &self, reply: OpenReply, now: ClientTimestamp, - ) -> Result { + ) -> Result { let negotiated = negotiate_params( &self.requested, reply.params, self.config.negotiation_policy, - )?; - self.state = MachineState::NoTestCompleted; + ) + .map_err(OpenAcceptanceFailure::without_cleanup)?; let event = ClientEvent::NoTestCompleted { remote: self.remote, negotiated: negotiated.clone(), at: now, }; - Ok(OpenOutcome::NoTestCompleted { + let outcome = OpenOutcome::NoTestCompleted { remote: self.remote, negotiated, event, + }; + Ok(PreparedOpenAcceptance { + next_state: MachineState::NoTestCompleted, + outcome, + cleanup_close: None, }) } @@ -620,6 +766,15 @@ impl SessionMachine { } } + fn ensure_connected(&self) -> Result<(), ClientError> { + match self.state { + MachineState::Connected => Ok(()), + MachineState::Open(_) => Err(ClientError::AlreadyOpen), + MachineState::Closed { .. } => Err(ClientError::AlreadyClosed), + MachineState::NoTestCompleted => Err(ClientError::AlreadyCompleted), + } + } + fn open_session_mut(&mut self) -> Result<&mut ActiveSession, ClientError> { match &mut self.state { MachineState::Open(session) => Ok(session), @@ -896,6 +1051,268 @@ mod tests { } } + fn connected_machine(config: ClientConfig) -> SessionMachine { + SessionMachine::new(config, "127.0.0.1:2112".parse().unwrap()).unwrap() + } + + fn normal_open_reply(machine: &SessionMachine, token: u64) -> OpenReply { + OpenReply { + flags: flags::FLAG_OPEN | flags::FLAG_REPLY, + token, + params: machine.requested.clone(), + } + } + + fn encoded_open_reply( + machine: &SessionMachine, + reply: &OpenReply, + ) -> Result, irtt_proto::ProtoError> { + irtt_proto::encode_open_reply(reply, machine.config.hmac_key.as_deref()) + } + + #[test] + fn prepared_open_request_is_exact_and_inert_when_dropped() { + let machine = connected_machine(ClientConfig::default()); + let first = machine.prepare_open_request().unwrap(); + let second = machine.prepare_open_request().unwrap(); + + assert_eq!(first.bytes, second.bytes); + assert!(!first.bytes.is_empty()); + drop(first); + drop(second); + assert!(matches!(machine.state, MachineState::Connected)); + } + + #[test] + fn malformed_wrong_direction_and_invalid_flags_are_ignored() { + let machine = connected_machine(ClientConfig::default()); + assert!(matches!( + machine.inspect_open_datagram(&[0_u8]), + Ok(OpenDatagramDisposition::Ignore) + )); + + let request = machine.prepare_open_request().unwrap(); + assert!(matches!( + machine.inspect_open_datagram(&request.bytes), + Ok(OpenDatagramDisposition::Ignore) + )); + + let reply = normal_open_reply(&machine, 0x1020_3040_5060_7080); + let mut reserved = encoded_open_reply(&machine, &reply).unwrap(); + reserved[3] |= 0x10; + assert!(matches!( + machine.inspect_open_datagram(&reserved), + Ok(OpenDatagramDisposition::Ignore) + )); + assert!(matches!(machine.state, MachineState::Connected)); + } + + #[test] + fn missing_unexpected_and_bad_hmac_are_ignored() { + let key = b"open-key".to_vec(); + let authenticated = connected_machine(ClientConfig { + hmac_key: Some(key.clone()), + ..ClientConfig::default() + }); + let reply = normal_open_reply(&authenticated, 0x1020_3040_5060_7080); + let plain = irtt_proto::encode_open_reply(&reply, None).unwrap(); + assert!(matches!( + authenticated.inspect_open_datagram(&plain), + Ok(OpenDatagramDisposition::Ignore) + )); + + let mut bad = encoded_open_reply(&authenticated, &reply).unwrap(); + *bad.last_mut().unwrap() ^= 0x80; + assert!(matches!( + authenticated.inspect_open_datagram(&bad), + Ok(OpenDatagramDisposition::Ignore) + )); + + let plain_machine = connected_machine(ClientConfig::default()); + let unexpected = irtt_proto::encode_open_reply(&reply, Some(&key)).unwrap(); + assert!(matches!( + plain_machine.inspect_open_datagram(&unexpected), + Ok(OpenDatagramDisposition::Ignore) + )); + } + + #[test] + fn authenticated_zero_token_is_trusted_and_terminal() { + let key = b"open-key".to_vec(); + let machine = connected_machine(ClientConfig { + hmac_key: Some(key.clone()), + ..ClientConfig::default() + }); + let reply = normal_open_reply(&machine, 0x1020_3040_5060_7080); + let mut packet = encoded_open_reply(&machine, &reply).unwrap(); + let token_offset = 4 + irtt_proto::HMAC_SIZE; + packet[token_offset..token_offset + 8].fill(0); + irtt_proto::compute_hmac_in_place(&key, &mut packet, 4).unwrap(); + + assert!(matches!( + machine.inspect_open_datagram(&packet), + Err(ClientError::ZeroToken) + )); + assert!(matches!(machine.state, MachineState::Connected)); + } + + #[test] + fn trusted_invalid_parameter_encoding_is_terminal() { + let machine = connected_machine(ClientConfig::default()); + let mut packet = irtt_proto::MAGIC.to_vec(); + packet.push(flags::FLAG_OPEN | flags::FLAG_REPLY); + packet.extend_from_slice(&0x1020_3040_5060_7080_u64.to_le_bytes()); + packet.push(0x80); + + assert!(matches!( + machine.inspect_open_datagram(&packet), + Err(ClientError::Protocol( + irtt_proto::ProtoError::TruncatedVarint + )) + )); + assert!(matches!(machine.state, MachineState::Connected)); + } + + #[test] + fn trusted_rejection_version_and_run_mode_failures_do_not_open() { + let machine = connected_machine(ClientConfig::default()); + let rejection = OpenReply { + flags: flags::FLAG_OPEN | flags::FLAG_REPLY | flags::FLAG_CLOSE, + token: 0, + params: machine.requested.clone(), + }; + assert!(matches!( + machine + .prepare_open_acceptance(rejection, timestamp(Instant::now())) + .unwrap_err() + .primary, + ClientError::ServerRejected + )); + + let mut version = normal_open_reply(&machine, 0x1020_3040_5060_7080); + version.params.protocol_version += 1; + assert!(matches!( + machine + .prepare_open_acceptance(version, timestamp(Instant::now())) + .unwrap_err() + .primary, + ClientError::ProtocolVersionMismatch { .. } + )); + + let no_test = connected_machine(ClientConfig { + run_mode: RunMode::NoTest, + ..ClientConfig::default() + }); + assert!(matches!( + no_test + .prepare_open_acceptance( + normal_open_reply(&no_test, 0x1020_3040_5060_7080), + timestamp(Instant::now()), + ) + .unwrap_err() + .primary, + ClientError::UnexpectedNoTestReply + )); + assert!(matches!(machine.state, MachineState::Connected)); + assert!(matches!(no_test.state, MachineState::Connected)); + } + + #[test] + fn dropping_prepared_acceptance_leaves_connected_state() { + let machine = connected_machine(ClientConfig::default()); + let prepared = machine + .prepare_open_acceptance( + normal_open_reply(&machine, 0x1020_3040_5060_7080), + timestamp(Instant::now()), + ) + .unwrap(); + + assert!(prepared.has_prepared_active_session()); + assert!(prepared.cleanup_close_packet().is_some()); + drop(prepared); + assert!(matches!(machine.state, MachineState::Connected)); + } + + #[test] + fn commit_open_assigns_prebuilt_state_once() { + let mut machine = connected_machine(ClientConfig::default()); + let token = 0x1020_3040_5060_7080; + let prepared = machine + .prepare_open_acceptance( + normal_open_reply(&machine, token), + timestamp(Instant::now()), + ) + .unwrap(); + + let outcome = machine.commit_open(prepared); + + assert!(matches!( + outcome, + OpenOutcome::Started { + token: outcome_token, + .. + } if outcome_token == token + )); + assert!(machine.is_open()); + assert!(matches!( + machine.prepare_open_request(), + Err(ClientError::AlreadyOpen) + )); + } + + #[test] + fn no_test_preparation_completes_without_active_session() { + let config = ClientConfig { + run_mode: RunMode::NoTest, + ..ClientConfig::default() + }; + let mut machine = connected_machine(config); + let reply = OpenReply { + flags: flags::FLAG_OPEN | flags::FLAG_REPLY | flags::FLAG_CLOSE, + token: 0, + params: machine.requested.clone(), + }; + let prepared = machine + .prepare_open_acceptance(reply, timestamp(Instant::now())) + .unwrap(); + + assert!(prepared.normal_negotiated().is_none()); + assert!(prepared.cleanup_close_packet().is_none()); + assert!(matches!(machine.state, MachineState::Connected)); + let outcome = machine.commit_open(prepared); + assert!(matches!(outcome, OpenOutcome::NoTestCompleted { .. })); + assert!(matches!(machine.state, MachineState::NoTestCompleted)); + } + + #[test] + fn cleanup_close_is_preencoded_with_token_and_hmac() { + let key = b"cleanup-key".to_vec(); + let machine = connected_machine(ClientConfig { + hmac_key: Some(key.clone()), + ..ClientConfig::default() + }); + let token = 0x1020_3040_5060_7080; + let prepared = machine + .prepare_open_acceptance( + normal_open_reply(&machine, token), + timestamp(Instant::now()), + ) + .unwrap(); + let cleanup = prepared.cleanup_close_packet().unwrap(); + + assert_eq!(cleanup[3], flags::FLAG_CLOSE | flags::FLAG_HMAC); + irtt_proto::verify_hmac(&key, cleanup, 4).unwrap(); + assert_eq!( + u64::from_le_bytes( + cleanup[4 + irtt_proto::HMAC_SIZE..12 + irtt_proto::HMAC_SIZE] + .try_into() + .unwrap() + ), + token + ); + assert!(matches!(machine.state, MachineState::Connected)); + } + #[test] fn dropping_prepared_probe_changes_nothing() { let machine = open_machine(4, Duration::from_secs(1));