Attempt handshakes if keepalive is set

In two separate cases it is possible enter
the state of no automatic handshake initiation
triggered by the keepalive which should happen.

This is done by having an expiration guard which is
a correct approach if there's no keepalive, but incorrect
if there's keepalive set. This makes it so that only
a packet which is sent to the tunnel for encapsulation
triggers the handshake.

Signed-off-by: Lukas Pukenis <[email protected]>

neptun/src/noise/mod.rs

@@ -676,7 +676,9 @@ impl Tunn {
 #[cfg(test)]
 mod tests {
     #[cfg(feature = "mock-instant")]
-    use crate::noise::timers::{REKEY_AFTER_TIME, REKEY_TIMEOUT};
+    use crate::noise::timers::{
+        REJECT_AFTER_TIME, REKEY_AFTER_TIME, REKEY_ATTEMPT_TIME, REKEY_TIMEOUT,
+    };
 
     use super::*;
     use rand_core::{OsRng, RngCore};
@@ -818,9 +820,128 @@ mod tests {
         assert!(matches!(their_tun.update_timers(&mut []), TunnResult::Done));
     }
 
+    #[cfg(feature = "mock-instant")]
+    #[test]
+    fn no_handshake_if_no_keepalive() {
+        // If no keepalive is set then handshake is initiated only on `encapsulate()` after reading
+        // the packet from the tunnel(app data -> tun)
+        let (mut my_tun, _) = create_two_tuns();
+
+        my_tun.set_persistent_keepalive(0);
+
+        let mut my_dst = [0u8; 1024];
+        let res = my_tun.update_timers(&mut my_dst);
+        assert!(matches!(res, TunnResult::Done));
+    }
+
+    #[cfg(feature = "mock-instant")]
+    #[test]
+    fn handshake_if_keepalive() {
+        // No data is needed to trigger the handshake if keepalive is set
+        let (mut my_tun, _) = create_two_tuns();
+
+        my_tun.set_persistent_keepalive(25);
+
+        let mut my_dst = [0u8; 1024];
+        let res = my_tun.update_timers(&mut my_dst);
+        assert!(matches!(res, TunnResult::WriteToNetwork(_)));
+    }
+
+    #[cfg(feature = "mock-instant")]
+    #[test]
+    fn handshake_if_keepalive_after_wg_rekey_attempt_time() {
+        let (mut my_tun, _) = create_two_tuns();
+
+        my_tun.set_persistent_keepalive(25);
+
+        let mut my_dst = [0u8; 1024];
+        let res = my_tun.update_timers(&mut my_dst);
+        assert!(matches!(res, TunnResult::WriteToNetwork(_)));
+
+        mock_instant::MockClock::advance(REKEY_ATTEMPT_TIME.into());
+        let res = my_tun.update_timers(&mut my_dst);
+        assert!(matches!(res, TunnResult::WriteToNetwork(_)));
+    }
+
+    #[cfg(feature = "mock-instant")]
+    #[test]
+    fn assert_no_handshake_if_no_keepalive_after_3x_wg_reject_after_time() {
+        // WireGuard#6.3:
+        // <...>
+        // If no new secure session is created after (Reject-After-Time × 3) seconds,
+        // the current secure session, the previous secure session, and potentially
+        // the next secure session are discarded and zeroed out, in addition to any
+        // possible partially-completed handshake states and ephemeral keys.
+        // <...>
+        //
+        // Though it doesn't mention explicitly but observing Linux kernel implementation
+        // shows that non-zero Persistent-Keepalive must trigger the handshake.
+
+        let (mut my_tun, mut their_tun) = create_two_tuns();
+
+        // Persistent-Keepalive value doesn't matter, while it's non zero
+        my_tun.set_persistent_keepalive(0);
+
+        let mut my_dst = [0u8; 1024];
+
+        // Establish the session
+        mock_instant::MockClock::advance(Duration::from_secs(1));
+        assert!(matches!(their_tun.update_timers(&mut []), TunnResult::Done));
+        assert!(matches!(
+            my_tun.update_timers(&mut my_dst),
+            TunnResult::Done
+        ));
+        let sent_packet_buf = create_ipv4_udp_packet();
+        let data = my_tun.encapsulate(&sent_packet_buf, &mut my_dst);
+        assert!(matches!(data, TunnResult::WriteToNetwork(_)));
+
+        // Expire the session
+        mock_instant::MockClock::advance((REJECT_AFTER_TIME * 3).into());
+
+        assert!(matches!(
+            my_tun.update_timers(&mut []),
+            TunnResult::Err(WireGuardError::ConnectionExpired)
+        ));
+    }
+
+    #[cfg(feature = "mock-instant")]
+    #[test]
+    fn todo_assert_handshake_if_keepalive_after_3x_wg_reject_after_time() {
+        // WireGuard#6.3:
+        // <...>
+        // If no new secure session is created after (Reject-After-Time × 3) seconds,
+        // the current secure session, the previous secure session, and potentially
+        // the next secure session are discarded and zeroed out, in addition to any
+        // possible partially-completed handshake states and ephemeral keys.
+        // <...>
+        //
+        // Though it doesn't mention explicitly but observing Linux kernel implementation
+        // shows that non-zero Persistent-Keepalive must trigger the handshake.
+
+        let (mut my_tun, _) = create_two_tuns();
+
+        // Persistent-Keepalive value doesn't matter, while it's non zero
+        my_tun.set_persistent_keepalive(25);
+
+        let mut my_dst = [0u8; 1024];
+
+        let res = my_tun.update_timers(&mut my_dst);
+        assert!(matches!(res, TunnResult::WriteToNetwork(_)));
+
+        // Expire the session
+        mock_instant::MockClock::advance((REJECT_AFTER_TIME * 3).into());
+
+        assert!(matches!(
+            my_tun.update_timers(&mut my_dst),
+            TunnResult::WriteToNetwork(_)
+        ));
+    }
     #[test]
     #[cfg(feature = "mock-instant")]
     fn new_handshake_after_two_mins() {
+        // TODO: this test asserts for incorrect behaviour according to the spec. Since there's
+        // no keepalive set, the initiator should not perform a handshake request as there's
+        // no data packet sent. This behaviour is not aligned with the spec(WireGuard#6.2)
         let (mut my_tun, mut their_tun) = create_two_tuns_and_handshake();
         let mut my_dst = [0u8; 1024];
 

neptun/src/noise/timers.rs

@@ -36,8 +36,8 @@ use _instant_boottime::Instant;
 // Some constants, represent time in seconds
 // https://www.wireguard.com/papers/wireguard.pdf#page=14
 pub(crate) const REKEY_AFTER_TIME: Duration = Duration::from_secs(120);
-const REJECT_AFTER_TIME: Duration = Duration::from_secs(180);
-const REKEY_ATTEMPT_TIME: Duration = Duration::from_secs(90);
+pub(crate) const REJECT_AFTER_TIME: Duration = Duration::from_secs(180);
+pub(crate) const REKEY_ATTEMPT_TIME: Duration = Duration::from_secs(90);
 pub(crate) const REKEY_TIMEOUT: Duration = Duration::from_secs(5);
 const KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10);
 const COOKIE_EXPIRATION_TIME: Duration = Duration::from_secs(120);
@@ -233,9 +233,14 @@ impl Tunn {
             // (REJECT_AFTER_TIME * 3) ms if no new keys have been exchanged.
             if now - session_established >= REJECT_AFTER_TIME * 3 {
                 tracing::error!("CONNECTION_EXPIRED(REJECT_AFTER_TIME * 3)");
-                self.handshake.set_expired();
                 self.clear_all();
-                return TunnResult::Err(WireGuardError::ConnectionExpired);
+
+                if persistent_keepalive > 0 {
+                    handshake_initiation_required = true;
+                } else {
+                    self.handshake.set_expired();
+                    return TunnResult::Err(WireGuardError::ConnectionExpired);
+                }
             }
 
             if let Some(time_init_sent) = self.handshake.timer() {
@@ -245,10 +250,16 @@ impl Tunn {
                     // the retries give up and cease, and clear all existing packets queued
                     // up to be sent. If a packet is explicitly queued up to be sent, then
                     // this timer is reset.
+
                     tracing::error!("CONNECTION_EXPIRED(REKEY_ATTEMPT_TIME)");
-                    self.handshake.set_expired();
                     self.clear_all();
-                    return TunnResult::Err(WireGuardError::ConnectionExpired);
+
+                    if persistent_keepalive > 0 {
+                        handshake_initiation_required = true;
+                    } else {
+                        self.handshake.set_expired();
+                        return TunnResult::Err(WireGuardError::ConnectionExpired);
+                    }
                 }
 
                 if time_init_sent.elapsed() >= REKEY_TIMEOUT {
@@ -257,6 +268,7 @@ impl Tunn {
                     // A handshake initiation is retried after REKEY_TIMEOUT + jitter ms,
                     // if a response has not been received, where jitter is some random
                     // value between 0 and 333 ms.
+
                     tracing::warn!("HANDSHAKE(REKEY_TIMEOUT)");
                     handshake_initiation_required = true;
                 }
@@ -267,9 +279,22 @@ impl Tunn {
                     // ms old, we initiate a new handshake. If the sender was the original
                     // responder of the handshake, it does not re-initiate a new handshake
                     // after REKEY_AFTER_TIME ms like the original initiator does.
+
                     if session_established < data_packet_sent
                         && now - session_established >= REKEY_AFTER_TIME
                     {
+                        // TODO: When peer A initiates a handshake,
+                        // sends a packet, then sleeps past REKEY_AFTER_TIME, it should only reinitiate
+                        // if there's outgoing data. Currently, due to timestamp checks, a new session
+                        // is always attempted. If no keepalive is needed, it should stay silent.
+                        // Though as a side effect it makes the connectivity even better but
+                        // leaving this comment to be cautious about it.
+                        //
+                        // WireGuard#6.2:
+                        // <...Likewise, if a peer is the initiator of a current secure session,
+                        // WireGuard will send a handshake initiation message to begin a new secure session
+                        // if, after transmitting a transport data message, the current secure session
+                        // is Rekey-After-Time seconds old...>
                         tracing::debug!("HANDSHAKE(REKEY_AFTER_TIME (on send))");
                         handshake_initiation_required = true;
                     }