refactor(session): prepare close transactions

crates/irtt-client/src/session/machine.rs

@@ -53,6 +53,7 @@ enum CloseSource {
 struct ActiveSession {
     token: u64,
     negotiated: NegotiatedParams,
+    local_close_packet: Box<[u8]>,
     next_wire_seq: u32,
     highest_received_seq: Option<u32>,
     packets_sent: u64,
@@ -76,7 +77,6 @@ pub(crate) enum OpenDatagramDisposition {
 pub(crate) struct PreparedOpenAcceptance {
     next_state: MachineState,
     outcome: OpenOutcome,
-    cleanup_close: Option<Box<[u8]>>,
 }
 
 impl PreparedOpenAcceptance {
@@ -91,7 +91,13 @@ impl PreparedOpenAcceptance {
     }
 
     pub(crate) fn cleanup_close_packet(&self) -> Option<&[u8]> {
-        self.cleanup_close.as_deref()
+        match &self.next_state {
+            MachineState::Open(session) => Some(&session.local_close_packet),
+            MachineState::NoTestCompleted => None,
+            MachineState::Connected | MachineState::Closed { .. } => {
+                unreachable!("open acceptance only prepares open or no-test state")
+            }
+        }
     }
 
     #[cfg(test)]
@@ -100,6 +106,18 @@ impl PreparedOpenAcceptance {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct PreparedClose<'a> {
+    pub(crate) bytes: &'a [u8],
+    pub(crate) commit: CloseCommit,
+}
+
+#[derive(Debug)]
+pub(crate) struct CloseCommit {
+    packets_sent: u64,
+    event: ClientEvent,
+}
+
 #[derive(Debug)]
 pub(crate) struct OpenAcceptanceFailure {
     pub(crate) primary: ClientError,
@@ -401,28 +419,40 @@ impl SessionMachine {
         Ok(events)
     }
 
+    pub(crate) fn prepare_close(&self) -> Result<PreparedClose<'_>, ClientError> {
+        let session = self.open_session()?;
+        Ok(PreparedClose {
+            bytes: &session.local_close_packet,
+            commit: CloseCommit {
+                packets_sent: session.packets_sent,
+                event: ClientEvent::SessionClosed {
+                    remote: self.remote,
+                    token: session.token,
+                    at: ClientTimestamp::now(),
+                },
+            },
+        })
+    }
+
+    pub(crate) fn commit_local_close(&mut self, commit: CloseCommit) -> ClientEvent {
+        debug_assert!(
+            matches!(self.state, MachineState::Open(_)),
+            "local close commits only from open state"
+        );
+        self.state = MachineState::Closed {
+            source: CloseSource::Local,
+            packets_sent: commit.packets_sent,
+        };
+        commit.event
+    }
+
     pub(crate) fn close_with<F>(&mut self, send: F) -> Result<Vec<ClientEvent>, ClientError>
     where
         F: FnOnce(&[u8]) -> Result<(), ClientError>,
     {
-        let token = match &self.state {
+        let prepared = self.prepare_close()?;
-            MachineState::Open(session) => session.token,
+        send(prepared.bytes)?;
-            MachineState::Closed { .. } => return Err(ClientError::AlreadyClosed),
+        Ok(vec![self.commit_local_close(prepared.commit)])
-            MachineState::Connected | MachineState::NoTestCompleted => {
-                return Err(ClientError::NotOpen)
-            }
-        };
-
-        let packet =
-            encode_close_request(&CloseRequest { token }, self.config.hmac_key.as_deref())?;
-        send(&packet)?;
-        self.transition_to_closed(CloseSource::Local);
-
-        Ok(vec![ClientEvent::SessionClosed {
-            remote: self.remote,
-            token,
-            at: ClientTimestamp::now(),
-        }])
     }
 
     pub(crate) fn pending_is_empty(&self) -> bool {
@@ -488,9 +518,12 @@ impl SessionMachine {
             Ok(negotiated) => negotiated,
             Err(primary) => return Err(OpenAcceptanceFailure::new(primary, cleanup_close)),
         };
+        let local_close_packet =
+            cleanup_close.expect("normal non-zero-token replies prepare a cleanup close");
         let next_state = MachineState::Open(Box::new(ActiveSession {
             token,
             negotiated: negotiated.clone(),
+            local_close_packet,
             next_wire_seq: 0,
             highest_received_seq: None,
             packets_sent: 0,
@@ -515,7 +548,6 @@ impl SessionMachine {
         Ok(PreparedOpenAcceptance {
             next_state,
             outcome,
-            cleanup_close,
         })
     }
 
@@ -543,7 +575,6 @@ impl SessionMachine {
         Ok(PreparedOpenAcceptance {
             next_state: MachineState::NoTestCompleted,
             outcome,
-            cleanup_close: None,
         })
     }
 
@@ -981,6 +1012,14 @@ mod tests {
         machine.state = MachineState::Open(Box::new(ActiveSession {
             token: 0x0102_0304_0506_0708,
             negotiated,
+            local_close_packet: encode_close_request(
+                &CloseRequest {
+                    token: 0x0102_0304_0506_0708,
+                },
+                None,
+            )
+            .unwrap()
+            .into_boxed_slice(),
             next_wire_seq: 0,
             highest_received_seq: None,
             packets_sent: 0,
@@ -1268,6 +1307,7 @@ mod tests {
             )
             .unwrap();
 
+        let expected_close = prepared.cleanup_close_packet().unwrap().to_vec();
         let outcome = machine.commit_open(prepared);
 
         assert!(matches!(
@@ -1278,6 +1318,7 @@ mod tests {
             } if outcome_token == token
         ));
         assert!(machine.is_open());
+        assert_eq!(active(&machine).local_close_packet.as_ref(), expected_close);
         assert!(matches!(
             machine.prepare_open_request(),
             Err(ClientError::AlreadyOpen)
@@ -1337,6 +1378,129 @@ mod tests {
         assert!(matches!(machine.state, MachineState::Connected));
     }
 
+    #[test]
+    fn prepared_close_is_stable_inert_and_commits_preserved_state() {
+        let mut machine = open_machine(4, Duration::from_secs(1));
+        let remote = machine.remote;
+        active_mut(&mut machine).packets_sent = 7;
+
+        let first_bytes = {
+            let prepared = machine.prepare_close().unwrap();
+            assert_eq!(
+                prepared.bytes,
+                encode_close_request(
+                    &CloseRequest {
+                        token: 0x0102_0304_0506_0708,
+                    },
+                    None,
+                )
+                .unwrap()
+            );
+            prepared.bytes.to_vec()
+        };
+        assert!(machine.is_open());
+        assert_eq!(active(&machine).packets_sent, 7);
+
+        let prepared = machine.prepare_close().unwrap();
+        assert_eq!(prepared.bytes, first_bytes);
+        let event = machine.commit_local_close(prepared.commit);
+
+        assert!(matches!(
+            event,
+            ClientEvent::SessionClosed {
+                remote: event_remote,
+                token: 0x0102_0304_0506_0708,
+                ..
+            } if event_remote == remote
+        ));
+        assert!(matches!(
+            machine.state,
+            MachineState::Closed {
+                source: CloseSource::Local,
+                packets_sent: 7,
+            }
+        ));
+        assert!(matches!(
+            machine.prepare_close(),
+            Err(ClientError::AlreadyClosed)
+        ));
+    }
+
+    #[test]
+    fn prepared_close_reuses_hmac_packet_built_during_open_acceptance() {
+        let key = b"close-key".to_vec();
+        let mut machine = connected_machine(ClientConfig {
+            hmac_key: Some(key.clone()),
+            ..ClientConfig::default()
+        });
+        let token = 0x1020_3040_5060_7080;
+        let prepared_open = machine
+            .prepare_open_acceptance(
+                normal_open_reply(&machine, token),
+                timestamp(Instant::now()),
+            )
+            .unwrap();
+        let opening_close = prepared_open.cleanup_close_packet().unwrap().to_vec();
+        machine.commit_open(prepared_open);
+
+        let prepared_close = machine.prepare_close().unwrap();
+        assert_eq!(prepared_close.bytes, opening_close);
+        assert_eq!(
+            prepared_close.bytes[3],
+            flags::FLAG_CLOSE | flags::FLAG_HMAC
+        );
+        irtt_proto::verify_hmac(&key, prepared_close.bytes, 4).unwrap();
+        assert_eq!(
+            u64::from_le_bytes(
+                prepared_close.bytes[4 + irtt_proto::HMAC_SIZE..12 + irtt_proto::HMAC_SIZE]
+                    .try_into()
+                    .unwrap()
+            ),
+            token
+        );
+    }
+
+    #[test]
+    fn close_preparation_preserves_non_open_errors_and_peer_source() {
+        let connected = connected_machine(ClientConfig::default());
+        assert!(matches!(
+            connected.prepare_close(),
+            Err(ClientError::NotOpen)
+        ));
+
+        let mut no_test = connected_machine(ClientConfig {
+            run_mode: RunMode::NoTest,
+            ..ClientConfig::default()
+        });
+        let reply = OpenReply {
+            flags: flags::FLAG_OPEN | flags::FLAG_REPLY | flags::FLAG_CLOSE,
+            token: 0,
+            params: no_test.requested.clone(),
+        };
+        let prepared = no_test
+            .prepare_open_acceptance(reply, timestamp(Instant::now()))
+            .unwrap();
+        no_test.commit_open(prepared);
+        assert!(matches!(
+            no_test.prepare_close(),
+            Err(ClientError::AlreadyCompleted)
+        ));
+
+        let mut peer_closed = open_machine(4, Duration::from_secs(1));
+        peer_closed.transition_to_closed(CloseSource::Peer);
+        assert!(matches!(
+            peer_closed.state,
+            MachineState::Closed {
+                source: CloseSource::Peer,
+                ..
+            }
+        ));
+        assert!(matches!(
+            peer_closed.prepare_close(),
+            Err(ClientError::AlreadyClosed)
+        ));
+    }
+
     #[test]
     fn dropping_prepared_probe_changes_nothing() {
         let machine = open_machine(4, Duration::from_secs(1));