From 824f17f37a39ea6a74f8f567ec970654ca56b2bb Mon Sep 17 00:00:00 2001 From: David Anderson Date: Tue, 18 Aug 2026 15:26:31 -0700 Subject: [PATCH] ts_tunnel: move packet queue storage back into Peer The packet queue's lifecycle is tied to both handshakes and session state, in ways that make it awkward for either to own the queue. Also delete a number of Session tests that are duplicative of Endpoint tests, and are becoming increasingly hard to make robust as the endpoint state machines become more entangled. Updates #339 Signed-off-by: David Anderson Change-Id: I2f7273fd3949c1a9f8bec3707438db356a6a6964 --- ts_tunnel/src/endpoint.rs | 89 +++++---- ts_tunnel/src/queue.rs | 14 ++ ts_tunnel/src/session.rs | 366 +++++--------------------------------- 3 files changed, 116 insertions(+), 353 deletions(-) diff --git a/ts_tunnel/src/endpoint.rs b/ts_tunnel/src/endpoint.rs index 31f29a39..8ed539ef 100644 --- a/ts_tunnel/src/endpoint.rs +++ b/ts_tunnel/src/endpoint.rs @@ -12,6 +12,7 @@ use crate::{ ids::IdMap, macs::MACReceiver, messages::{HandshakeResponse, Message, MessageMut, SessionId}, + queue::Queue, session::Session, time::TAI64NClock, }; @@ -26,8 +27,8 @@ struct Peer { config: PeerConfig, session: Session, handshake: Handshake, + queue: Queue, keepalive: Option>, - session_cleanup: Option>, send_another_keepalive: bool, } @@ -40,7 +41,7 @@ impl From for Peer { session: Default::default(), keepalive: None, - session_cleanup: None, + queue: Default::default(), send_another_keepalive: false, } } @@ -60,14 +61,20 @@ impl Peer { fn send( &mut self, endpoint: &mut EndpointState, - packets: Vec, + mut packets: Vec, now: Instant, out: &mut SendResult, ) { - if let Some(packets) = self.session.send(packets, now) { - tracing::trace!("enqueueing packets to peer"); - out.queue_to_peer(self.config.id, packets); - // Fall through to check if the session is in need of rotation. + self.check_invariants(now); + + match self.session.send(&mut packets, now) { + Ok(_) => { + tracing::trace!("enqueueing packets to peer"); + out.queue_to_peer(self.config.id, packets); + } + Err(_) => { + self.queue.append(packets); + } } if self.handshake.is_active() { @@ -92,6 +99,8 @@ impl Peer { now: Instant, out: &mut RecvResult, ) { + self.check_invariants(now); + let mut dropped = 0; packets.retain_mut(|packet| match MessageMut::try_from(packet.as_mut()) { Err(()) => { @@ -137,16 +146,17 @@ impl Peer { return; }; - let (expiry, packets) = self.session.activate(session, now, true); + self.session + .activate(endpoint, self.config.id, session, now); + let mut packets = self.queue.drain(); + if packets.is_empty() { + // Initiator must transmit one packet to confirm session. If there are none queued, + // send a keepalive. + packets.push(PacketMut::new(0)); + } + // Session was just activated, so can always send. + self.session.send(&mut packets, now).unwrap(); out.queue_to_peer(self.config.id, packets); - if let Some(handle) = self.session_cleanup.take() { - handle.cancel(); - }; - self.session_cleanup = Some( - endpoint - .scheduler - .add(expiry, Event::ExpireSession(self.config.id)), - ); } fn recv_transport_data( @@ -182,18 +192,14 @@ impl Peer { out.queue_to_local(self.config.id, packets); self.schedule_keepalive(&mut endpoint.scheduler, now); - let (expiry, packets_for_peer) = self.session.activate(session, now, false); - if !packets_for_peer.is_empty() { - out.queue_to_peer(self.config.id, packets_for_peer); - } - if let Some(handle) = self.session_cleanup.take() { - handle.cancel(); + self.session + .activate(endpoint, self.config.id, session, now); + let mut packets = self.queue.drain(); + if !packets.is_empty() { + // Session was just activated, so can always send. + self.session.send(&mut packets, now).unwrap(); + out.queue_to_peer(self.config.id, packets); } - self.session_cleanup = Some( - endpoint - .scheduler - .add(expiry, Event::ExpireSession(self.config.id)), - ); } fn respond_to_handshake( @@ -217,6 +223,8 @@ impl Peer { now: Instant, out: &mut EventResult, ) { + self.check_invariants(now); + if !self.handshake.is_active() { // Handshake completed prior to timeout firing. return; @@ -246,22 +254,18 @@ impl Peer { } fn cleanup_expired(&mut self, now: Instant) { - self.session.cleanup_expired(now) + self.check_invariants(now); + self.session.cleanup_expired(now); } fn shutdown(&mut self) { self.session.deactivate(); self.handshake.abandon(); - if let Some(handle) = self.session_cleanup.take() { - handle.cancel(); - } if let Some(handle) = self.keepalive.take() { handle.cancel(); } } - /// (Soft) precondition: `self.handshake == HandshakeState::None` (previous handshake is lost, but - /// that shouldn't cause anything terrible to happen). fn start_handshake( &mut self, endpoint: &mut EndpointState, @@ -271,6 +275,25 @@ impl Peer { let packet = self.handshake.initiate(endpoint, &self.config, now); out.queue_to_peer(self.config.id, [packet]); } + + fn check_invariants(&self, now: Instant) { + if !cfg!(debug_assertions) { + return; + } + if !self.queue.is_empty() { + assert!( + !self.session.is_active(now), + "peer {:?}: packets in queue with active session", + self.config.id + ); + + assert!( + self.handshake.is_active(), + "peer {:?}: packets in queue with no handshake in flight", + self.config.id + ); + } + } } /// A WireGuard endpoint capable of communicating with multiple remote peers. diff --git a/ts_tunnel/src/queue.rs b/ts_tunnel/src/queue.rs index a0e277e5..8e589b74 100644 --- a/ts_tunnel/src/queue.rs +++ b/ts_tunnel/src/queue.rs @@ -30,6 +30,20 @@ impl Queue { self.0.clear(); self.0.shrink_to_fit(); } + + /// Drain all packets from the queue into a `Vec`. + /// + /// The queue's memory footprint is shrunk to its minimum, on the assumption that + /// it is unlikely to be used again soon. + pub fn drain(&mut self) -> Vec { + let ret = self.0.drain(..).collect(); + self.clear(); + ret + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } } impl IntoIterator for Queue { diff --git a/ts_tunnel/src/session.rs b/ts_tunnel/src/session.rs index 7225a005..395c65b6 100644 --- a/ts_tunnel/src/session.rs +++ b/ts_tunnel/src/session.rs @@ -14,9 +14,10 @@ use zerocopy::{ }; use crate::{ + Event, PeerId, + endpoint::EndpointState, ids::SessionHandle, messages::{SessionId, TransportDataHeader}, - queue::Queue, replay::ReplayWindow, }; @@ -372,72 +373,41 @@ impl ActiveSession { } /// A communication session to a peer. -pub enum Session { - /// No session established yet. - /// - /// This session cannot receive packets. Sent packets are queued if possible, and will be - /// transmitted if the session becomes active in the future. - None(Queue), - /// Active session capable of bidirectional communication. - Active(ActiveSession), -} - -impl Default for Session { - fn default() -> Self { - Self::None(Queue::default()) - } -} +#[derive(Default)] +pub struct Session(Option); impl Session { - fn take(&mut self) -> Session { - std::mem::take(self) - } - - /// Return a reference to the active session, if any. - /// - /// Calls [`Session::maybe_expire`], so callers can assume that the returned session - /// consists only of unexpired state. - fn as_active(&mut self, now: Instant) -> Option<&mut ActiveSession> { - self.cleanup_expired(now); - if let Self::Active(session) = self { - Some(session) - } else { - None + pub fn is_active(&self, now: Instant) -> bool { + match &self.0 { + Some(session) => !session.expired(now), + None => false, } } /// Activate the session with the given keys. - /// - /// Returns the expiry time for the session. If any packets were queued waiting - /// for an active session, they are encrypted and returned. pub fn activate( &mut self, + endpoint: &mut EndpointState, + peer_id: PeerId, next: BidiSession, now: Instant, - need_keepalive: bool, - ) -> (TimeRange, Vec) { + ) { tracing::trace!(recv_id = ?next.recv.id(), "activating new session"); - let (active, mut packets) = match self.take() { - Self::None(queue) => (next.into(), queue.into()), - Self::Active(mut session) => { + match &mut self.0 { + Some(session) => { session.rotate(next, now); - (session, vec![]) } - }; - - if need_keepalive && packets.is_empty() { - packets.push(PacketMut::new(0)); + None => self.0 = Some(next.into()), } - active.cur.encrypt(&mut packets); - *self = Self::Active(active); - ( - TimeRange::new( - now + SESSION_LIFETIME, - now + SESSION_LIFETIME + SESSION_CLEANUP_GRACE, - ), - packets, - ) + + let cleanup = TimeRange::new( + now + SESSION_LIFETIME, + now + SESSION_LIFETIME + SESSION_CLEANUP_GRACE, + ); + endpoint + .scheduler + .add(cleanup, Event::ExpireSession(peer_id)); } /// Discard all state for this session. @@ -449,7 +419,8 @@ impl Session { /// /// Returns None if the session is inactive (and thus no keepalive is necessary). pub fn send_keepalive(&mut self, now: Instant) -> Option { - let session = self.as_active(now)?; + self.cleanup_expired(now); + let session = self.0.as_mut()?; let mut packet = vec![PacketMut::new(0)]; session.cur.encrypt(&mut packet); packet.pop() @@ -457,27 +428,18 @@ impl Session { /// Send packets to the peer. /// - /// If the session is inactive, packets are queued for future transmission. - /// - /// Returns None to indicate that packets were queued, indicating the caller may need to - /// initiate a handshake. - pub fn send(&mut self, mut packets: Vec, now: Instant) -> Option> { + /// Returns Err to indicate that no session is available for transmission. + pub fn send(&mut self, packets: &mut Vec, now: Instant) -> Result<(), ()> { self.cleanup_expired(now); - match self { - Self::None(queue) => { - queue.append(packets); - None - } - Self::Active(session) => { - session.cur.encrypt(&mut packets); - Some(packets) - } - } + let session = self.0.as_mut().ok_or(())?; + session.cur.encrypt(packets); + Ok(()) } /// Get the ReceiveSession for the given receiving ID, if any. pub fn get_recv(&mut self, id: SessionId, now: Instant) -> Option<&mut ReceiveSession> { - let session = self.as_active(now)?; + self.cleanup_expired(now); + let session = self.0.as_mut()?; if session.cur.recv_id() == id { Some(&mut session.cur.recv) } else if let Some(prev) = session.prev.as_mut() @@ -491,24 +453,26 @@ impl Session { /// Reports whether the session is old enough to require rotation. pub fn needs_rotation(&self, now: Instant) -> bool { - match self { - Self::None(_) => true, - Self::Active(session) => session.cur.needs_rotation(now), + match &self.0 { + None => true, + Some(session) => session.cur.needs_rotation(now), } } /// Clean up expired session state, if any. pub fn cleanup_expired(&mut self, now: Instant) { - if let Self::Active(session) = self { - if session.expired(now) { - *self = Self::default(); - return; - } - if let Some(prev) = session.prev.as_ref() - && prev.expired(now) - { - session.prev = None; - } + let Some(session) = self.0.as_mut() else { + return; + }; + + if session.expired(now) { + *self = Self::default(); + return; + } + if let Some(prev) = session.prev.as_ref() + && prev.expired(now) + { + session.prev = None; } } } @@ -570,244 +534,6 @@ mod tests { assert_eq!(pkt[0].as_ref(), CLEARTEXT); } - fn packet(payload: &str) -> Vec { - vec![PacketMut::from(payload.as_bytes())] - } - - #[derive(Default)] - struct PeerSession { - ids: IdMap, - recv_id: Option, - recv_id_prev: Option, - session: Session, - } - - impl PeerSession { - fn allocate_handle(&mut self) -> SessionHandle { - self.recv_id_prev = self.recv_id.take(); - let ret = self.ids.allocate_session(PeerId(1)); - self.recv_id = Some(ret.id()); - ret - } - - fn handshake_with( - &mut self, - other: &mut Self, - now: Instant, - ) -> (Vec, Vec) { - let (k1, k2): ([u8; 32], [u8; 32]) = rand::random(); - let sid1 = self.allocate_handle(); - let sid1_id = sid1.id(); - let sid2 = other.allocate_handle(); - let s1 = BidiSession::new_initiator( - ts_noise::core::Session { - initiator_to_responder: k1.into(), - responder_to_initiator: k2.into(), - role: Role::Initiator, - }, - sid1, - sid2.id(), - now, - ); - let s2 = BidiSession::new_responder( - ts_noise::core::Session { - initiator_to_responder: k1.into(), - responder_to_initiator: k2.into(), - role: Role::Responder, - }, - sid2, - sid1_id, - now, - ); - let (_, p1) = self.session.activate(s1, now, false); - let (_, p2) = other.session.activate(s2, now, false); - (p1, p2) - } - - fn send(&mut self, now: Instant, packets: Vec) -> Option> { - self.session.send(packets, now) - } - - fn get_recv(&mut self, now: Instant, packets: &[PacketMut]) -> Option<&mut ReceiveSession> { - let (hdr, _) = - TransportDataHeader::try_ref_from_prefix(packets.first()?.as_ref()).unwrap(); - self.session.get_recv(hdr.receiver_id, now) - } - - fn recv(&mut self, now: Instant, packets: Vec) -> Vec { - self.get_recv(now, &packets).unwrap().decrypt(packets) - } - - fn unknown_recv(&mut self, now: Instant, packets: Vec) -> bool { - self.get_recv(now, &packets).is_none() - } - - fn needs_handshake(&self, now: Instant) -> bool { - self.session.needs_rotation(now) - } - } - - #[test] - fn test_session() { - let mut a = PeerSession::default(); - let mut b = PeerSession::default(); - - let now = Instant::now(); - - assert_eq!(a.send(now, packet("foobar")), None); - assert_eq!(b.send(now, packet("qux")), None); - - assert!(a.needs_handshake(now)); - assert!(b.needs_handshake(now)); - - // Establish a session between the peers. We're cheating and not doing any of the - // handshake lifecycle. - let (a_to_b, b_to_a) = a.handshake_with(&mut b, now); - assert_eq!(a_to_b.len(), 1); - assert_eq!(b_to_a.len(), 1); - - // Verify that the packets queued prior to session activation transmit correctly. - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("foobar")); - - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("qux")); - - assert!(!a.needs_handshake(now)); - assert!(!b.needs_handshake(now)); - - // Transmit with established session. - let now = now + Duration::from_secs(60); - - let a_to_b = a.send(now, packet("frobozz")).unwrap(); - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("frobozz")); - - let b_to_a = b.send(now, packet("xyzzy")).unwrap(); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("xyzzy")); - - assert!(!a.needs_handshake(now)); - assert!(!b.needs_handshake(now)); - - // Transmit with stale session. - let now = now + Duration::from_secs(70); - - let a_to_b = a.send(now, packet("foo")).unwrap(); - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("foo")); - - let b_to_a = b.send(now, packet("bar")).unwrap(); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("bar")); - - assert!(a.needs_handshake(now)); - assert!(!b.needs_handshake(now)); - - // Transmit with expired session. - let now = now + Duration::from_secs(120); - - let a_to_b = a.send(now, packet("no")); - assert_eq!(a_to_b, None); - - let b_to_a = b.send(now, packet("nope")); - assert_eq!(b_to_a, None); - - assert!(a.needs_handshake(now)); - assert!(b.needs_handshake(now)); - } - - #[test] - fn test_rotation() { - let mut a = PeerSession::default(); - let mut b = PeerSession::default(); - - let start = Instant::now(); - let epsilon = Duration::from_secs(1); - - let now = start; - let (a_to_b, b_to_a) = a.handshake_with(&mut b, now); - assert!(a_to_b.is_empty()); - assert!(b_to_a.is_empty()); - - let now = start + SESSION_FRESH_LIFETIME - epsilon; - let a_to_b = a.send(now, packet("foo")).unwrap(); - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("foo")); - - let b_to_a = b.send(now, packet("bar")).unwrap(); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("bar")); - - // Rotate session, with packets delivered across the rotation. - let a_to_b = a.send(now, packet("before rotate A->B")).unwrap(); - let b_to_a = b.send(now, packet("before rotate B->A")).unwrap(); - - let very_delayed_a_to_b = a.send(now, packet("delayed A->B")).unwrap(); - let very_delayed_b_to_a = b.send(now, packet("delayed B->A")).unwrap(); - - let to_send = a.handshake_with(&mut b, now); - assert_eq!(to_send, (vec![], vec![])); - - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("before rotate A->B")); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("before rotate B->A")); - - let a_to_b = a.send(now, packet("after rotate A->B")).unwrap(); - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("after rotate A->B")); - - let b_to_a = b.send(now, packet("after rotate B->A")).unwrap(); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("after rotate B->A")); - - // Rotate again, delayed packets should not decrypt anymore. - let now = start + SESSION_FRESH_LIFETIME + epsilon; - let a_to_b = a.send(now, packet("before rotate2 A->B")).unwrap(); - let b_to_a = b.send(now, packet("before rotate2 B->A")).unwrap(); - - let to_send = a.handshake_with(&mut b, now); - assert_eq!(to_send, (vec![], vec![])); - - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("before rotate2 A->B")); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("before rotate2 B->A")); - - let a_to_b = a.send(now, packet("after rotate2 A->B")).unwrap(); - let a_to_b = b.recv(now, a_to_b); - assert_eq!(a_to_b, packet("after rotate2 A->B")); - - let b_to_a = b.send(now, packet("after rotate2 B->A")).unwrap(); - let b_to_a = a.recv(now, b_to_a); - assert_eq!(b_to_a, packet("after rotate2 B->A")); - - assert!(b.unknown_recv(now, very_delayed_a_to_b)); - assert!(a.unknown_recv(now, very_delayed_b_to_a)); - } - - #[test] - fn test_expiration() { - let mut a = PeerSession::default(); - let mut b = PeerSession::default(); - - let now = Instant::now(); - let to_send = a.handshake_with(&mut b, now); - assert_eq!(to_send, (vec![], vec![])); - - let a_to_b = a.send(now, packet("A->B")).unwrap(); - let b_to_a = b.send(now, packet("B->A")).unwrap(); - - let now = now + SESSION_LIFETIME + SESSION_LIFETIME; - - assert!(b.unknown_recv(now, a_to_b)); - assert!(a.unknown_recv(now, b_to_a)); - - assert_eq!(a.send(now, packet("expired A->B")), None); - assert_eq!(b.send(now, packet("expired B->A")), None); - } - #[test] fn test_session_timers() { let k: [u8; 32] = rand::random();