Make `tunnel_close_event` optional

3 files changed, 25 insertions, 10 deletions

diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index 2e1e613460..fd26b20bc3 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -20,7 +20,7 @@ pub struct ConnectedStateBootstrap {
     pub metadata: TunnelMetadata,
     pub tunnel_events: mpsc::UnboundedReceiver<TunnelEvent>,
     pub tunnel_parameters: TunnelParameters,
-    pub tunnel_close_event: oneshot::Receiver<Option<BlockReason>>,
+    pub tunnel_close_event: Option<oneshot::Receiver<Option<BlockReason>>>,
     pub close_handle: Option<CloseHandle>,
 }
 
@@ -29,7 +29,7 @@ pub struct ConnectedState {
     metadata: TunnelMetadata,
     tunnel_events: mpsc::UnboundedReceiver<TunnelEvent>,
     tunnel_parameters: TunnelParameters,
-    tunnel_close_event: oneshot::Receiver<Option<BlockReason>>,
+    tunnel_close_event: Option<oneshot::Receiver<Option<BlockReason>>>,
     close_handle: Option<CloseHandle>,
 }
 
@@ -175,7 +175,12 @@ impl ConnectedState {
     ) -> EventConsequence<Self> {
         use self::EventConsequence::*;
 
-        match self.tunnel_close_event.poll() {
+        let poll_result = match &mut self.tunnel_close_event {
+            Some(tunnel_close_event) => tunnel_close_event.poll(),
+            None => Ok(Async::NotReady),
+        };
+
+        match poll_result {
             Ok(Async::Ready(block_reason)) => {
                 if let Some(reason) = block_reason {
                     return NewState(BlockedState::enter(shared_values, reason));
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index a8969cb122..8c4e7578dd 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -34,7 +34,7 @@ const MIN_TUNNEL_ALIVE_TIME: Duration = Duration::from_millis(1000);
 pub struct ConnectingState {
     tunnel_events: mpsc::UnboundedReceiver<TunnelEvent>,
     tunnel_parameters: TunnelParameters,
-    tunnel_close_event: oneshot::Receiver<Option<BlockReason>>,
+    tunnel_close_event: Option<oneshot::Receiver<Option<BlockReason>>>,
     close_handle: Option<CloseHandle>,
     retry_attempt: u32,
 }
@@ -92,7 +92,7 @@ impl ConnectingState {
 
     fn spawn_tunnel_monitor_wait_thread(
         tunnel_monitor: TunnelMonitor,
-    ) -> oneshot::Receiver<Option<BlockReason>> {
+    ) -> Option<oneshot::Receiver<Option<BlockReason>>> {
         let (tunnel_close_event_tx, tunnel_close_event_rx) = oneshot::channel();
 
         thread::spawn(move || {
@@ -113,7 +113,7 @@ impl ConnectingState {
             trace!("Tunnel monitor thread exit");
         });
 
-        tunnel_close_event_rx
+        Some(tunnel_close_event_rx)
     }
 
     fn wait_for_tunnel_monitor(tunnel_monitor: TunnelMonitor) -> Option<BlockReason> {
@@ -270,7 +270,12 @@ impl ConnectingState {
         mut self,
         shared_values: &mut SharedTunnelStateValues,
     ) -> EventConsequence<Self> {
-        match self.tunnel_close_event.poll() {
+        let poll_result = match &mut self.tunnel_close_event {
+            Some(tunnel_close_event) => tunnel_close_event.poll(),
+            None => Ok(Async::NotReady),
+        };
+
+        match poll_result {
             Ok(Async::Ready(block_reason)) => {
                 if let Some(reason) = block_reason {
                     return EventConsequence::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 5361bebf36..733a6e2448 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -16,7 +16,7 @@ use talpid_types::{
 /// This state is active from when we manually trigger a tunnel kill until the tunnel wait
 /// operation (TunnelExit) returned.
 pub struct DisconnectingState {
-    exited: oneshot::Receiver<Option<BlockReason>>,
+    exited: Option<oneshot::Receiver<Option<BlockReason>>>,
     after_disconnect: AfterDisconnect,
 }
 
@@ -101,7 +101,12 @@ impl DisconnectingState {
     ) -> EventConsequence<Self> {
         use self::EventConsequence::*;
 
-        match self.exited.poll() {
+        let poll_result = match &mut self.exited {
+            Some(exited) => exited.poll(),
+            None => Ok(Async::Ready(None)),
+        };
+
+        match poll_result {
             Ok(Async::NotReady) => NoEvents(self),
             Ok(Async::Ready(block_reason)) => {
                 NewState(self.after_disconnect(block_reason, shared_values))
@@ -132,7 +137,7 @@ impl DisconnectingState {
 impl TunnelState for DisconnectingState {
     type Bootstrap = (
         Option<CloseHandle>,
-        oneshot::Receiver<Option<BlockReason>>,
+        Option<oneshot::Receiver<Option<BlockReason>>>,
         AfterDisconnect,
     );