Refactor to add `ConnectingState::disconnect`

1 files changed, 31 insertions, 53 deletions

diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 720e2d3729..2aec258503 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -162,6 +162,17 @@ impl ConnectingState {
         }
     }
 
+    fn disconnect(
+        self,
+        shared_values: &mut SharedTunnelStateValues,
+        after_disconnect: AfterDisconnect,
+    ) -> EventConsequence<Self> {
+        EventConsequence::NewState(DisconnectingState::enter(
+            shared_values,
+            (self.close_handle, self.tunnel_close_event, after_disconnect),
+        ))
+    }
+
     fn handle_commands(
         self,
         commands: &mut mpsc::UnboundedReceiver<TunnelCommand>,
@@ -182,14 +193,10 @@ impl ConnectingState {
                             )
                         );
 
-                        NewState(DisconnectingState::enter(
+                        self.disconnect(
                             shared_values,
-                            (
-                                self.close_handle,
-                                self.tunnel_close_event,
-                                AfterDisconnect::Block(ErrorStateCause::SetFirewallPolicyError),
-                            ),
-                        ))
+                            AfterDisconnect::Block(ErrorStateCause::SetFirewallPolicyError),
+                        )
                     }
                 }
             }
@@ -200,42 +207,23 @@ impl ConnectingState {
             Ok(TunnelCommand::IsOffline(is_offline)) => {
                 shared_values.is_offline = is_offline;
                 if is_offline {
-                    NewState(DisconnectingState::enter(
+                    self.disconnect(
                         shared_values,
-                        (
-                            self.close_handle,
-                            self.tunnel_close_event,
-                            AfterDisconnect::Block(ErrorStateCause::IsOffline),
-                        ),
-                    ))
+                        AfterDisconnect::Block(ErrorStateCause::IsOffline),
+                    )
                 } 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,
-                (
-                    self.close_handle,
-                    self.tunnel_close_event,
-                    AfterDisconnect::Nothing,
-                ),
-            )),
-            Ok(TunnelCommand::Block(reason)) => NewState(DisconnectingState::enter(
-                shared_values,
-                (
-                    self.close_handle,
-                    self.tunnel_close_event,
-                    AfterDisconnect::Block(reason),
-                ),
-            )),
+            Ok(TunnelCommand::Connect) => {
+                self.disconnect(shared_values, AfterDisconnect::Reconnect(0))
+            }
+            Ok(TunnelCommand::Disconnect) | Err(_) => {
+                self.disconnect(shared_values, AfterDisconnect::Nothing)
+            }
+            Ok(TunnelCommand::Block(reason)) => {
+                self.disconnect(shared_values, AfterDisconnect::Block(reason))
+            }
         }
     }
 
@@ -247,14 +235,10 @@ impl ConnectingState {
         use self::EventConsequence::*;
 
         match try_handle_event!(self, self.tunnel_events.poll()) {
-            Ok(TunnelEvent::AuthFailed(reason)) => NewState(DisconnectingState::enter(
+            Ok(TunnelEvent::AuthFailed(reason)) => self.disconnect(
                 shared_values,
-                (
-                    self.close_handle,
-                    self.tunnel_close_event,
-                    AfterDisconnect::Block(ErrorStateCause::AuthFailed(reason)),
-                ),
-            )),
+                AfterDisconnect::Block(ErrorStateCause::AuthFailed(reason)),
+            ),
             Ok(TunnelEvent::Up(metadata)) => NewState(ConnectedState::enter(
                 shared_values,
                 self.into_connected_state_bootstrap(metadata),
@@ -262,14 +246,8 @@ impl ConnectingState {
             Ok(_) => SameState(self),
             Err(_) => {
                 debug!("The tunnel disconnected unexpectedly");
-                NewState(DisconnectingState::enter(
-                    shared_values,
-                    (
-                        self.close_handle,
-                        self.tunnel_close_event,
-                        AfterDisconnect::Reconnect(self.retry_attempt + 1),
-                    ),
-                ))
+                let retry_attempt = self.retry_attempt + 1;
+                self.disconnect(shared_values, AfterDisconnect::Reconnect(retry_attempt))
             }
         }
     }