Skip to content

ts_tunnel: make IdMap hand out handles for newly allocated sessions #366

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

Open
danderson wants to merge 1 commit into from
Open

ts_tunnel: make IdMap hand out handles for newly allocated sessions #366

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
32 changes: 13 additions & 19 deletions ts_tunnel/src/endpoint.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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();
Expand All @@ -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);
Expand All @@ -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);
}
Expand Down Expand Up @@ -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(
Expand All @@ -255,7 +252,6 @@ impl Peer {
return;
}

endpoint.ids.remove_handshake_session(&self.handshake);
self.handshake = Handshake::None;

self.start_handshake(endpoint, now, out);
Expand All @@ -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;
};
Expand All @@ -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();
Expand All @@ -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,
Expand All @@ -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));

Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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;
};
Expand Down Expand Up @@ -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);
}
}
}
Expand Down
39 changes: 17 additions & 22 deletions ts_tunnel/src/handshake.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use zerocopy::IntoBytes;
use crate::{
config::Psk,
endpoint::Event,
ids::SessionHandle,
macs::{MACReceiver, MACSender, Mac},
messages::*,
session::BidiSession,
Expand Down Expand Up @@ -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,
Expand All @@ -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()
};

Expand All @@ -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<TAI64N>,
}

Expand All @@ -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<SessionId> {
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<Event>, Mac)> {
match std::mem::replace(self, Handshake::None) {
Handshake::Initiated(sent, timeout, mac) => Some((sent, timeout, mac)),
Expand All @@ -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,
Expand Down Expand Up @@ -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,
);

Expand Down Expand Up @@ -268,6 +260,7 @@ mod tests {
use zerocopy::TryFromBytes;

use super::*;
use crate::{PeerId, ids::IdMap};

#[test]
fn test_handshake() {
Expand All @@ -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);
Expand All @@ -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());

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

Expand Down
79 changes: 53 additions & 26 deletions ts_tunnel/src/ids.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<SessionId, PeerId>,
sessions: Arc<Mutex<HashMap<SessionId, PeerId>>>,
// TODO: track recently abandoned session IDs, avoid reusing them for
// one or two session lifetimes to avoid confusion with reordered packets.
node_keys: HashMap<NodePublicKey, PeerId>,
Expand All @@ -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<PeerId> {
self.sessions.lock().unwrap().get(key).copied()
}

/// Add a peer handle for communicating with the given peer pubkey.
Expand All @@ -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<Mutex<HashMap<SessionId, PeerId>>>,
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<SessionId> 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();
}
}
Loading