Skip to content

ts_tunnel: use carrier structs to pass global and peer state to handshakes #376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danderson merged 1 commit into from
Aug 19, 2026
Merged

ts_tunnel: use carrier structs to pass global and peer state to handshakes #376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions ts_runtime/src/dataplane.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,13 +306,11 @@ impl Message<Arc<PeerState>> 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 {
Expand Down
10 changes: 2 additions & 8 deletions ts_tunnel/examples/handshake.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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?;
Expand Down
9 changes: 9 additions & 0 deletions ts_tunnel/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
}
}
124 changes: 50 additions & 74 deletions ts_tunnel/src/endpoint.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ use crate::{
const KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10);

struct Peer {
id: PeerId,
config: PeerConfig,
session: Session,
handshake: Handshake,
Expand All @@ -32,11 +31,10 @@ struct Peer {
send_another_keepalive: bool,
}

impl Peer {
fn new(id: PeerId, config: PeerConfig) -> Self {
impl From<PeerConfig> for Peer {
fn from(config: PeerConfig) -> Self {
let handshake = Handshake::new(&config.key);
Self {
id,
config,
handshake,

Expand All @@ -46,14 +44,16 @@ impl Peer {
send_another_keepalive: false,
}
}
}

impl Peer {
fn schedule_keepalive(&mut self, scheduler: &mut Scheduler<Event>, 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.
Expand All @@ -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.
}

Expand Down Expand Up @@ -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)),
);
}

Expand All @@ -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);
Expand All @@ -185,20 +179,20 @@ 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();
}
self.session_cleanup = Some(
endpoint
.scheduler
.add(expiry, Event::ExpireSession(self.id)),
.add(expiry, Event::ExpireSession(self.config.id)),
);
}

Expand All @@ -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]);
}
}

Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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]);
}
}

Expand All @@ -297,32 +279,38 @@ pub struct Endpoint {
peers: HashMap<PeerId, Peer>,
}

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<Event>,
}

my_cookie: MACReceiver,
ids: IdMap,
timestamps: TAI64NClock,
scheduler: Scheduler<Event>,
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.
Expand All @@ -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<PeerConfig> {
match self.peers.get_mut(&id) {
pub fn upsert_peer(&mut self, mut cfg: PeerConfig) -> Option<PeerConfig> {
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
}
}
Expand Down Expand Up @@ -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,
Expand Down
Loading