ts_tunnel: manage event lifecycle better, and recover from the stuck responder state

ts_tunnel/src/endpoint.rs

@@ -247,6 +247,11 @@ impl Peer {
     fn cleanup_expired(&mut self, now: Instant) {
         self.check_invariants(now);
         self.session.cleanup_expired(now);
+        self.handshake.cleanup_expired(now);
+        // We may have packets queued that were waiting on session confirmation from the initiator.
+        if !self.handshake.is_active() {
+            self.queue.clear();
+        }
     }
 
     fn shutdown(&mut self) {
@@ -597,6 +602,7 @@ mod tests {
         config::PeerConfig,
         handshake::{HANDSHAKE_FAILURE_TIMEOUT, HANDSHAKE_RETRY_TIMEOUT},
         messages::{HandshakeInitiation, TransportDataHeader},
+        session::SESSION_LIFETIME,
     };
 
     /// Matches different shapes of packets.
@@ -832,7 +838,6 @@ mod tests {
         ///
         /// This may result in the queuing of encrypted packets from B to A, which can be inspected
         /// with [`EndpointPair::assert_b_to_a`].
-        #[allow(dead_code)]
         pub fn event_b(&mut self, now: Instant) {
             EndpointPair::events(now, &mut self.b, &mut self.b_to_a);
         }
@@ -843,7 +848,6 @@ mod tests {
         }
 
         /// Drop in-flight packets from B to A.
-        #[allow(dead_code)]
         pub fn drop_b_to_a(&mut self) {
             self.b_to_a.clear();
         }
@@ -1048,4 +1052,68 @@ mod tests {
         // B never receives packet(1), it was dropped when the first handshake failed.
         p.assert_received_at_b([packet(2)]);
     }
+
+    #[test]
+    fn responder_timeout() {
+        use PacketMatcher::*;
+
+        let mut p = EndpointPair::new();
+        let t = TestClock::new();
+
+        // A initiates, B responds.
+        p.send_from_a(t.at(0), [packet(1)]);
+        p.assert_a_to_b([HandshakeInitiation]);
+        p.recv_at_b(t.at(1));
+        p.assert_b_to_a([HandshakeResponse]);
+        p.drop_b_to_a();
+
+        // A vanishes, never confirms the handshake.
+        // After session timeout, B abandons handshake state.
+        p.event_b(t.at(1 + SESSION_LIFETIME.as_secs() + 1));
+
+        // B tries to send. Triggers a handshake initiation since there's no more handshake state.
+        p.send_from_b(t.at(1 + SESSION_LIFETIME.as_secs() + 2), [packet(1)]);
+        p.assert_b_to_a([HandshakeInitiation]);
+    }
+
+    #[test]
+    fn responder_timeout_with_packets() {
+        use PacketMatcher::*;
+
+        let mut p = EndpointPair::new();
+        let t = TestClock::new();
+
+        // A initiates, B responds.
+        p.send_from_a(t.at(0), [packet(1)]);
+        p.assert_a_to_b([HandshakeInitiation]);
+        p.recv_at_b(t.at(1));
+        p.assert_b_to_a([HandshakeResponse]);
+        p.drop_b_to_a();
+
+        // A vanishes, never confirms the handshake.
+
+        // B sends while handshake is still pending. Packet queued.
+        p.send_from_b(t.at(2), [packet(2)]);
+        assert_no_packets(&p.b_to_a);
+
+        // Pending handshake expires.
+        let expiry = 1 + SESSION_LIFETIME.as_secs();
+        p.event_b(t.at(expiry + 1));
+        assert_no_packets(&p.b_to_a);
+
+        // Expire A's handshake, to get everything back into an idle state.
+        p.event_a(t.at(expiry + 1));
+        assert_no_packets(&p.a_to_b);
+
+        // B sends again, becomes initiator for a full completed handshake.
+        p.send_from_b(t.at(expiry + 2), [packet(3)]);
+        p.recv_at_a(t.at(expiry + 3));
+        p.recv_at_b(t.at(expiry + 4));
+        p.assert_b_to_a([TransportData]);
+        p.recv_at_a(t.at(expiry + 5));
+
+        // Only the latest sent packet survives, packet(1) and packet(2) are lost.
+        assert_no_packets(&p.received_at_b);
+        p.assert_received_at_a([packet(3)]);
+    }
 }

ts_tunnel/src/handshake.rs

@@ -296,6 +296,8 @@ impl Handshake {
         };
 
         let session = BidiSession::new_initiator(
+            endpoint,
+            peer,
             session_keys,
             sent_handshake.responder_to_initiator_handle,
             packet.sender_id,
@@ -338,6 +340,8 @@ impl Handshake {
         self.cookie_sender.write_macs(pkt.as_mut());
 
         let session = Box::new(BidiSession::new_responder(
+            endpoint,
+            peer,
             session_keys,
             session_handle,
             handshake.responder_to_initiator_id,
@@ -398,6 +402,16 @@ impl Handshake {
         };
         self.cookie_sender.receive_cookie(packet, &handshake.mac1);
     }
+
+    /// clean up expired responded handshake state, if any.
+    pub fn cleanup_expired(&mut self, now: Instant) {
+        if let State::Responded(tentative) = &self.state
+            && tentative.expired(now)
+        {
+            tracing::trace!("pending responded handshake expired");
+            self.state = State::None;
+        }
+    }
 }
 
 #[cfg(test)]

ts_tunnel/src/session.rs

@@ -7,14 +7,14 @@ use std::{
 use aead::AeadInPlace;
 use chacha20poly1305::{ChaCha20Poly1305, KeyInit};
 use ts_packet::PacketMut;
-use ts_time::TimeRange;
+use ts_time::{Handle, TimeRange};
 use zerocopy::{
     FromBytes, Immutable, IntoBytes, KnownLayout, TryFromBytes, Unaligned,
     little_endian::{U32, U64},
 };
 
 use crate::{
-    Event, PeerId,
+    Event, PeerConfig, PeerId,
     endpoint::EndpointState,
     ids::SessionHandle,
     messages::{SessionId, TransportDataHeader},
@@ -141,6 +141,7 @@ pub struct ReceiveSession {
     session_handle: SessionHandle,
     expiry: Instant,
     window: ReplayWindow,
+    _cleanup: Handle<Event>,
 }
 
 impl Debug for ReceiveSession {
@@ -152,12 +153,24 @@ impl Debug for ReceiveSession {
 }
 
 impl ReceiveSession {
-    pub fn new(key: SessionKey, session_handle: SessionHandle, now: Instant) -> Self {
+    pub fn new(
+        endpoint: &mut EndpointState,
+        peer: &PeerConfig,
+        key: SessionKey,
+        session_handle: SessionHandle,
+        now: Instant,
+    ) -> Self {
+        let expiry = now + SESSION_LIFETIME;
+
         ReceiveSession {
             cipher: ChaCha20Poly1305::new(&key),
             session_handle,
-            expiry: now + SESSION_LIFETIME,
+            expiry,
             window: ReplayWindow::default(),
+            _cleanup: endpoint.scheduler.add(
+                TimeRange::new(expiry, expiry + SESSION_CLEANUP_GRACE),
+                Event::ExpireSession(peer.id),
+            ),
         }
     }
 
@@ -245,13 +258,17 @@ pub struct BidiSession {
 impl BidiSession {
     /// Create a new session in the initiator role.
     pub fn new_initiator(
+        endpoint: &mut EndpointState,
+        peer: &PeerConfig,
         keys: ts_noise::core::Session,
         responder_to_initiator_handle: SessionHandle,
         initiator_to_responder_id: SessionId,
         now: Instant,
     ) -> Self {
         Self {
             recv: ReceiveSession::new(
+                endpoint,
+                peer,
                 keys.responder_to_initiator,
                 responder_to_initiator_handle,
                 now,
@@ -265,13 +282,17 @@ impl BidiSession {
 
     /// Create a new session in the responder role.
     pub fn new_responder(
+        endpoint: &mut EndpointState,
+        peer: &PeerConfig,
         keys: ts_noise::core::Session,
         initiator_to_responder_handle: SessionHandle,
         responder_to_initiator_id: SessionId,
         now: Instant,
     ) -> Self {
         Self {
             recv: ReceiveSession::new(
+                endpoint,
+                peer,
                 keys.initiator_to_responder,
                 initiator_to_responder_handle,
                 now,
@@ -479,24 +500,30 @@ impl Session {
 
 #[cfg(test)]
 mod tests {
+    use ts_keys::{NodeKeyPair, NodePublicKey};
     use ts_noise::core::Role;
 
     use super::*;
-    use crate::{PeerId, ids::IdMap, messages::Message};
+    use crate::{PeerId, Psk, ids::IdMap, messages::Message};
 
     #[test]
     fn test_session_parts() {
         let k: [u8; 32] = rand::random();
         let mut ids = IdMap::default();
 
-        let initiator_session = ids.allocate_session(PeerId(1));
-        let responder_session = ids.allocate_session(PeerId(2));
+        let initiator_cfg = PeerConfig::new(PeerId(1), NodePublicKey::default(), Psk::default());
+        let initiator_session = ids.allocate_session(initiator_cfg.id);
+        let responder_cfg = PeerConfig::new(PeerId(2), NodePublicKey::default(), Psk::default());
+        let responder_session = ids.allocate_session(responder_cfg.id);
         let responder_session_id = responder_session.id();
         let now = Instant::now();
+        let mut endpoint = EndpointState::from(NodeKeyPair::new());
         // NOTE: this would be catastrophically insecure in non-test code, because it reuses the
         // same key in both directions, which leads to catastrophic nonce reuse. It's okay here
         // because (a) it's a test and (b) we only ever transmit in one direction.
         let send = BidiSession::new_initiator(
+            &mut endpoint,
+            &responder_cfg,
             ts_noise::core::Session {
                 initiator_to_responder: k.into(),
                 responder_to_initiator: k.into(),
@@ -506,7 +533,13 @@ mod tests {
             responder_session_id,
             now,
         );
-        let mut recv = ReceiveSession::new(k.into(), responder_session, now);
+        let mut recv = ReceiveSession::new(
+            &mut endpoint,
+            &initiator_cfg,
+            k.into(),
+            responder_session,
+            now,
+        );
 
         const CLEARTEXT: &[u8] = b"foobar";
         let mut pkt = [PacketMut::from(CLEARTEXT)];
@@ -538,13 +571,15 @@ mod tests {
     fn test_session_timers() {
         let k: [u8; 32] = rand::random();
         let mut ids = IdMap::default();
-        let recv_session = ids.allocate_session(PeerId(1));
+        let mut endpoint = EndpointState::from(NodeKeyPair::new());
+        let recv_cfg = PeerConfig::new(PeerId(1), NodePublicKey::default(), Psk::default());
+        let recv_session = ids.allocate_session(recv_cfg.id);
         let recv_session_id = recv_session.id();
         let bidi_session = ids.allocate_session(PeerId(2));
         let now = Instant::now();
         let epsilon = Duration::from_secs(1);
 
-        let recv = ReceiveSession::new(k.into(), recv_session, now);
+        let recv = ReceiveSession::new(&mut endpoint, &recv_cfg, k.into(), recv_session, now);
         assert!(!recv.expired(now));
         assert!(!recv.expired(now + SESSION_FRESH_LIFETIME - epsilon));
         assert!(!recv.expired(now + SESSION_FRESH_LIFETIME + epsilon));
@@ -553,6 +588,8 @@ mod tests {
         let k2: [u8; 32] = rand::random();
 
         let bidi = BidiSession::new_initiator(
+            &mut endpoint,
+            &recv_cfg,
             ts_noise::core::Session {
                 initiator_to_responder: k.into(),
                 responder_to_initiator: k2.into(),