diff --git a/ts_tunnel/src/endpoint.rs b/ts_tunnel/src/endpoint.rs index cdc3606a..0a8ffa3f 100644 --- a/ts_tunnel/src/endpoint.rs +++ b/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)]); + } } diff --git a/ts_tunnel/src/handshake.rs b/ts_tunnel/src/handshake.rs index 3ffbf1d0..0dfc4870 100644 --- a/ts_tunnel/src/handshake.rs +++ b/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)] diff --git a/ts_tunnel/src/session.rs b/ts_tunnel/src/session.rs index cea6f97f..e7aeeb81 100644 --- a/ts_tunnel/src/session.rs +++ b/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, } 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,6 +258,8 @@ 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, @@ -252,6 +267,8 @@ impl BidiSession { ) -> Self { Self { recv: ReceiveSession::new( + endpoint, + peer, keys.responder_to_initiator, responder_to_initiator_handle, now, @@ -265,6 +282,8 @@ 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, @@ -272,6 +291,8 @@ impl BidiSession { ) -> 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(),