Merge branch 'reconnect-to-different-relay'

6 files changed, 108 insertions, 93 deletions

diff --git a/talpid-core/src/tunnel_state_machine/blocked_state.rs b/talpid-core/src/tunnel_state_machine/blocked_state.rs
index 33cdd3a1e1..8c73ccd0c5 100644
--- a/talpid-core/src/tunnel_state_machine/blocked_state.rs
+++ b/talpid-core/src/tunnel_state_machine/blocked_state.rs
@@ -55,9 +55,7 @@ impl TunnelState for BlockedState {
                 Self::set_security_policy(shared_values);
                 SameState(self)
             }
-            Ok(TunnelCommand::Connect(parameters)) => {
-                NewState(ConnectingState::enter(shared_values, parameters))
-            }
+            Ok(TunnelCommand::Connect) => NewState(ConnectingState::enter(shared_values, 0)),
             Ok(TunnelCommand::Disconnect) | Err(_) => {
                 NewState(DisconnectedState::enter(shared_values, ()))
             }
diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index 7d04bf0c3b..a969ece39f 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -79,20 +79,14 @@ impl ConnectedState {
                     }
                 }
             }
-            Ok(TunnelCommand::Connect(parameters)) => {
-                if parameters != self.tunnel_parameters {
-                    NewState(DisconnectingState::enter(
-                        shared_values,
-                        (
-                            self.close_handle,
-                            self.tunnel_close_event,
-                            AfterDisconnect::Reconnect(parameters),
-                        ),
-                    ))
-                } else {
-                    SameState(self)
-                }
-            }
+            Ok(TunnelCommand::Connect) => NewState(DisconnectingState::enter(
+                shared_values,
+                (
+                    self.close_handle,
+                    self.tunnel_close_event,
+                    AfterDisconnect::Reconnect(0),
+                ),
+            )),
             Ok(TunnelCommand::Disconnect) | Err(_) => NewState(DisconnectingState::enter(
                 shared_values,
                 (
@@ -124,7 +118,7 @@ impl ConnectedState {
                 (
                     self.close_handle,
                     self.tunnel_close_event,
-                    AfterDisconnect::Reconnect(self.tunnel_parameters),
+                    AfterDisconnect::Reconnect(0),
                 ),
             )),
             Ok(_) => SameState(self),
@@ -144,10 +138,7 @@ impl ConnectedState {
         }
 
         info!("Tunnel closed. Reconnecting.");
-        NewState(ConnectingState::enter(
-            shared_values,
-            self.tunnel_parameters,
-        ))
+        NewState(ConnectingState::enter(shared_values, 0))
     }
 }
 
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index bcf1c422e4..1e0880141e 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -47,6 +47,7 @@ pub struct ConnectingState {
     tunnel_parameters: TunnelParameters,
     tunnel_close_event: oneshot::Receiver<()>,
     close_handle: CloseHandle,
+    retry_attempt: u32,
 }
 
 impl ConnectingState {
@@ -68,6 +69,7 @@ impl ConnectingState {
         parameters: TunnelParameters,
         log_dir: &Option<PathBuf>,
         resource_dir: &Path,
+        retry_attempt: u32,
     ) -> Result<Self> {
         let (event_tx, event_rx) = mpsc::unbounded();
         let monitor = Self::spawn_tunnel_monitor(&parameters, log_dir, resource_dir, event_tx)?;
@@ -79,6 +81,7 @@ impl ConnectingState {
             tunnel_parameters: parameters,
             tunnel_close_event,
             close_handle,
+            retry_attempt,
         })
     }
 
@@ -185,20 +188,14 @@ impl ConnectingState {
                     }
                 }
             }
-            Ok(TunnelCommand::Connect(parameters)) => {
-                if parameters != self.tunnel_parameters {
-                    NewState(DisconnectingState::enter(
-                        shared_values,
-                        (
-                            self.close_handle,
-                            self.tunnel_close_event,
-                            AfterDisconnect::Reconnect(parameters),
-                        ),
-                    ))
-                } else {
-                    SameState(self)
-                }
-            }
+            Ok(TunnelCommand::Connect) => NewState(DisconnectingState::enter(
+                shared_values,
+                (
+                    self.close_handle,
+                    self.tunnel_close_event,
+                    AfterDisconnect::Reconnect(0),
+                ),
+            )),
             Ok(TunnelCommand::Disconnect) | Err(_) => NewState(DisconnectingState::enter(
                 shared_values,
                 (
@@ -245,7 +242,7 @@ impl ConnectingState {
                     (
                         self.close_handle,
                         self.tunnel_close_event,
-                        AfterDisconnect::Reconnect(self.tunnel_parameters),
+                        AfterDisconnect::Reconnect(self.retry_attempt + 1),
                     ),
                 ))
             }
@@ -262,47 +259,61 @@ impl ConnectingState {
             Err(_cancelled) => warn!("Tunnel monitor thread has stopped unexpectedly"),
         }
 
-        info!("Tunnel closed. Reconnecting.");
+        info!(
+            "Tunnel closed. Reconnecting, attempt {}.",
+            self.retry_attempt + 1
+        );
         EventConsequence::NewState(ConnectingState::enter(
             shared_values,
-            self.tunnel_parameters,
+            self.retry_attempt + 1,
         ))
     }
 }
 
 impl TunnelState for ConnectingState {
-    type Bootstrap = TunnelParameters;
+    type Bootstrap = u32;
 
     fn enter(
         shared_values: &mut SharedTunnelStateValues,
-        parameters: Self::Bootstrap,
+        retry_attempt: u32,
     ) -> (TunnelStateWrapper, TunnelStateTransition) {
-        if let Err(error) = Self::set_security_policy(shared_values, parameters.endpoint) {
-            error!("{}", error.display_chain());
-            return BlockedState::enter(shared_values, BlockReason::StartTunnelError);
-        }
-
-        match Self::start_tunnel(
-            parameters,
-            &shared_values.log_dir,
-            &shared_values.resource_dir,
-        ) {
-            Ok(connecting_state) => (
-                TunnelStateWrapper::from(connecting_state),
-                TunnelStateTransition::Connecting,
-            ),
-            Err(error) => {
-                let block_reason = match *error.kind() {
-                    ErrorKind::TunnelMonitorError(tunnel::ErrorKind::EnableIpv6Error) => {
-                        BlockReason::Ipv6Unavailable
-                    }
-                    _ => BlockReason::StartTunnelError,
-                };
+        match shared_values
+            .tunnel_parameters_generator
+            .generate(retry_attempt)
+        {
+            None => BlockedState::enter(shared_values, BlockReason::NoMatchingRelay),
+            Some(tunnel_parameters) => {
+                if let Err(error) =
+                    Self::set_security_policy(shared_values, tunnel_parameters.endpoint)
+                {
+                    error!("{}", error.display_chain());
+                    BlockedState::enter(shared_values, BlockReason::StartTunnelError)
+                } else {
+                    match Self::start_tunnel(
+                        tunnel_parameters,
+                        &shared_values.log_dir,
+                        &shared_values.resource_dir,
+                        retry_attempt,
+                    ) {
+                        Ok(connecting_state) => (
+                            TunnelStateWrapper::from(connecting_state),
+                            TunnelStateTransition::Connecting,
+                        ),
+                        Err(error) => {
+                            let block_reason = match *error.kind() {
+                                ErrorKind::TunnelMonitorError(
+                                    tunnel::ErrorKind::EnableIpv6Error,
+                                ) => BlockReason::Ipv6Unavailable,
+                                _ => BlockReason::StartTunnelError,
+                            };
 
-                let chained_error = error.chain_err(|| "Failed to start tunnel");
-                error!("{}", chained_error.display_chain());
+                            let chained_error = error.chain_err(|| "Failed to start tunnel");
+                            error!("{}", chained_error.display_chain());
 
-                BlockedState::enter(shared_values, block_reason)
+                            BlockedState::enter(shared_values, block_reason)
+                        }
+                    }
+                }
             }
         }
     }
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index 3d4518a050..227e019323 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -46,9 +46,7 @@ impl TunnelState for DisconnectedState {
                 shared_values.allow_lan = allow_lan;
                 SameState(self)
             }
-            Ok(TunnelCommand::Connect(parameters)) => {
-                NewState(ConnectingState::enter(shared_values, parameters))
-            }
+            Ok(TunnelCommand::Connect) => NewState(ConnectingState::enter(shared_values, 0)),
             Ok(TunnelCommand::Block(reason)) => {
                 NewState(BlockedState::enter(shared_values, reason))
             }
diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index 5028792c11..864b11ab4b 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -8,8 +8,7 @@ use talpid_types::tunnel::{ActionAfterDisconnect, BlockReason};
 
 use super::{
     BlockedState, ConnectingState, DisconnectedState, EventConsequence, ResultExt,
-    SharedTunnelStateValues, TunnelCommand, TunnelParameters, TunnelState, TunnelStateTransition,
-    TunnelStateWrapper,
+    SharedTunnelStateValues, TunnelCommand, TunnelState, TunnelStateTransition, TunnelStateWrapper,
 };
 use tunnel::CloseHandle;
 
@@ -26,8 +25,6 @@ impl DisconnectingState {
         commands: &mut mpsc::UnboundedReceiver<TunnelCommand>,
         shared_values: &mut SharedTunnelStateValues,
     ) -> EventConsequence<Self> {
-        use self::AfterDisconnect::*;
-
         let event = try_handle_event!(self, commands.poll());
         let after_disconnect = self.after_disconnect;
 
@@ -35,30 +32,30 @@ impl DisconnectingState {
             AfterDisconnect::Nothing => match event {
                 Ok(TunnelCommand::AllowLan(allow_lan)) => {
                     shared_values.allow_lan = allow_lan;
-                    Nothing
+                    AfterDisconnect::Nothing
                 }
-                Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
-                Ok(TunnelCommand::Block(reason)) => Block(reason),
-                _ => Nothing,
+                Ok(TunnelCommand::Connect) => AfterDisconnect::Reconnect(0),
+                Ok(TunnelCommand::Block(reason)) => AfterDisconnect::Block(reason),
+                _ => AfterDisconnect::Nothing,
             },
             AfterDisconnect::Block(reason) => match event {
                 Ok(TunnelCommand::AllowLan(allow_lan)) => {
                     shared_values.allow_lan = allow_lan;
-                    Block(reason)
+                    AfterDisconnect::Block(reason)
                 }
-                Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
-                Ok(TunnelCommand::Disconnect) => Nothing,
-                Ok(TunnelCommand::Block(new_reason)) => Block(new_reason),
-                Err(_) => Block(reason),
+                Ok(TunnelCommand::Connect) => AfterDisconnect::Reconnect(0),
+                Ok(TunnelCommand::Disconnect) => AfterDisconnect::Nothing,
+                Ok(TunnelCommand::Block(new_reason)) => AfterDisconnect::Block(new_reason),
+                Err(_) => AfterDisconnect::Block(reason),
             },
-            AfterDisconnect::Reconnect(tunnel_parameters) => match event {
+            AfterDisconnect::Reconnect(retry_attempt) => match event {
                 Ok(TunnelCommand::AllowLan(allow_lan)) => {
                     shared_values.allow_lan = allow_lan;
-                    Reconnect(tunnel_parameters)
+                    AfterDisconnect::Reconnect(retry_attempt)
                 }
-                Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
-                Ok(TunnelCommand::Disconnect) | Err(_) => Nothing,
-                Ok(TunnelCommand::Block(reason)) => Block(reason),
+                Ok(TunnelCommand::Connect) => AfterDisconnect::Reconnect(retry_attempt),
+                Ok(TunnelCommand::Disconnect) | Err(_) => AfterDisconnect::Nothing,
+                Ok(TunnelCommand::Block(reason)) => AfterDisconnect::Block(reason),
             },
         };
 
@@ -84,8 +81,8 @@ impl DisconnectingState {
         match self.after_disconnect {
             AfterDisconnect::Nothing => DisconnectedState::enter(shared_values, ()),
             AfterDisconnect::Block(reason) => BlockedState::enter(shared_values, reason),
-            AfterDisconnect::Reconnect(tunnel_parameters) => {
-                ConnectingState::enter(shared_values, tunnel_parameters)
+            AfterDisconnect::Reconnect(retry_attempt) => {
+                ConnectingState::enter(shared_values, retry_attempt)
             }
         }
     }
@@ -133,7 +130,7 @@ impl TunnelState for DisconnectingState {
 pub enum AfterDisconnect {
     Nothing,
     Block(BlockReason),
-    Reconnect(TunnelParameters),
+    Reconnect(u32),
 }
 
 impl AfterDisconnect {
@@ -142,7 +139,7 @@ impl AfterDisconnect {
         match self {
             AfterDisconnect::Nothing => ActionAfterDisconnect::Nothing,
             AfterDisconnect::Block(..) => ActionAfterDisconnect::Block,
-            AfterDisconnect::Reconnect(_) => ActionAfterDisconnect::Reconnect,
+            AfterDisconnect::Reconnect(..) => ActionAfterDisconnect::Reconnect,
         }
     }
 }
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 91ca9aa7ad..66e86de9de 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -44,6 +44,7 @@ error_chain! {
 /// Spawn the tunnel state machine thread, returning a channel for sending tunnel commands.
 pub fn spawn<P, T>(
     allow_lan: bool,
+    tunnel_parameters_generator: impl TunnelParametersGenerator,
     log_dir: Option<PathBuf>,
     resource_dir: PathBuf,
     cache_dir: P,
@@ -59,6 +60,7 @@ where
     thread::spawn(move || {
         match create_event_loop(
             allow_lan,
+            tunnel_parameters_generator,
             log_dir,
             resource_dir,
             cache_dir,
@@ -92,6 +94,7 @@ where
 
 fn create_event_loop<T>(
     allow_lan: bool,
+    tunnel_parameters_generator: impl TunnelParametersGenerator,
     log_dir: Option<PathBuf>,
     resource_dir: PathBuf,
     cache_dir: impl AsRef<Path>,
@@ -102,8 +105,14 @@ where
     T: From<TunnelStateTransition> + Send + 'static,
 {
     let reactor = Core::new().chain_err(|| ErrorKind::ReactorError)?;
-    let state_machine =
-        TunnelStateMachine::new(allow_lan, log_dir, resource_dir, cache_dir, commands)?;
+    let state_machine = TunnelStateMachine::new(
+        allow_lan,
+        tunnel_parameters_generator,
+        log_dir,
+        resource_dir,
+        cache_dir,
+        commands,
+    )?;
 
     let future = state_machine.for_each(move |state_change_event| {
         state_change_listener
@@ -119,7 +128,7 @@ pub enum TunnelCommand {
     /// Enable or disable LAN access in the firewall.
     AllowLan(bool),
     /// Open tunnel connection.
-    Connect(TunnelParameters),
+    Connect,
     /// Close tunnel connection.
     Disconnect,
     /// Disconnect any open tunnel and block all network access
@@ -152,6 +161,7 @@ struct TunnelStateMachine {
 impl TunnelStateMachine {
     fn new(
         allow_lan: bool,
+        tunnel_parameters_generator: impl TunnelParametersGenerator,
         log_dir: Option<PathBuf>,
         resource_dir: PathBuf,
         cache_dir: impl AsRef<Path>,
@@ -162,6 +172,7 @@ impl TunnelStateMachine {
         let mut shared_values = SharedTunnelStateValues {
             security,
             allow_lan,
+            tunnel_parameters_generator: Box::new(tunnel_parameters_generator),
             log_dir,
             resource_dir,
         };
@@ -225,12 +236,21 @@ impl<T: TunnelState> From<EventConsequence<T>> for TunnelStateMachineAction {
     }
 }
 
+/// Trait for any type that can provide a stream of `TunnelParameters` to the `TunnelStateMachine`.
+pub trait TunnelParametersGenerator: Send + 'static {
+    /// Given the number of consecutive failed retry attempts, it should yield a `TunnelParameters`
+    /// to establish a tunnel with.
+    /// If this returns `None` then the state machine goes into the `Blocked` state.
+    fn generate(&mut self, retry_attempt: u32) -> Option<TunnelParameters>;
+}
 
 /// Values that are common to all tunnel states.
 struct SharedTunnelStateValues {
     security: NetworkSecurity,
     /// Should LAN access be allowed outside the tunnel.
     allow_lan: bool,
+    /// The generator of new `TunnelParameter`s
+    tunnel_parameters_generator: Box<dyn TunnelParametersGenerator>,
     /// Directory to store tunnel log file.
     log_dir: Option<PathBuf>,
     /// Resource directory path.