Include state transition in `NewState`

5 files changed, 38 insertions, 31 deletions

diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index 2881e0409c..d207e9a9fb 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -5,7 +5,8 @@ use talpid_types::net::TunnelEndpoint;
 
 use super::{
     AfterDisconnect, ConnectingState, DisconnectingState, EventConsequence, Result, ResultExt,
-    SharedTunnelStateValues, TunnelCommand, TunnelParameters, TunnelState, TunnelStateWrapper,
+    SharedTunnelStateValues, TunnelCommand, TunnelParameters, TunnelState, TunnelStateTransition,
+    TunnelStateWrapper,
 };
 use security::{NetworkSecurity, SecurityPolicy};
 use tunnel::{CloseHandle, TunnelEvent, TunnelMetadata};
@@ -151,11 +152,14 @@ impl TunnelState for ConnectedState {
     fn enter(
         shared_values: &mut SharedTunnelStateValues,
         bootstrap: Self::Bootstrap,
-    ) -> TunnelStateWrapper {
+    ) -> (TunnelStateWrapper, TunnelStateTransition) {
         let connected_state = ConnectedState::from(bootstrap);
 
         match connected_state.set_security_policy(shared_values) {
-            Ok(()) => TunnelStateWrapper::from(connected_state),
+            Ok(()) => (
+                TunnelStateWrapper::from(connected_state),
+                TunnelStateTransition::Connected,
+            ),
             Err(error) => {
                 error!("{}", error.chain_err(|| "Failed to set security policy"));
                 DisconnectingState::enter(
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 15e851e351..d1c16bac1d 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -13,7 +13,7 @@ use talpid_types::net::{TunnelEndpoint, TunnelEndpointData};
 use super::{
     AfterDisconnect, ConnectedState, ConnectedStateBootstrap, DisconnectedState,
     DisconnectingState, EventConsequence, Result, ResultExt, SharedTunnelStateValues,
-    TunnelCommand, TunnelParameters, TunnelState, TunnelStateWrapper,
+    TunnelCommand, TunnelParameters, TunnelState, TunnelStateTransition, TunnelStateWrapper,
 };
 use logging;
 use security::{NetworkSecurity, SecurityPolicy};
@@ -269,9 +269,12 @@ impl TunnelState for ConnectingState {
     fn enter(
         shared_values: &mut SharedTunnelStateValues,
         parameters: Self::Bootstrap,
-    ) -> TunnelStateWrapper {
+    ) -> (TunnelStateWrapper, TunnelStateTransition) {
         match Self::new(shared_values, parameters) {
-            Ok(connecting_state) => TunnelStateWrapper::from(connecting_state),
+            Ok(connecting_state) => (
+                TunnelStateWrapper::from(connecting_state),
+                TunnelStateTransition::Connecting,
+            ),
             Err(error) => {
                 error!("{}", error.chain_err(|| "Failed to start tunnel"));
                 DisconnectedState::enter(shared_values, ())
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index 7662f941c6..f92630606c 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -4,7 +4,7 @@ use futures::Stream;
 
 use super::{
     ConnectingState, Error, EventConsequence, SharedTunnelStateValues, TunnelCommand, TunnelState,
-    TunnelStateWrapper,
+    TunnelStateTransition, TunnelStateWrapper,
 };
 use security::NetworkSecurity;
 
@@ -27,10 +27,13 @@ impl TunnelState for DisconnectedState {
     fn enter(
         shared_values: &mut SharedTunnelStateValues,
         _: Self::Bootstrap,
-    ) -> TunnelStateWrapper {
+    ) -> (TunnelStateWrapper, TunnelStateTransition) {
         Self::reset_security_policy(shared_values);
 
-        TunnelStateWrapper::from(DisconnectedState)
+        (
+            TunnelStateWrapper::from(DisconnectedState),
+            TunnelStateTransition::Disconnected,
+        )
     }
 
     fn handle_event(
diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index ffcbd51a65..0f65818257 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -6,7 +6,7 @@ use futures::{Async, Future, Stream};
 
 use super::{
     ConnectingState, DisconnectedState, EventConsequence, ResultExt, SharedTunnelStateValues,
-    TunnelCommand, TunnelParameters, TunnelState, TunnelStateWrapper,
+    TunnelCommand, TunnelParameters, TunnelState, TunnelStateTransition, TunnelStateWrapper,
 };
 use tunnel::CloseHandle;
 
@@ -57,7 +57,10 @@ impl DisconnectingState {
         }
     }
 
-    fn after_disconnect(self, shared_values: &mut SharedTunnelStateValues) -> TunnelStateWrapper {
+    fn after_disconnect(
+        self,
+        shared_values: &mut SharedTunnelStateValues,
+    ) -> (TunnelStateWrapper, TunnelStateTransition) {
         match self.after_disconnect {
             AfterDisconnect::Nothing => DisconnectedState::enter(shared_values, ()),
             AfterDisconnect::Reconnect(tunnel_parameters) => {
@@ -73,7 +76,7 @@ impl TunnelState for DisconnectingState {
     fn enter(
         _: &mut SharedTunnelStateValues,
         (close_handle, exited, after_disconnect): Self::Bootstrap,
-    ) -> TunnelStateWrapper {
+    ) -> (TunnelStateWrapper, TunnelStateTransition) {
         thread::spawn(move || {
             let close_result = close_handle
                 .close()
@@ -84,10 +87,13 @@ impl TunnelState for DisconnectingState {
             }
         });
 
-        TunnelStateWrapper::from(DisconnectingState {
-            exited,
-            after_disconnect,
-        })
+        (
+            TunnelStateWrapper::from(DisconnectingState {
+                exited,
+                after_disconnect,
+            }),
+            TunnelStateTransition::Disconnecting,
+        )
     }
 
     fn handle_event(
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 68dfa17082..700f082146 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -198,8 +198,7 @@ impl<T: TunnelState> From<EventConsequence<T>> for TunnelStateMachineAction {
         use self::TunnelStateMachineAction::*;
 
         match event_consequence {
-            NewState(state_wrapper) => {
-                let transition = state_wrapper.info();
+            NewState((state_wrapper, transition)) => {
                 Notify(Some(state_wrapper), Ok(Async::Ready(Some(transition))))
             }
             SameState(state) => Repeat(state.into()),
@@ -217,7 +216,7 @@ struct SharedTunnelStateValues {
 /// Asynchronous result of an attempt to progress a state.
 enum EventConsequence<T: TunnelState> {
     /// Transition to a new state.
-    NewState(TunnelStateWrapper),
+    NewState((TunnelStateWrapper, TunnelStateTransition)),
     /// An event was received, but it was ignored by the state so no transition is performed.
     SameState(T),
     /// No events were received, the event loop should block until one becomes available.
@@ -259,7 +258,7 @@ trait TunnelState: Into<TunnelStateWrapper> + Sized {
     fn enter(
         shared_values: &mut SharedTunnelStateValues,
         bootstrap: Self::Bootstrap,
-    ) -> TunnelStateWrapper;
+    ) -> (TunnelStateWrapper, TunnelStateTransition);
 
     /// Main state function.
     ///
@@ -294,7 +293,9 @@ impl TunnelStateWrapper {
         shared_values: &mut SharedTunnelStateValues,
         bootstrap: <DisconnectedState as TunnelState>::Bootstrap,
     ) -> TunnelStateWrapper {
-        DisconnectedState::enter(shared_values, bootstrap)
+        let (new_state, _transition) = DisconnectedState::enter(shared_values, bootstrap);
+
+        new_state
     }
 
     fn handle_event(
@@ -322,16 +323,6 @@ impl TunnelStateWrapper {
             Disconnecting,
         }
     }
-
-    /// Returns information describing the state.
-    fn info(&self) -> TunnelStateTransition {
-        match *self {
-            TunnelStateWrapper::Disconnected(_) => TunnelStateTransition::Disconnected,
-            TunnelStateWrapper::Connecting(_) => TunnelStateTransition::Connecting,
-            TunnelStateWrapper::Connected(_) => TunnelStateTransition::Connected,
-            TunnelStateWrapper::Disconnecting(_) => TunnelStateTransition::Disconnecting,
-        }
-    }
 }
 
 macro_rules! impl_from_for_tunnel_state {