ts_tunnel: give up on handshake initiation after enough retries

ts_tunnel/src/endpoint.rs

@@ -8,7 +8,7 @@ use ts_time::{Handle, Scheduler, TimeRange};
 
 use crate::{
     config::{PeerConfig, PeerId},
-    handshake::{Handshake, ReceivedHandshake},
+    handshake::{Handshake, ReceivedHandshake, TimeoutError},
     ids::IdMap,
     macs::MACReceiver,
     messages::{HandshakeResponse, Message, MessageMut, SessionId},
@@ -77,17 +77,7 @@ impl Peer {
             }
         }
 
-        if self.handshake.is_active() {
-            tracing::trace!("handshake is already in-flight, bail");
-            return;
-        }
-
-        if !self.session.needs_rotation(now) {
-            tracing::trace!("session does not need new handshake");
-            return;
-        }
-
-        self.start_handshake(endpoint, now, out);
+        self.maybe_start_handshake(endpoint, now, out);
     }
 
     #[tracing::instrument(skip_all, fields(?session_id, n_packets = packets.len()))]
@@ -172,9 +162,7 @@ impl Peer {
             if !packets.is_empty() {
                 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);
-                }
+                self.maybe_start_handshake(endpoint, now, out);
             }
             return;
         }
@@ -217,6 +205,7 @@ impl Peer {
         }
     }
 
+    #[tracing::instrument(skip_all, fields(now, peer_id = ?self.config.id))]
     fn handshake_timeout(
         &mut self,
         endpoint: &mut EndpointState,
@@ -225,12 +214,14 @@ impl Peer {
     ) {
         self.check_invariants(now);
 
-        if !self.handshake.is_active() {
-            // Handshake completed prior to timeout firing.
-            return;
+        match self.handshake.timeout(endpoint, &self.config, now) {
+            Ok(packet) => out.queue_to_peer(self.config.id, [packet]),
+            Err(TimeoutError::WrongHandshakeState) => {}
+            Err(TimeoutError::HandshakeTimeout) => {
+                tracing::warn!("timeout while waiting for handshake response");
+                self.queue.clear();
+            }
         }
-
-        self.start_handshake(endpoint, now, out);
     }
 
     fn send_keepalive(
@@ -266,12 +257,23 @@ impl Peer {
         }
     }
 
-    fn start_handshake(
+    #[tracing::instrument(skip_all, fields(now, peer_id = ?self.config.id))]
+    fn maybe_start_handshake(
         &mut self,
         endpoint: &mut EndpointState,
         now: Instant,
         out: &mut impl QueueToPeer,
     ) {
+        if self.handshake.is_active() {
+            tracing::trace!("handshake is already in flight, bail");
+            return;
+        }
+
+        if !self.session.needs_rotation(now) {
+            tracing::trace!("session does not need new handshake");
+            return;
+        }
+
         let packet = self.handshake.initiate(endpoint, &self.config, now);
         out.queue_to_peer(self.config.id, [packet]);
     }
@@ -594,6 +596,7 @@ mod tests {
     use super::*;
     use crate::{
         config::PeerConfig,
+        handshake::{HANDSHAKE_FAILURE_TIMEOUT, HANDSHAKE_RETRY_TIMEOUT},
         messages::{HandshakeInitiation, TransportDataHeader},
     };
 
@@ -759,6 +762,19 @@ mod tests {
             }
         }
 
+        fn events(now: Instant, ep: &mut Endpoint, to_peer: &mut Vec<PacketMut>) {
+            let id = PeerId(1);
+            let mut acts = ep.dispatch_events(now);
+            match acts.to_peers.len() {
+                0 => (),
+                1 => to_peer.extend(acts.to_peers.remove(&id).unwrap()),
+                _ => panic!(
+                    "got packets for {} peers, expected 1 or 0",
+                    acts.to_peers.len()
+                ),
+            }
+        }
+
         /// Send cleartext packets from endpoint A to endpoint B.
         ///
         /// This may result in the queuing of encrypted packets from A to B, which can be
@@ -805,6 +821,34 @@ mod tests {
             );
         }
 
+        /// Process pending events at A.
+        ///
+        /// This may result in the queuing of encrypted packets from A to B, which can be inspected
+        /// with [`EndpointPair::assert_a_to_b`].
+        pub fn event_a(&mut self, now: Instant) {
+            EndpointPair::events(now, &mut self.a, &mut self.a_to_b);
+        }
+
+        /// Process pending events at B.
+        ///
+        /// This may result in the queuing of encrypted packets from B to A, which can be inspected
+        /// with [`EndpointPair::assert_b_to_a`].
+        #[allow(dead_code)]
+        pub fn event_b(&mut self, now: Instant) {
+            EndpointPair::events(now, &mut self.b, &mut self.b_to_a);
+        }
+
+        /// Drop in-flight packets from A to B.
+        pub fn drop_a_to_b(&mut self) {
+            self.a_to_b.clear();
+        }
+
+        /// Drop in-flight packets from B to A.
+        #[allow(dead_code)]
+        pub fn drop_b_to_a(&mut self) {
+            self.b_to_a.clear();
+        }
+
         /// Assert that packets currently in flight from A to B match the expected packet shapes.
         ///
         /// The inspection is non-destructive, in-flight packets will be delivered by the next call
@@ -959,4 +1003,50 @@ mod tests {
         p.recv_at_b(t.at(129));
         assert_no_packets(&p.received_at_b);
     }
+
+    #[test]
+    fn test_handshake_timeout() {
+        use PacketMatcher::*;
+
+        let mut p = EndpointPair::new();
+        let t = TestClock::new();
+
+        // A sends to B. Triggers a handshake initiation, which is lost.
+        p.send_from_a(t.at(0), [packet(1)]);
+        p.assert_a_to_b([HandshakeInitiation]);
+        p.drop_a_to_b();
+
+        // Timeout fires repeatedly, A retries.
+        let mut at = 0;
+        let timeout = HANDSHAKE_RETRY_TIMEOUT.as_secs();
+        let deadline = HANDSHAKE_FAILURE_TIMEOUT.as_secs();
+        while at + timeout < deadline {
+            at += timeout;
+            p.event_a(t.at(at));
+            p.assert_a_to_b([HandshakeInitiation]);
+            p.drop_a_to_b();
+        }
+
+        // Timeout fires for the final time, A gives up.
+        p.event_a(t.at(at + timeout));
+        assert_no_packets(&p.a_to_b);
+
+        // A explicitly sends to B again. Triggers a new handshake initiation, which completes
+        // this time.
+        p.send_from_a(t.at(at + 10), [packet(2)]);
+        p.assert_a_to_b([HandshakeInitiation]);
+
+        p.recv_at_b(t.at(at + 11));
+        p.assert_b_to_a([HandshakeResponse]);
+        assert_no_packets(&p.received_at_b);
+
+        p.recv_at_a(t.at(at + 12));
+        p.assert_a_to_b([TransportData]);
+        assert_no_packets(&p.received_at_a);
+
+        p.recv_at_b(t.at(at + 13));
+        assert_no_packets(&p.b_to_a);
+        // B never receives packet(1), it was dropped when the first handshake failed.
+        p.assert_received_at_b([packet(2)]);
+    }
 }

ts_tunnel/src/handshake.rs

@@ -1,5 +1,6 @@
 use std::{
     mem::replace,
+    ops::Add,
     time::{Duration, Instant},
 };
 
@@ -70,9 +71,12 @@ struct SentHandshake {
     timeout: Handle<Event>,
     // The mac1 of the transmitted handshake, required to process CookieReply messages.
     mac1: Mac,
+    // The final deadline for this handshake attempt, including all retries.
+    deadline: Instant,
 }
 
-const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(5);
+pub(crate) const HANDSHAKE_RETRY_TIMEOUT: Duration = Duration::from_secs(5);
+pub(crate) const HANDSHAKE_FAILURE_TIMEOUT: Duration = Duration::from_secs(90);
 
 /// The state machine portion of a handshake.
 #[derive(Default)]
@@ -135,6 +139,19 @@ pub struct Handshake {
     state: State,
 }
 
+/// Error returned from [`Handshake::timeout`].
+#[derive(Copy, Clone, Debug)]
+pub enum TimeoutError {
+    /// Handshake is not in the initiated state. Likely this means the timeout notification raced
+    /// with handshake completion. Either way, no action is required.
+    WrongHandshakeState,
+    /// The handshake has reached the maximum number of timeout retries, and has permanently failed.
+    /// Caller should clear out any state that was waiting for a session (e.g. packet buffers) and
+    /// not try another handshake initiation until [`Endpoint::send`] is called again with a new
+    /// packet for the peer.
+    HandshakeTimeout,
+}
+
 impl Handshake {
     pub fn new(peer_key: &NodePublicKey) -> Self {
         let cookie_sender = MACSender::new(peer_key);
@@ -169,11 +186,48 @@ impl Handshake {
     /// Start a new handshake in the initiator role.
     ///
     /// Starting a new handshake abandons any other handshake that was already in flight.
+    ///
+    /// [`Handshake::initiate`] schedules an [`Event::HandshakeTimeout`] event. The caller must call
+    /// [`Handshake::timeout`] when that event fires to continue advancing the handshake state machine.
     pub fn initiate(
         &mut self,
         endpoint: &mut EndpointState,
         peer: &PeerConfig,
         now: Instant,
+    ) -> PacketMut {
+        let deadline = now.add(HANDSHAKE_FAILURE_TIMEOUT);
+        self.initiate_inner(endpoint, peer, now, deadline)
+    }
+
+    /// Process a handshake initiation timeout.
+    ///
+    /// Returns `Ok(packet)` if the handshake can continue to retry,
+    /// `Err(TimeoutError::HandshakeTimeout)` if the maximum number of retries has been reached,
+    /// or `Err(TimeoutError::WrongHandshakeState)` if the timeout is no longer applicable.
+    pub fn timeout(
+        &mut self,
+        endpoint: &mut EndpointState,
+        peer: &PeerConfig,
+        now: Instant,
+    ) -> Result<PacketMut, TimeoutError> {
+        let handshake = self
+            .state
+            .take_if_initiated()
+            .ok_or(TimeoutError::WrongHandshakeState)?;
+
+        if now >= handshake.deadline {
+            return Err(TimeoutError::HandshakeTimeout);
+        }
+
+        Ok(self.initiate_inner(endpoint, peer, now, handshake.deadline))
+    }
+
+    fn initiate_inner(
+        &mut self,
+        endpoint: &mut EndpointState,
+        peer: &PeerConfig,
+        now: Instant,
+        deadline: Instant,
     ) -> PacketMut {
         let session_handle = endpoint.ids.allocate_session(peer.id);
 
@@ -192,14 +246,15 @@ impl Handshake {
         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 tr = TimeRange::new_around(now + HANDSHAKE_RETRY_TIMEOUT, Duration::from_millis(500));
         let timeout = endpoint.scheduler.add(tr, Event::HandshakeTimeout(peer.id));
 
         self.state = State::Initiated(SentHandshake {
             responder_to_initiator_handle: session_handle,
             noise,
             timeout,
             mac1,
+            deadline,
         });
 
         pkt