Gradually increase retry interval for wireguard tunnels

4 files changed, 37 insertions, 9 deletions

diff --git a/talpid-core/src/tunnel/mod.rs b/talpid-core/src/tunnel/mod.rs
index ea04a9d425..0ed39d7bee 100644
--- a/talpid-core/src/tunnel/mod.rs
+++ b/talpid-core/src/tunnel/mod.rs
@@ -106,6 +106,7 @@ impl TunnelMonitor {
         on_event: L,
         tun_provider: &mut TunProvider,
         route_manager: &mut RouteManager,
+        retry_attempt: u32,
     ) -> Result<Self>
     where
         L: (Fn(TunnelEvent) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>)
@@ -133,6 +134,7 @@ impl TunnelMonitor {
                 on_event,
                 tun_provider,
                 route_manager,
+                retry_attempt,
             ),
         }
     }
@@ -165,6 +167,7 @@ impl TunnelMonitor {
         on_event: L,
         tun_provider: &mut TunProvider,
         route_manager: &mut RouteManager,
+        retry_attempt: u32,
     ) -> Result<Self>
     where
         L: (Fn(TunnelEvent) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>)
@@ -182,6 +185,7 @@ impl TunnelMonitor {
             on_event,
             tun_provider,
             route_manager,
+            retry_attempt,
         )?;
         Ok(TunnelMonitor {
             monitor: InternalTunnelMonitor::Wireguard(monitor),
diff --git a/talpid-core/src/tunnel/wireguard/connectivity_check.rs b/talpid-core/src/tunnel/wireguard/connectivity_check.rs
index 5a2118fb27..ab91896ca0 100644
--- a/talpid-core/src/tunnel/wireguard/connectivity_check.rs
+++ b/talpid-core/src/tunnel/wireguard/connectivity_check.rs
@@ -3,6 +3,7 @@ use crate::{
     tunnel::wireguard::stats::StatsMap,
 };
 use std::{
+    cmp,
     net::Ipv4Addr,
     sync::{mpsc, Mutex, Weak},
     time::{Duration, Instant},
@@ -25,6 +26,12 @@ const TRAFFIC_TIMEOUT: Duration = Duration::from_secs(120);
 /// Timeout for waiting on receiving traffic after sending the first ICMP packet.  Once this
 /// timeout is reached, it is assumed that the connection is lost.
 const PING_TIMEOUT: Duration = Duration::from_secs(15);
+/// Timeout for receiving traffic when establishing a connection.
+const ESTABLISH_TIMEOUT: Duration = Duration::from_secs(4);
+/// `ESTABLISH_TIMEOUT` is multiplied by this after each failed connection attempt.
+const ESTABLISH_TIMEOUT_MULTIPLIER: u32 = 2;
+/// Maximum timeout for establishing a connection.
+const MAX_ESTABLISH_TIMEOUT: Duration = PING_TIMEOUT;
 /// Number of seconds to wait between sending ICMP packets
 const SECONDS_PER_PING: Duration = Duration::from_secs(3);
 
@@ -99,14 +106,20 @@ impl ConnectivityMonitor {
 
     // checks if the tunnel has ever worked. Intended to check if a connection to a tunnel is
     // successfull at the start of a connection.
-    pub(super) fn establish_connectivity(&mut self) -> Result<bool, Error> {
+    pub(super) fn establish_connectivity(&mut self, retry_attempt: u32) -> Result<bool, Error> {
         if self.conn_state.connected() {
             return Ok(true);
         }
 
+        let check_timeout = cmp::min(
+            MAX_ESTABLISH_TIMEOUT,
+            ESTABLISH_TIMEOUT
+                .saturating_mul(ESTABLISH_TIMEOUT_MULTIPLIER.saturating_pow(retry_attempt)),
+        );
+
         let start = Instant::now();
-        while start.elapsed() < PING_TIMEOUT {
-            if self.check_connectivity(Instant::now())? {
+        while start.elapsed() < check_timeout {
+            if self.check_connectivity_interval(Instant::now(), check_timeout)? {
                 return Ok(true);
             }
             if self.should_shut_down(DELAY_ON_INITIAL_SETUP) {
@@ -155,6 +168,15 @@ impl ConnectivityMonitor {
 
     /// Returns true if connection is established
     fn check_connectivity(&mut self, now: Instant) -> Result<bool, Error> {
+        self.check_connectivity_interval(now, PING_TIMEOUT)
+    }
+
+    /// Returns true if connection is established
+    fn check_connectivity_interval(
+        &mut self,
+        now: Instant,
+        timeout: Duration,
+    ) -> Result<bool, Error> {
         match self.get_stats() {
             None => Ok(false),
             Some(new_stats) => {
@@ -166,7 +188,7 @@ impl ConnectivityMonitor {
                 }
 
                 self.maybe_send_ping(now)?;
-                Ok(!self.ping_timed_out() && self.conn_state.connected())
+                Ok(!self.ping_timed_out(timeout) && self.conn_state.connected())
             }
         }
     }
@@ -203,9 +225,9 @@ impl ConnectivityMonitor {
         Ok(())
     }
 
-    fn ping_timed_out(&self) -> bool {
+    fn ping_timed_out(&self, timeout: Duration) -> bool {
         self.initial_ping_timestamp
-            .map(|initial_ping_timestamp| initial_ping_timestamp.elapsed() > PING_TIMEOUT)
+            .map(|initial_ping_timestamp| initial_ping_timestamp.elapsed() > timeout)
             .unwrap_or(false)
     }
 
@@ -662,7 +684,7 @@ mod test {
             let start = now - Duration::from_secs(1);
             let mut monitor = mock_monitor(start, Box::new(pinger), tunnel, stop_rx);
 
-            let start_result = monitor.establish_connectivity();
+            let start_result = monitor.establish_connectivity(0);
             result_tx.send(start_result).unwrap();
 
             let result = monitor.run().map(|_| true);
@@ -715,7 +737,7 @@ mod test {
             let now = Instant::now();
             let start = now - Duration::from_secs(1);
             let mut monitor = mock_monitor(start, Box::new(pinger), tunnel, stop_rx);
-            let start_result = monitor.establish_connectivity();
+            let start_result = monitor.establish_connectivity(0);
             result_tx.send(start_result).unwrap();
             let end_result = monitor.run().map(|_| true);
             result_tx.send(end_result).expect("Failed to send result");
diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 848c309da4..0e14ea9269 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -168,6 +168,7 @@ impl WireguardMonitor {
         on_event: F,
         tun_provider: &mut TunProvider,
         route_manager: &mut routing::RouteManager,
+        retry_attempt: u32,
     ) -> Result<WireguardMonitor> {
         let mut tcp_proxies = vec![];
         let mut endpoint_addrs = vec![];
@@ -287,7 +288,7 @@ impl WireguardMonitor {
                 return;
             }
 
-            match connectivity_monitor.establish_connectivity() {
+            match connectivity_monitor.establish_connectivity(retry_attempt) {
                 Ok(true) => {
                     runtime.block_on((on_event)(TunnelEvent::Up(metadata)));
 
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 28a80260dc..8b4c9f5855 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -118,6 +118,7 @@ impl ConnectingState {
             on_tunnel_event,
             tun_provider,
             route_manager,
+            retry_attempt,
         )?;
         let close_handle = Some(monitor.close_handle());
         let tunnel_close_event =