Include after disconnect action in new state event

1 files changed, 15 insertions, 2 deletions

diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index a7f20ca5ea..47467e3090 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -4,7 +4,7 @@ use error_chain::ChainedError;
 use futures::sync::{mpsc, oneshot};
 use futures::{Async, Future, Stream};
 
-use talpid_types::tunnel::BlockReason;
+use talpid_types::tunnel::{ActionAfterDisconnect, BlockReason};
 
 use super::{
     BlockedState, ConnectingState, DisconnectedState, EventConsequence, ResultExt,
@@ -103,12 +103,14 @@ impl TunnelState for DisconnectingState {
             }
         });
 
+        let action_after_disconnect = after_disconnect.action();
+
         (
             TunnelStateWrapper::from(DisconnectingState {
                 exited,
                 after_disconnect,
             }),
-            TunnelStateTransition::Disconnecting,
+            TunnelStateTransition::Disconnecting(action_after_disconnect),
         )
     }
 
@@ -128,3 +130,14 @@ pub enum AfterDisconnect {
     Block(BlockReason, bool),
     Reconnect(TunnelParameters),
 }
+
+impl AfterDisconnect {
+    /// Build event representation of the action that will be taken after the disconnection.
+    pub fn action(&self) -> ActionAfterDisconnect {
+        match self {
+            AfterDisconnect::Nothing => ActionAfterDisconnect::Nothing,
+            AfterDisconnect::Block(..) => ActionAfterDisconnect::Block,
+            AfterDisconnect::Reconnect(_) => ActionAfterDisconnect::Reconnect,
+        }
+    }
+}