Return next scheduled update from next_timer_update

Author: thomaseizinger

boringtun/src/noise/timers.rs

@@ -237,6 +237,13 @@ impl Tunn {
         }
     }
 
+    fn next_expired_session(&self) -> Option<Instant> {
+        self.sessions
+            .iter()
+            .flat_map(|s| Some(s.as_ref()?.established_at() + REJECT_AFTER_TIME))
+            .min()
+    }
+
     #[deprecated(note = "Prefer `Timers::update_timers_at` to avoid time-impurity")]
     pub fn update_timers<'a>(&mut self, dst: &'a mut [u8]) -> TunnResult<'a> {
         self.update_timers_at(dst, Instant::now())
@@ -420,7 +427,36 @@ impl Tunn {
     ///
     /// If this returns `None`, you may call it at your usual desired precision (usually once a second is enough).
     pub fn next_timer_update(&self) -> Option<Instant> {
-        self.timers.send_handshake_at
+        // Mimic the `update_timers_at` function: If we have a handshake scheduled, other timers don't matter.
+        if let Some(scheduled_handshake) = self.timers.send_handshake_at {
+            return Some(scheduled_handshake);
+        }
+
+        let common_timers = [
+            self.next_expired_session(),
+            self.handshake.cookie_expiration(),
+            Some(self.timers.reject_after_time()),
+        ]
+        .into_iter();
+
+        if let Some(rekey_timeout) = self.handshake.rekey_timeout() {
+            earliest(
+                common_timers.chain([Some(rekey_timeout), Some(self.timers.rekey_attempt_time())]),
+            )
+        } else {
+            // Persistent keep-alive only makes sense if the current session is active.
+            let persistent_keepalive = self.sessions[self.current % N_SESSIONS]
+                .as_ref()
+                .and_then(|_| self.timers.next_persistent_keepalive());
+
+            earliest(common_timers.chain([
+                self.timers.rekey_after_time_on_send(),
+                self.timers.reject_after_time_on_receive(),
+                self.timers.rekey_after_time_without_response(),
+                self.timers.keepalive_after_time_without_send(),
+                persistent_keepalive,
+            ]))
+        }
     }
 
     #[deprecated(note = "Prefer `Tunn::time_since_last_handshake_at` to avoid time-impurity")]
@@ -449,3 +485,7 @@ impl Tunn {
         }
     }
 }
+
+fn earliest(instants: impl IntoIterator<Item = Option<Instant>>) -> Option<Instant> {
+    instants.into_iter().flatten().min()
+}