From 91d6e59adec7e6facb11273c07e380fb113708a5 Mon Sep 17 00:00:00 2001 From: David Anderson Date: Thu, 6 Aug 2026 12:38:02 -0700 Subject: [PATCH] ts_tunnel: make IdMap hand out handles for newly allocated sessions Every allocated session needs to be owned either by an in-flight handshake or an established session anyway, so making it RAII makes the cleanup of expired or abandoned sessions much more robust. Updates #339 Signed-off-by: David Anderson Change-Id: I3e03401ecc0d7882f5a794583b6b36d66a6a6964 --- ts_tunnel/src/endpoint.rs | 32 +++---- ts_tunnel/src/handshake.rs | 39 ++++---- ts_tunnel/src/ids.rs | 79 ++++++++++------ ts_tunnel/src/session.rs | 178 ++++++++++++++++--------------------- 4 files changed, 162 insertions(+), 166 deletions(-) diff --git a/ts_tunnel/src/endpoint.rs b/ts_tunnel/src/endpoint.rs index 188a69a4..0e7c68c0 100644 --- a/ts_tunnel/src/endpoint.rs +++ b/ts_tunnel/src/endpoint.rs @@ -68,7 +68,7 @@ impl Peer { now: Instant, out: &mut SendResult, ) { - if let Some(packets) = self.session.send(packets, &mut endpoint.ids, now) { + if let Some(packets) = self.session.send(packets, now) { tracing::trace!("enqueueing packets to peer"); out.queue_to_peer(self.id, packets); // Fall through to check if the session is in need of rotation. @@ -155,7 +155,7 @@ impl Peer { return; }; - let (expiry, packets) = self.session.activate(session, &mut endpoint.ids, now, true); + let (expiry, packets) = self.session.activate(session, now, true); out.queue_to_peer(self.id, packets); if let Some(handle) = self.session_cleanup.take() { handle.cancel(); @@ -175,7 +175,7 @@ impl Peer { now: Instant, out: &mut RecvResult, ) { - if let Some(recv) = self.session.get_recv(session_id, &mut endpoint.ids, now) { + if let Some(recv) = self.session.get_recv(session_id, now) { let packets = recv.decrypt(packets); if !packets.is_empty() { out.queue_to_local(self.id, packets); @@ -195,9 +195,7 @@ impl Peer { out.queue_to_local(self.id, packets); self.schedule_keepalive(&mut endpoint.scheduler, now); - let (expiry, packets_for_peer) = - self.session - .activate(session, &mut endpoint.ids, now, false); + let (expiry, packets_for_peer) = self.session.activate(session, now, false); if !packets_for_peer.is_empty() { out.queue_to_peer(self.id, packets_for_peer); } @@ -231,7 +229,6 @@ impl Peer { } self.last_seen_timestamp = Some(handshake.timestamp); - endpoint.ids.remove_handshake_session(&self.handshake); let session_id = endpoint.ids.allocate_session(self.id); let packet = self.handshake.respond( @@ -255,7 +252,6 @@ impl Peer { return; } - endpoint.ids.remove_handshake_session(&self.handshake); self.handshake = Handshake::None; self.start_handshake(endpoint, now, out); @@ -267,7 +263,7 @@ impl Peer { now: Instant, out: &mut EventResult, ) { - let Some(packet) = self.session.send_keepalive(&mut endpoint.ids, now) else { + let Some(packet) = self.session.send_keepalive(now) else { tracing::trace!("send keepalive: session expired, skipping"); return; }; @@ -281,14 +277,13 @@ impl Peer { } } - fn cleanup_expired(&mut self, endpoint: &mut EndpointState, now: Instant) { - self.session.cleanup_expired(&mut endpoint.ids, now) + fn cleanup_expired(&mut self, now: Instant) { + self.session.cleanup_expired(now) } - fn shutdown(&mut self, endpoint: &mut EndpointState) { - self.session.deactivate(&mut endpoint.ids); + fn shutdown(&mut self) { + self.session.deactivate(); - endpoint.ids.remove_handshake_session(&self.handshake); self.handshake = Handshake::None; if let Some(handle) = self.session_cleanup.take() { handle.cancel(); @@ -308,6 +303,7 @@ impl Peer { ) { // TODO most of this logic might be better in the `handshake` module. let session_id = endpoint.ids.allocate_session(self.id); + tracing::debug!(peer_id = ?self.id, ?session_id, "enqueue handshake start"); let (handshake, packet) = initiate_handshake( &endpoint.my_key, &self.config.key, @@ -318,8 +314,6 @@ impl Peer { let mut packet = PacketMut::from(packet.as_bytes()); let mac = self.cookie_sender.write_macs(packet.as_mut()); - tracing::debug!(peer_id = ?self.id, ?session_id, "enqueue handshake start"); - out.queue_to_peer(self.id, [packet]); let tr = TimeRange::new_around(now + HANDSHAKE_TIMEOUT, Duration::from_millis(500)); @@ -397,7 +391,7 @@ impl Endpoint { match self.peers.remove(&peer) { None => false, Some(mut peer) => { - peer.shutdown(&mut self.state); + peer.shutdown(); self.state.ids.remove_peer(&peer.config.key); true } @@ -463,7 +457,7 @@ impl Endpoint { tracing::warn!(?session_id, "session not found"); continue; }; - let Some(peer) = self.peers.get_mut(peer_id) else { + let Some(peer) = self.peers.get_mut(&peer_id) else { tracing::warn!(?peer_id, "no peer found"); continue; }; @@ -523,7 +517,7 @@ impl Endpoint { let Some(peer) = self.peers.get_mut(&peer_id) else { continue; }; - peer.cleanup_expired(&mut self.state, now); + peer.cleanup_expired(now); } } } diff --git a/ts_tunnel/src/handshake.rs b/ts_tunnel/src/handshake.rs index 23e4699a..350e12d0 100644 --- a/ts_tunnel/src/handshake.rs +++ b/ts_tunnel/src/handshake.rs @@ -9,6 +9,7 @@ use zerocopy::IntoBytes; use crate::{ config::Psk, endpoint::Event, + ids::SessionHandle, macs::{MACReceiver, MACSender, Mac}, messages::*, session::BidiSession, @@ -49,20 +50,20 @@ impl ReceivedHandshake { /// Finalize the handshake, producing a HandshakeResponse. pub fn respond( self, - initiator_to_responder_id: SessionId, + initiator_to_responder_id: SessionHandle, psk: &Psk, macs: &MACSender, now: Instant, ) -> (BidiSession, PacketMut) { let mut response = HandshakeResponse { - sender_id: initiator_to_responder_id, + sender_id: initiator_to_responder_id.id(), receiver_id: self.responder_to_initiator_id, ..Default::default() }; let session_keys = self.noise.finish(psk, response.noise.as_mut_bytes()); - let session = BidiSession::new( + let session = BidiSession::new_responder( session_keys, initiator_to_responder_id, self.responder_to_initiator_id, @@ -84,11 +85,11 @@ impl ReceivedHandshake { pub fn initiate_handshake( endpoint_static: &NodeKeyPair, peer_static: &NodePublicKey, - session_id: SessionId, + session_id: SessionHandle, timestamp: TAI64N, ) -> (SentHandshake, HandshakeInitiation) { let mut pkt = HandshakeInitiation { - sender_id: session_id, + sender_id: session_id.id(), ..Default::default() }; @@ -110,7 +111,7 @@ pub fn initiate_handshake( /// A partially completed sent handshake. pub struct SentHandshake { - pub responder_to_initiator_id: SessionId, + pub responder_to_initiator_id: SessionHandle, noise: ikpsk2::SentHandshake, } @@ -132,15 +133,6 @@ impl Handshake { !matches!(self, Handshake::None) } - /// Return the session id of the handshake, if any. - pub(crate) fn session_id(&self) -> Option { - match self { - Handshake::Initiated(handshake, ..) => Some(handshake.responder_to_initiator_id), - Handshake::Responded(tentative) => Some(tentative.recv_id()), - Handshake::None => None, - } - } - pub(crate) fn take_initiated(&mut self) -> Option<(SentHandshake, Handle, Mac)> { match std::mem::replace(self, Handshake::None) { Handshake::Initiated(sent, timeout, mac) => Some((sent, timeout, mac)), @@ -157,7 +149,7 @@ impl Handshake { /// Responding replaces any other handshake state unconditionally. pub(crate) fn respond( &mut self, - session_id: SessionId, + session_id: SessionHandle, handshake: ReceivedHandshake, psk: &Psk, cookie_sender: &MACSender, @@ -214,10 +206,10 @@ impl Handshake { } }; - let session = BidiSession::new( + let session = BidiSession::new_initiator( session_keys, - packet.sender_id, sent_handshake.responder_to_initiator_id, + packet.sender_id, now, ); @@ -268,6 +260,7 @@ mod tests { use zerocopy::TryFromBytes; use super::*; + use crate::{PeerId, ids::IdMap}; #[test] fn test_handshake() { @@ -277,7 +270,8 @@ mod tests { // Peer A sends a handshake initiation... let a_mac_send = MACSender::new(&b_static.public); let a_mac_recv = MACReceiver::new(&a_static.public); - let a_session = SessionId::random(); // A wants to receive at this ID + let mut ids = IdMap::default(); + let a_session = ids.allocate_session(PeerId(1)); // A wants to receive at this ID let a_init_time = TAI64N::now(); let (a_handshake, init_pkt) = initiate_handshake(&a_static, &b_static.public, a_session, a_init_time); @@ -301,7 +295,7 @@ mod tests { .expect("peer B should successfully process A's handshake initiation"); assert_eq!(b_handshake.peer_static(), a_static.public); assert_eq!(b_handshake.timestamp, a_init_time); - let b_session = SessionId::random(); // B wants to receive at this ID + let b_session = ids.allocate_session(PeerId(2)); // B wants to receive at this ID let (mut b_session, mut response_pkt) = b_handshake.respond(b_session, &psk, &b_mac_send, Instant::now()); @@ -333,11 +327,12 @@ mod tests { fn test_invalid_response_ignored() { let (a_static, b_static) = (NodeKeyPair::new(), NodeKeyPair::new()); let psk = rand::random(); + let mut ids = IdMap::default(); // A sends a handshake let a_mac_send = MACSender::new(&b_static.public); let a_mac_recv = MACReceiver::new(&a_static.public); - let a_session = SessionId::random(); // A wants to receive at this ID + let a_session = ids.allocate_session(PeerId(1)); // A wants to receive at this ID let a_init_time = TAI64N::now(); let (a_handshake, init_pkt) = initiate_handshake(&a_static, &b_static.public, a_session, a_init_time); @@ -358,7 +353,7 @@ mod tests { let b_mac_recv = MACReceiver::new(&b_static.public); let b_handshake = ReceivedHandshake::new(init_pkt, &b_static, &b_mac_recv) .expect("peer B should successfully process A's handshake initiation"); - let b_session = SessionId::random(); // B wants to receive at this ID + let b_session = ids.allocate_session(PeerId(2)); // B wants to receive at this ID let (_b_session, mut response_pkt) = b_handshake.respond(b_session, &psk, &b_mac_send, Instant::now()); diff --git a/ts_tunnel/src/ids.rs b/ts_tunnel/src/ids.rs index 2804d789..c72bf472 100644 --- a/ts_tunnel/src/ids.rs +++ b/ts_tunnel/src/ids.rs @@ -1,13 +1,17 @@ -use std::collections::HashMap; +use core::fmt::{Debug, Formatter}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex, Weak}, +}; use ts_keys::NodePublicKey; -use crate::{PeerId, handshake::Handshake, messages::SessionId}; +use crate::{PeerId, messages::SessionId}; /// Tracks and allocates session IDs for peer sessions. #[derive(Default)] pub struct IdMap { - sessions: HashMap, + sessions: Arc>>, // TODO: track recently abandoned session IDs, avoid reusing them for // one or two session lifetimes to avoid confusion with reordered packets. node_keys: HashMap, @@ -20,8 +24,8 @@ impl IdMap { } /// Return the peer handle for a session, if any. - pub fn get_by_session_id(&self, key: &SessionId) -> Option<&PeerId> { - self.sessions.get(key) + pub fn get_by_session_id(&self, key: &SessionId) -> Option { + self.sessions.lock().unwrap().get(key).copied() } /// Add a peer handle for communicating with the given peer pubkey. @@ -36,40 +40,63 @@ impl IdMap { true } + /// Delete the peer handle for the given key. + /// + /// Panics if there is no peer currently using that key. + pub fn remove_peer(&mut self, key: &NodePublicKey) { + self.node_keys.remove(key).unwrap(); + } + /// Allocate a new session ID for communication with the given peer. /// /// Note that due to key rotation, a peer can have multiple session IDs in use at once. - pub fn allocate_session(&mut self, peer: PeerId) -> SessionId { + pub fn allocate_session(&mut self, peer: PeerId) -> SessionHandle { + let mut sessions = self.sessions.lock().unwrap(); loop { let ret = SessionId::random(); - if self.sessions.contains_key(&ret) { + if sessions.contains_key(&ret) { continue; } - self.sessions.insert(ret, peer); - return ret; + sessions.insert(ret, peer); + return SessionHandle { + sessions: Arc::downgrade(&self.sessions), + id: ret, + }; } } +} - /// Abandon the given session ID. - /// - /// Panics if the session ID isn't currently in use. - pub fn remove_session(&mut self, id: SessionId) { - self.sessions.remove(&id).unwrap(); +/// A handle for a receiving session. +pub struct SessionHandle { + sessions: Weak>>, + id: SessionId, +} + +impl SessionHandle { + /// Return the wire ID for this session. + pub fn id(&self) -> SessionId { + self.id } +} - /// Abandon the session ID associated with a handshake. - /// - /// Panics if the handshake's ID wasn't allocated in this IdMap. - pub fn remove_handshake_session(&mut self, handshake: &Handshake) { - if let Some(id) = handshake.session_id() { - self.remove_session(id); - } +impl Debug for SessionHandle { + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + core::fmt::Debug::fmt(&self.id, f) } +} - /// Delete the peer handle for the given key. - /// - /// Panics if there is no peer currently using that key. - pub fn remove_peer(&mut self, key: &NodePublicKey) { - self.node_keys.remove(key).unwrap(); +impl AsRef for SessionHandle { + fn as_ref(&self) -> &SessionId { + &self.id + } +} + +impl Drop for SessionHandle { + fn drop(&mut self) { + let Some(sessions) = self.sessions.upgrade() else { + return; + }; + let mut sessions = sessions.lock().unwrap(); + sessions.remove(&self.id).unwrap(); } } diff --git a/ts_tunnel/src/session.rs b/ts_tunnel/src/session.rs index 7fd39b3c..9c4f98f5 100644 --- a/ts_tunnel/src/session.rs +++ b/ts_tunnel/src/session.rs @@ -8,7 +8,6 @@ use std::{ use aead::AeadInPlace; use chacha20poly1305::{ChaCha20Poly1305, KeyInit}; -use ts_noise::core::Role; use ts_packet::PacketMut; use ts_time::TimeRange; use zerocopy::{ @@ -17,7 +16,7 @@ use zerocopy::{ }; use crate::{ - ids::IdMap, + ids::SessionHandle, messages::{SessionId, TransportDataHeader}, replay::ReplayWindow, }; @@ -139,7 +138,7 @@ pub const SESSION_CLEANUP_GRACE: Duration = Duration::from_secs(5); /// Established session that can only receive. pub struct ReceiveSession { cipher: ChaCha20Poly1305, - id: SessionId, + id: SessionHandle, expiry: Instant, window: ReplayWindow, } @@ -147,13 +146,13 @@ pub struct ReceiveSession { impl Debug for ReceiveSession { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("ReceiveSession") - .field("id", &self.id) + .field("id", self.id.as_ref()) .finish_non_exhaustive() } } impl ReceiveSession { - pub fn new(key: SessionKey, id: SessionId, now: Instant) -> Self { + pub fn new(key: SessionKey, id: SessionHandle, now: Instant) -> Self { ReceiveSession { cipher: ChaCha20Poly1305::new(&key), id, @@ -181,7 +180,7 @@ impl ReceiveSession { let _guard = tracing::trace_span!("header_parsed", ?header).entered(); - if header.receiver_id != self.id { + if header.receiver_id != self.id() { // Technically an unnecessary check, because a bespoke session is created for each // session ID, with different AEAD keys. So, if the caller mistakenly hands the wrong // packet to a session, it'll always fail to decrypt below. But, comparing one u32 @@ -223,7 +222,7 @@ impl ReceiveSession { /// Return the session ID that will appear on received packets meant for this session. pub fn id(&self) -> SessionId { - self.id + self.id.id() } /// Report whether the session is expired. @@ -244,35 +243,35 @@ pub struct BidiSession { } impl BidiSession { - pub fn new( + /// Create a new session in the initiator role. + pub fn new_initiator( keys: ts_noise::core::Session, + responder_to_initiator_id: SessionHandle, initiator_to_responder_id: SessionId, + now: Instant, + ) -> Self { + Self { + recv: ReceiveSession::new(keys.responder_to_initiator, responder_to_initiator_id, now), + send_id: initiator_to_responder_id, + send_cipher: ChaCha20Poly1305::new(&keys.initiator_to_responder), + send_nonce: Default::default(), + is_initiator: true, + } + } + + /// Create a new session in the responder role. + pub fn new_responder( + keys: ts_noise::core::Session, + initiator_to_responder_id: SessionHandle, responder_to_initiator_id: SessionId, now: Instant, ) -> Self { - match keys.role { - Role::Initiator => Self { - recv: ReceiveSession::new( - keys.responder_to_initiator, - responder_to_initiator_id, - now, - ), - send_id: initiator_to_responder_id, - send_cipher: ChaCha20Poly1305::new(&keys.initiator_to_responder), - send_nonce: Default::default(), - is_initiator: true, - }, - Role::Responder => Self { - recv: ReceiveSession::new( - keys.initiator_to_responder, - initiator_to_responder_id, - now, - ), - send_id: responder_to_initiator_id, - send_cipher: ChaCha20Poly1305::new(&keys.responder_to_initiator), - send_nonce: Default::default(), - is_initiator: false, - }, + Self { + recv: ReceiveSession::new(keys.initiator_to_responder, initiator_to_responder_id, now), + send_id: responder_to_initiator_id, + send_cipher: ChaCha20Poly1305::new(&keys.responder_to_initiator), + send_nonce: Default::default(), + is_initiator: false, } } @@ -310,7 +309,7 @@ impl BidiSession { /// Return the session ID that will appear on received packets meant for this session. pub fn recv_id(&self) -> SessionId { - self.recv.id + self.recv.id.id() } pub fn rotation_time(&self) -> Instant { @@ -385,25 +384,13 @@ impl ActiveSession { /// /// The prior receive session is rotated into the previous slot, and will continue to accept /// packets until the next rotation (or the hard session expiry deadline). - fn rotate(&mut self, next: BidiSession, ids: &mut IdMap, now: Instant) { - if let Some(prev) = self.prev.as_ref() { - ids.remove_session(prev.id()); - } + fn rotate(&mut self, next: BidiSession, now: Instant) { let prev = std::mem::replace(self.cur.as_mut(), next); - if prev.expired(now) { - ids.remove_session(prev.recv_id()); - } else { + if !prev.expired(now) { self.prev = Some(Box::new(prev.into())); } } - fn cleanup_ids(&mut self, ids: &mut IdMap) { - ids.remove_session(self.cur.recv_id()); - if let Some(prev) = self.prev.as_ref() { - ids.remove_session(prev.id()); - } - } - fn expired(&self, now: Instant) -> bool { self.cur.expired(now) } @@ -435,8 +422,8 @@ impl Session { /// /// Calls [`Session::maybe_expire`], so callers can assume that the returned session /// consists only of unexpired state. - fn as_active(&mut self, ids: &mut IdMap, now: Instant) -> Option<&mut ActiveSession> { - self.cleanup_expired(ids, now); + fn as_active(&mut self, now: Instant) -> Option<&mut ActiveSession> { + self.cleanup_expired(now); if let Self::Active(session) = self { Some(session) } else { @@ -451,7 +438,6 @@ impl Session { pub fn activate( &mut self, next: BidiSession, - ids: &mut IdMap, now: Instant, need_keepalive: bool, ) -> (TimeRange, Vec) { @@ -460,7 +446,7 @@ impl Session { let (active, mut packets) = match self.take() { Self::None(queue) => (next.into(), queue.into()), Self::Active(mut session) => { - session.rotate(next, ids, now); + session.rotate(next, now); (session, vec![]) } }; @@ -480,18 +466,15 @@ impl Session { } /// Discard all state for this session. - pub fn deactivate(&mut self, ids: &mut IdMap) { - if let Self::Active(mut session) = self.take() { - session.cleanup_ids(ids); - } + pub fn deactivate(&mut self) { *self = Self::default(); } /// Encrypt a keepalive packet for the peer. /// /// Returns None if the session is inactive (and thus no keepalive is necessary). - pub fn send_keepalive(&mut self, ids: &mut IdMap, now: Instant) -> Option { - let session = self.as_active(ids, now)?; + pub fn send_keepalive(&mut self, now: Instant) -> Option { + let session = self.as_active(now)?; let mut packet = vec![PacketMut::new(0)]; session.cur.encrypt(&mut packet); packet.pop() @@ -503,13 +486,8 @@ impl Session { /// /// 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, - ids: &mut IdMap, - now: Instant, - ) -> Option> { - self.cleanup_expired(ids, now); + pub fn send(&mut self, mut packets: Vec, now: Instant) -> Option> { + self.cleanup_expired(now); match self { Self::None(queue) => { queue.append(packets); @@ -523,13 +501,8 @@ impl Session { } /// Get the ReceiveSession for the given receiving ID, if any. - pub fn get_recv( - &mut self, - id: SessionId, - ids: &mut IdMap, - now: Instant, - ) -> Option<&mut ReceiveSession> { - let session = self.as_active(ids, now)?; + pub fn get_recv(&mut self, id: SessionId, now: Instant) -> Option<&mut ReceiveSession> { + let session = self.as_active(now)?; if session.cur.recv_id() == id { Some(&mut session.cur.recv) } else if let Some(prev) = session.prev.as_mut() @@ -550,17 +523,15 @@ impl Session { } /// Clean up expired session state, if any. - pub fn cleanup_expired(&mut self, ids: &mut IdMap, now: Instant) { + pub fn cleanup_expired(&mut self, now: Instant) { if let Self::Active(session) = self { if session.expired(now) { - session.cleanup_ids(ids); *self = Self::default(); return; } if let Some(prev) = session.prev.as_ref() && prev.expired(now) { - ids.remove_session(prev.id()); session.prev = None; } } @@ -569,28 +540,34 @@ impl Session { #[cfg(test)] mod tests { + use ts_noise::core::Role; + use super::*; - use crate::{PeerId, messages::Message}; + use crate::{PeerId, ids::IdMap, messages::Message}; #[test] fn test_session_parts() { let k: [u8; 32] = rand::random(); - let session = SessionId::random(); + let mut ids = IdMap::default(); + + let initiator_session = ids.allocate_session(PeerId(1)); + let responder_session = ids.allocate_session(PeerId(2)); + let responder_session_id = responder_session.id(); let now = Instant::now(); // 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( + let send = BidiSession::new_initiator( ts_noise::core::Session { initiator_to_responder: k.into(), responder_to_initiator: k.into(), role: Role::Initiator, }, - session, - session, + initiator_session, + responder_session_id, now, ); - let mut recv = ReceiveSession::new(k.into(), session, now); + let mut recv = ReceiveSession::new(k.into(), responder_session, now); const CLEARTEXT: &[u8] = b"foobar"; let mut pkt = [PacketMut::from(CLEARTEXT)]; @@ -600,7 +577,7 @@ mod tests { let Ok(Message::TransportDataHeader(msg)) = Message::try_from(pkt[0].as_ref()) else { panic!("packet is not a valid TransportData message"); }; - assert_eq!(msg.receiver_id, session); + assert_eq!(msg.receiver_id, responder_session_id); assert_eq!(u64::from(msg.nonce), 0); assert!(recv.decrypt_one(&mut pkt[0])); @@ -611,7 +588,7 @@ mod tests { let Ok(Message::TransportDataHeader(msg)) = Message::try_from(pkt[0].as_ref()) else { panic!("packet is not a valid TransportData message"); }; - assert_eq!(msg.receiver_id, session); + assert_eq!(msg.receiver_id, responder_session_id); assert_eq!(u64::from(msg.nonce), 1); assert!(recv.decrypt_one(&mut pkt[0])); @@ -631,11 +608,11 @@ mod tests { } impl PeerSession { - fn allocate_id(&mut self) -> SessionId { + fn allocate_id(&mut self) -> SessionHandle { self.recv_id_prev = self.recv_id.take(); - let id = self.ids.allocate_session(PeerId(1)); - self.recv_id = Some(id); - id + let ret = self.ids.allocate_session(PeerId(1)); + self.recv_id = Some(ret.id()); + ret } fn handshake_with( @@ -645,40 +622,41 @@ mod tests { ) -> (Vec, Vec) { let (k1, k2): ([u8; 32], [u8; 32]) = rand::random(); let sid1 = self.allocate_id(); + let sid1_id = sid1.id(); let sid2 = other.allocate_id(); - let s1 = BidiSession::new( + let s1 = BidiSession::new_initiator( ts_noise::core::Session { initiator_to_responder: k1.into(), responder_to_initiator: k2.into(), role: Role::Initiator, }, - sid2, sid1, + sid2.id(), now, ); - let s2 = BidiSession::new( + let s2 = BidiSession::new_responder( ts_noise::core::Session { initiator_to_responder: k1.into(), responder_to_initiator: k2.into(), role: Role::Responder, }, sid2, - sid1, + sid1_id, now, ); - let (_, p1) = self.session.activate(s1, &mut self.ids, now, false); - let (_, p2) = other.session.activate(s2, &mut other.ids, now, false); + 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, &mut self.ids, now) + 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, &mut self.ids, now) + self.session.get_recv(hdr.receiver_id, now) } fn recv(&mut self, now: Instant, packets: Vec) -> Vec { @@ -858,27 +836,29 @@ mod tests { #[test] fn test_session_timers() { let k: [u8; 32] = rand::random(); - let id = SessionId::random(); + let mut ids = IdMap::default(); + let recv_session = ids.allocate_session(PeerId(1)); + 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(), id, now); + let recv = ReceiveSession::new(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)); assert!(recv.expired(now + SESSION_LIFETIME + epsilon)); let k2: [u8; 32] = rand::random(); - let id2 = SessionId::random(); - let bidi = BidiSession::new( + let bidi = BidiSession::new_initiator( ts_noise::core::Session { initiator_to_responder: k.into(), responder_to_initiator: k2.into(), role: Role::Initiator, }, - id, - id2, + bidi_session, + recv_session_id, now, ); assert!(!bidi.expired(now));