From a151aa936ee3727940a4a1546da31113f446de33 Mon Sep 17 00:00:00 2001 From: David Anderson Date: Tue, 18 Aug 2026 12:45:22 -0700 Subject: [PATCH] ts_tunnel: use carrier structs to pass global and peer state to handshakes Updates #339 Signed-off-by: David Anderson Change-Id: I3723948a6807c0c92946ca1c423c21be6a6a6964 --- ts_runtime/src/dataplane.rs | 10 +-- ts_tunnel/examples/handshake.rs | 10 +-- ts_tunnel/src/config.rs | 9 +++ ts_tunnel/src/endpoint.rs | 124 ++++++++++++----------------- ts_tunnel/src/handshake.rs | 134 ++++++++++++-------------------- 5 files changed, 116 insertions(+), 171 deletions(-) diff --git a/ts_runtime/src/dataplane.rs b/ts_runtime/src/dataplane.rs index ba149f36..d2107fcf 100644 --- a/ts_runtime/src/dataplane.rs +++ b/ts_runtime/src/dataplane.rs @@ -306,13 +306,11 @@ impl Message> for DataplaneActor { for &upsert in &msg.upserts { let (_, node) = msg.peers.get(&upsert).unwrap(); - wg.upsert_peer( + wg.upsert_peer(ts_tunnel::PeerConfig::new( ts_tunnel::PeerId(upsert.0), - ts_tunnel::PeerConfig { - key: node.node_key, - psk: [0u8; 32], - }, - ); + node.node_key, + [0u8; 32], + )); } for delete in &msg.deletions { diff --git a/ts_tunnel/examples/handshake.rs b/ts_tunnel/examples/handshake.rs index 289c319f..28fc2edf 100644 --- a/ts_tunnel/examples/handshake.rs +++ b/ts_tunnel/examples/handshake.rs @@ -87,14 +87,8 @@ async fn main() -> BoxResult<()> { let peer_id = ts_tunnel::PeerId(1); assert!( - ep.upsert_peer( - peer_id, - ts_tunnel::PeerConfig { - key: peer_key, - psk: [0; 32], - } - ) - .is_none() + ep.upsert_peer(ts_tunnel::PeerConfig::new(peer_id, peer_key, [0; 32])) + .is_none() ); let sock = tokio::net::UdpSocket::bind("0.0.0.0:0").await?; diff --git a/ts_tunnel/src/config.rs b/ts_tunnel/src/config.rs index 8b5d8079..cbb70f83 100644 --- a/ts_tunnel/src/config.rs +++ b/ts_tunnel/src/config.rs @@ -9,8 +9,17 @@ pub type Psk = ts_noise::core::Psk; /// The cryptographic configuration for a wireguard peer. pub struct PeerConfig { + /// The ID used to refer to this peer in [`crate::Endpoint`] APIs. + pub id: PeerId, /// The peer's public key. pub key: NodePublicKey, /// The pre-shared key to use for the peer, for post-quantum resistance. pub psk: Psk, } + +impl PeerConfig { + /// Return a [`PeerConfig`] with the given configuration. + pub fn new(id: PeerId, key: NodePublicKey, psk: Psk) -> Self { + Self { id, key, psk } + } +} diff --git a/ts_tunnel/src/endpoint.rs b/ts_tunnel/src/endpoint.rs index b9ebf872..31f29a39 100644 --- a/ts_tunnel/src/endpoint.rs +++ b/ts_tunnel/src/endpoint.rs @@ -23,7 +23,6 @@ use crate::{ const KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10); struct Peer { - id: PeerId, config: PeerConfig, session: Session, handshake: Handshake, @@ -32,11 +31,10 @@ struct Peer { send_another_keepalive: bool, } -impl Peer { - fn new(id: PeerId, config: PeerConfig) -> Self { +impl From for Peer { + fn from(config: PeerConfig) -> Self { let handshake = Handshake::new(&config.key); Self { - id, config, handshake, @@ -46,14 +44,16 @@ impl Peer { send_another_keepalive: false, } } +} +impl Peer { fn schedule_keepalive(&mut self, scheduler: &mut Scheduler, now: Instant) { if self.keepalive.is_some() { self.send_another_keepalive = true; return; } let tr = TimeRange::new_around(now + KEEPALIVE_TIMEOUT, Duration::from_secs(1)); - self.keepalive = Some(scheduler.add(tr, Event::MaybeSendKeepalive(self.id))); + self.keepalive = Some(scheduler.add(tr, Event::MaybeSendKeepalive(self.config.id))); } // TODO: consider replacing outparam with plain SendResult that supports merging. @@ -66,7 +66,7 @@ impl Peer { ) { if let Some(packets) = self.session.send(packets, now) { tracing::trace!("enqueueing packets to peer"); - out.queue_to_peer(self.id, packets); + out.queue_to_peer(self.config.id, packets); // Fall through to check if the session is in need of rotation. } @@ -132,26 +132,20 @@ impl Peer { now: Instant, out: &mut RecvResult, ) { - let Some(session) = self.handshake.finish( - packet, - &endpoint.my_key, - &self.config.psk, - &endpoint.my_cookie, - now, - ) else { + let Some(session) = self.handshake.finish(packet, endpoint, &self.config, now) else { tracing::error!("handshake failed to complete"); return; }; let (expiry, packets) = self.session.activate(session, now, true); - out.queue_to_peer(self.id, packets); + 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.id)), + .add(expiry, Event::ExpireSession(self.config.id)), ); } @@ -166,7 +160,7 @@ impl Peer { 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); + out.queue_to_local(self.config.id, packets); self.schedule_keepalive(&mut endpoint.scheduler, now); if self.session.needs_rotation(now) { self.start_handshake(endpoint, now, out); @@ -185,12 +179,12 @@ impl Peer { return; }; - out.queue_to_local(self.id, packets); + 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.id, packets_for_peer); + out.queue_to_peer(self.config.id, packets_for_peer); } if let Some(handle) = self.session_cleanup.take() { handle.cancel(); @@ -198,7 +192,7 @@ impl Peer { self.session_cleanup = Some( endpoint .scheduler - .add(expiry, Event::ExpireSession(self.id)), + .add(expiry, Event::ExpireSession(self.config.id)), ); } @@ -209,14 +203,11 @@ impl Peer { now: Instant, out: &mut RecvResult, ) { - let packet = self.handshake.respond( - handshake, - || endpoint.ids.allocate_session(self.id), - &self.config.psk, - now, - ); + let packet = self + .handshake + .respond(handshake, endpoint, &self.config, now); if let Some(packet) = packet { - out.queue_to_peer(self.id, [packet]); + out.queue_to_peer(self.config.id, [packet]); } } @@ -244,7 +235,7 @@ impl Peer { tracing::trace!("send keepalive: session expired, skipping"); return; }; - out.queue_to_peer(self.id, [packet]); + out.queue_to_peer(self.config.id, [packet]); self.keepalive = None; @@ -277,17 +268,8 @@ impl Peer { now: Instant, out: &mut impl QueueToPeer, ) { - let session_id = endpoint.ids.allocate_session(self.id); - let packet = self.handshake.initiate( - &mut endpoint.scheduler, - session_id, - &endpoint.my_key, - &self.config.key, - Event::HandshakeTimeout(self.id), - endpoint.timestamps.now(), - now, - ); - out.queue_to_peer(self.id, [packet]); + let packet = self.handshake.initiate(endpoint, &self.config, now); + out.queue_to_peer(self.config.id, [packet]); } } @@ -297,32 +279,38 @@ pub struct Endpoint { peers: HashMap, } -struct EndpointState { - my_key: NodeKeyPair, +pub struct EndpointState { + pub my_key: NodeKeyPair, + + pub my_cookie: MACReceiver, + pub ids: IdMap, + pub timestamps: TAI64NClock, + pub scheduler: Scheduler, +} - my_cookie: MACReceiver, - ids: IdMap, - timestamps: TAI64NClock, - scheduler: Scheduler, +impl EndpointState { + pub fn new(my_key: NodeKeyPair) -> Self { + let my_cookie = MACReceiver::new(&my_key.public); + Self { + my_key, + my_cookie, + ids: IdMap::default(), + timestamps: TAI64NClock::default(), + scheduler: Scheduler::default(), + } + } } impl Endpoint { /// Construct a new endpoint with the given keypair. pub fn new(my_key: NodeKeyPair) -> Self { - let my_cookie = MACReceiver::new(&my_key.public); Self { - state: EndpointState { - my_key, - my_cookie, - ids: Default::default(), - timestamps: Default::default(), - scheduler: Default::default(), - }, + state: EndpointState::new(my_key), peers: HashMap::new(), } } - /// Insert a peer if it doesn't exist, otherwise update the peer with the given `id` + /// Insert a peer if it doesn't exist, otherwise update the peer with the given `cfg.id` /// with the given config. /// /// Returns the old [`PeerConfig`] if there was one. @@ -331,23 +319,23 @@ impl Endpoint { /// /// If the [`NodePublicKey`] in the new [`PeerConfig`] collides with an existing key /// for a different [`PeerId`]. - pub fn upsert_peer(&mut self, id: PeerId, mut cfg: PeerConfig) -> Option { - match self.peers.get_mut(&id) { + pub fn upsert_peer(&mut self, mut cfg: PeerConfig) -> Option { + match self.peers.get_mut(&cfg.id) { Some(peer) => { if peer.config.key != cfg.key { self.state.ids.remove_peer(&peer.config.key); - self.state.ids.add_peer(id, &cfg.key); + self.state.ids.add_peer(cfg.id, &cfg.key); } core::mem::swap(&mut peer.config, &mut cfg); Some(cfg) } None => { - if !self.state.ids.add_peer(id, &cfg.key) { + if !self.state.ids.add_peer(cfg.id, &cfg.key) { panic!("nodekey collision"); } - self.peers.insert(id, Peer::new(id, cfg)); + self.peers.insert(cfg.id, Peer::from(cfg)); None } } @@ -673,24 +661,12 @@ mod tests { let mut b = Endpoint::new(key_b.clone()); let psk = rand::random(); assert!( - a.upsert_peer( - PeerId(1), - PeerConfig { - key: key_b.public, - psk, - } - ) - .is_none() + a.upsert_peer(PeerConfig::new(PeerId(1), key_b.public, psk)) + .is_none() ); assert!( - b.upsert_peer( - PeerId(1), - PeerConfig { - key: key_a.public, - psk, - } - ) - .is_none() + b.upsert_peer(PeerConfig::new(PeerId(1), key_a.public, psk)) + .is_none() ); Self { a, diff --git a/ts_tunnel/src/handshake.rs b/ts_tunnel/src/handshake.rs index d321d189..1016e569 100644 --- a/ts_tunnel/src/handshake.rs +++ b/ts_tunnel/src/handshake.rs @@ -6,12 +6,12 @@ use std::{ use ts_keys::{NodeKeyPair, NodePublicKey}; use ts_noise::ikpsk2; use ts_packet::PacketMut; -use ts_time::{Handle, Scheduler, TimeRange}; +use ts_time::{Handle, TimeRange}; use zerocopy::IntoBytes; use crate::{ - config::Psk, - endpoint::Event, + PeerConfig, + endpoint::{EndpointState, Event}, ids::SessionHandle, macs::{MACReceiver, MACSender, Mac}, messages::*, @@ -171,31 +171,29 @@ impl Handshake { /// Starting a new handshake abandons any other handshake that was already in flight. pub fn initiate( &mut self, - scheduler: &mut Scheduler, - session_handle: SessionHandle, - my_static: &NodeKeyPair, - peer_static: &NodePublicKey, - timeout_event: Event, - timestamp: TAI64N, + endpoint: &mut EndpointState, + peer: &PeerConfig, now: Instant, ) -> PacketMut { + let session_handle = endpoint.ids.allocate_session(peer.id); + let mut pkt = HandshakeInitiation { sender_id: session_handle.id(), ..Default::default() }; let noise = ikpsk2::SentHandshake::new( - my_static.into(), - peer_static.into(), + (&endpoint.my_key).into(), + peer.key.into(), PROLOGUE, - timestamp, + endpoint.timestamps.now(), pkt.noise.as_mut_bytes(), ); let mut pkt = PacketMut::from(pkt.as_bytes()); let mac1 = self.cookie_sender.write_macs(pkt.as_mut()); let tr = TimeRange::new_around(now + HANDSHAKE_TIMEOUT, Duration::from_millis(500)); - let timeout = scheduler.add(tr, timeout_event); + let timeout = endpoint.scheduler.add(tr, Event::HandshakeTimeout(peer.id)); self.state = State::Initiated(SentHandshake { responder_to_initiator_handle: session_handle, @@ -215,30 +213,29 @@ impl Handshake { pub fn finish( &mut self, packet: &mut HandshakeResponse, - endpoint_static: &NodeKeyPair, - psk: &Psk, - cookies: &MACReceiver, + endpoint: &mut EndpointState, + peer: &PeerConfig, now: Instant, ) -> Option { let mut sent_handshake = self.state.take_if_initiated()?; - if !cookies.verify_macs(packet.as_bytes()) { + if !endpoint.my_cookie.verify_macs(packet.as_bytes()) { self.state = State::Initiated(sent_handshake); return None; }; - let session_keys = - match sent_handshake - .noise - .try_finish(&mut packet.noise, endpoint_static.into(), psk) - { - Ok(session_keys) => session_keys, - Err(handshake) => { - sent_handshake.noise = handshake; - self.state = State::Initiated(sent_handshake); - return None; - } - }; + let session_keys = match sent_handshake.noise.try_finish( + &mut packet.noise, + (&endpoint.my_key).into(), + &peer.psk, + ) { + Ok(session_keys) => session_keys, + Err(handshake) => { + sent_handshake.noise = handshake; + self.state = State::Initiated(sent_handshake); + return None; + } + }; let session = BidiSession::new_initiator( session_keys, @@ -261,8 +258,8 @@ impl Handshake { pub fn respond( &mut self, handshake: ReceivedHandshake, - allocate_session_handle: impl FnOnce() -> SessionHandle, - psk: &Psk, + endpoint: &mut EndpointState, + peer: &PeerConfig, now: Instant, ) -> Option { if let Some(last_seen_timestamp) = self.last_seen_timestamp @@ -272,13 +269,15 @@ impl Handshake { return None; } - let session_handle = allocate_session_handle(); + let session_handle = endpoint.ids.allocate_session(peer.id); let mut response = HandshakeResponse { sender_id: session_handle.id(), receiver_id: handshake.responder_to_initiator_id, ..Default::default() }; - let session_keys = handshake.noise.finish(psk, response.noise.as_mut_bytes()); + let session_keys = handshake + .noise + .finish(&peer.psk, response.noise.as_mut_bytes()); let mut pkt = PacketMut::from(response.as_bytes()); self.cookie_sender.write_macs(pkt.as_mut()); @@ -348,11 +347,10 @@ impl Handshake { #[cfg(test)] mod tests { use ts_keys::NodeKeyPair; - use ts_time::Scheduler; use zerocopy::TryFromBytes; use super::*; - use crate::{Event::HandshakeTimeout, PeerId, ids::IdMap}; + use crate::PeerId; #[test] fn test_handshake() { @@ -360,42 +358,27 @@ mod tests { let psk = rand::random(); // Peer A sends a handshake initiation... - let a_mac_recv = MACReceiver::new(&a_static.public); - 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 mut a_sched = Scheduler::default(); - + let mut a_state = EndpointState::new(a_static.clone()); let mut a_handshake = Handshake::new(&b_static.public); - let init_pkt = a_handshake.initiate( - &mut a_sched, - a_session, - &a_static, - &b_static.public, - HandshakeTimeout(PeerId(0)), - a_init_time, - Instant::now(), - ); + let a_peer = PeerConfig::new(PeerId(1), b_static.public, psk); + let init_pkt = a_handshake.initiate(&mut a_state, &a_peer, Instant::now()); // Peer B receives it and responds let b_mac_recv = MACReceiver::new(&b_static.public); let init_pkt = ReceivedHandshake::new(init_pkt, &b_static, &b_mac_recv) .expect("B should parse the initiation message"); - let mut b_handshake = Handshake::new(&a_static.public); + let mut b_handshake = Handshake::new(&a_state.my_key.public); + let mut b_state = EndpointState::new(b_static); + let b_peer = PeerConfig::new(PeerId(2), a_static.public, psk); let mut response_pkt = b_handshake - .respond( - init_pkt, - || ids.allocate_session(PeerId(2)), - &psk, - Instant::now(), - ) + .respond(init_pkt, &mut b_state, &b_peer, Instant::now()) .expect("B should respond to handshake"); // Peer A receives response, sends confirmation let response_pkt = HandshakeResponse::try_mut_from_bytes(response_pkt.as_mut()) .expect("response_pkt should be a valid handshake response message"); let Some(mut a_session) = - a_handshake.finish(response_pkt, &a_static, &psk, &a_mac_recv, Instant::now()) + a_handshake.finish(response_pkt, &mut a_state, &a_peer, Instant::now()) else { panic!("failed to process handshake response from peer B"); }; @@ -419,37 +402,23 @@ 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_recv = MACReceiver::new(&a_static.public); - let a_session = ids.allocate_session(PeerId(1)); // A wants to receive at this ID - let a_init_time = TAI64N::now(); - let mut a_scheduler = Scheduler::default(); + let mut a_state = EndpointState::new(a_static.clone()); let mut a_handshake = Handshake::new(&b_static.public); - let init_pkt = a_handshake.initiate( - &mut a_scheduler, - a_session, - &a_static, - &b_static.public, - HandshakeTimeout(PeerId(0)), - a_init_time, - Instant::now(), - ); + let a_peer = PeerConfig::new(PeerId(0), b_static.public, psk); + let init_pkt = a_handshake.initiate(&mut a_state, &a_peer, Instant::now()); // B receives and responds let b_mac_recv = MACReceiver::new(&b_static.public); let init_pkt = ReceivedHandshake::new(init_pkt, &b_static, &b_mac_recv) .expect("B should parse the initiation message"); let mut b_handshake = Handshake::new(&a_static.public); + let mut b_state = EndpointState::new(b_static); + let b_peer = PeerConfig::new(PeerId(2), a_static.public, psk); let mut response_pkt = b_handshake - .respond( - init_pkt, - || ids.allocate_session(PeerId(2)), - &psk, - Instant::now(), - ) + .respond(init_pkt, &mut b_state, &b_peer, Instant::now()) .expect("B responds to handshake"); // A receives several invalid responses: one with bad MACs, one with a bad Noise handshake @@ -457,9 +426,8 @@ mod tests { a_handshake .finish( &mut HandshakeResponse::default(), - &a_static, - &psk, - &a_mac_recv, + &mut a_state, + &a_peer, Instant::now() ) .is_none() @@ -469,7 +437,7 @@ mod tests { corrupt_pkt.noise[3] = corrupt_pkt.noise[3].wrapping_add(1); assert!( a_handshake - .finish(corrupt_pkt, &a_static, &psk, &a_mac_recv, Instant::now()) + .finish(corrupt_pkt, &mut a_state, &a_peer, Instant::now()) .is_none() ); @@ -478,7 +446,7 @@ mod tests { .expect("response_pkt should be a valid handshake response message"); assert!( a_handshake - .finish(response_pkt, &a_static, &psk, &a_mac_recv, Instant::now()) + .finish(response_pkt, &mut a_state, &a_peer, Instant::now()) .is_some() ); }