From 1a5d36a3377bbe988cac903f23d865cb4e904b09 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Wed, 15 Jan 2025 00:19:31 +0100
Subject: [PATCH] Track time of first packet without reply

---
 boringtun/src/noise/timers.rs | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/boringtun/src/noise/timers.rs b/boringtun/src/noise/timers.rs
index 97d12955..0052650b 100644
--- a/boringtun/src/noise/timers.rs
+++ b/boringtun/src/noise/timers.rs
@@ -48,7 +48,9 @@ pub struct Timers {
     /// Did we receive data without sending anything back?
     want_keepalive: bool,
     /// Did we send data without hearing back?
-    want_handshake: bool,
+    ///
+    /// If `Some`, holds the _first_ instant when we sent a packet.
+    want_handshake_since: Option<Instant>,
     persistent_keepalive: usize,
     /// Should this timer call reset rr function (if not a shared rr instance)
     pub(super) should_reset_rr: bool,
@@ -69,7 +71,7 @@ impl Timers {
             is_initiator: false,
             timers: [now; TimerName::Top as usize],
             want_keepalive: Default::default(),
-            want_handshake: Default::default(),
+            want_handshake_since: Default::default(),
             persistent_keepalive: usize::from(persistent_keepalive.unwrap_or(0)),
             should_reset_rr: reset_rr,
             send_handshake_at: None,
@@ -118,9 +120,7 @@ impl Timers {
     }
 
     pub(crate) fn rekey_after_time_without_response(&self) -> Option<Instant> {
-        if !self.want_handshake {
-            return None;
-        }
+        let first_packet_without_reply = self.want_handshake_since?;
 
         let last_data_packet_sent = self[TimeLastDataPacketSent];
         let last_packet_received = self[TimeLastPacketReceived];
@@ -130,7 +130,7 @@ impl Timers {
             return None;
         }
 
-        Some(last_packet_received + KEEPALIVE_TIMEOUT + REKEY_TIMEOUT)
+        Some(first_packet_without_reply + KEEPALIVE_TIMEOUT + REKEY_TIMEOUT)
     }
 
     pub(crate) fn keepalive_after_time_without_send(&self) -> Option<Instant> {
@@ -169,7 +169,7 @@ impl Timers {
         for t in &mut self.timers[..] {
             *t = now;
         }
-        self.want_handshake = false;
+        self.want_handshake_since = None;
         self.want_keepalive = false;
     }
 }
@@ -192,12 +192,20 @@ impl Tunn {
         match timer_name {
             TimeLastPacketReceived => {
                 self.timers.want_keepalive = true;
-                self.timers.want_handshake = false;
+                self.timers.want_handshake_since = None;
             }
             TimeLastPacketSent => {
-                self.timers.want_handshake = true;
                 self.timers.want_keepalive = false;
             }
+            TimeLastDataPacketSent => {
+                match self.timers.want_handshake_since {
+                    Some(_) => {} // Already waiting for a reply, don't update the timestamp.
+                    None => {
+                        // This is the first packet to not be replied to, start tracking the time.
+                        self.timers.want_handshake_since = Some(now)
+                    }
+                }
+            }
             _ => {}
         }