Make BlockedState activate Blocked security policy on entry

8 files changed, 65 insertions, 34 deletions

diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index 2fb27c3048..b67658c2ab 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -604,8 +604,9 @@ impl Daemon {
     }
 
     fn connect_tunnel(&mut self) {
+        let allow_lan = self.settings.get_allow_lan();
         let command = match self.settings.get_account_token() {
-            None => TunnelCommand::Block(BlockReason::NoAccountToken),
+            None => TunnelCommand::Block(BlockReason::NoAccountToken, allow_lan),
             Some(account_token) => match self.settings.get_relay_settings() {
                 RelaySettings::CustomTunnelEndpoint(custom_relay) => custom_relay
                     .to_tunnel_endpoint()
@@ -622,7 +623,7 @@ impl Daemon {
             .map(|parameters| TunnelCommand::Connect(parameters))
             .unwrap_or_else(|error| {
                 error!("{}", error.display_chain());
-                TunnelCommand::Block(BlockReason::NoMatchingRelay)
+                TunnelCommand::Block(BlockReason::NoMatchingRelay, allow_lan)
             }),
         };
         self.send_tunnel_command(command);
diff --git a/talpid-core/src/security/windows/mod.rs b/talpid-core/src/security/windows/mod.rs
index 75f712ee50..16fd707030 100644
--- a/talpid-core/src/security/windows/mod.rs
+++ b/talpid-core/src/security/windows/mod.rs
@@ -277,9 +277,7 @@ mod winfw {
         ) -> ApplyConnectedResult;
 
         #[link_name(WinFw_ApplyPolicyBlocked)]
-        pub fn WinFw_ApplyPolicyBlocked(
-            settings: &WinFwSettings,
-        ) -> ApplyBlockedResult;
+        pub fn WinFw_ApplyPolicyBlocked(settings: &WinFwSettings) -> ApplyBlockedResult;
 
         #[link_name(WinFw_Reset)]
         pub fn WinFw_Reset() -> ResettingPolicyResult;
diff --git a/talpid-core/src/tunnel_state_machine/blocked_state.rs b/talpid-core/src/tunnel_state_machine/blocked_state.rs
index 50efb9e371..57203db275 100644
--- a/talpid-core/src/tunnel_state_machine/blocked_state.rs
+++ b/talpid-core/src/tunnel_state_machine/blocked_state.rs
@@ -1,23 +1,40 @@
+use error_chain::ChainedError;
 use futures::sync::mpsc;
 use futures::Stream;
 
 use talpid_types::tunnel::BlockReason;
 
 use super::{
-    ConnectingState, DisconnectedState, EventConsequence, SharedTunnelStateValues, TunnelCommand,
-    TunnelState, TunnelStateTransition, TunnelStateWrapper,
+    ConnectingState, DisconnectedState, EventConsequence, ResultExt, SharedTunnelStateValues,
+    TunnelCommand, TunnelState, TunnelStateTransition, TunnelStateWrapper,
 };
+use security::{NetworkSecurity, SecurityPolicy};
 
 /// No tunnel is running and all network connections are blocked.
 pub struct BlockedState;
 
+impl BlockedState {
+    fn set_security_policy(shared_values: &mut SharedTunnelStateValues, allow_lan: bool) {
+        let policy = SecurityPolicy::Blocked { allow_lan };
+        debug!("Setting security policy: {:?}", policy);
+        if let Err(error) = shared_values
+            .security
+            .apply_policy(policy)
+            .chain_err(|| "Failed to apply security policy for blocked state")
+        {
+            error!("{}", error.display_chain());
+        }
+    }
+}
+
 impl TunnelState for BlockedState {
-    type Bootstrap = BlockReason;
+    type Bootstrap = (BlockReason, bool);
 
     fn enter(
-        _: &mut SharedTunnelStateValues,
-        block_reason: Self::Bootstrap,
+        shared_values: &mut SharedTunnelStateValues,
+        (block_reason, allow_lan): Self::Bootstrap,
     ) -> (TunnelStateWrapper, TunnelStateTransition) {
+        Self::set_security_policy(shared_values, allow_lan);
         (
             TunnelStateWrapper::from(BlockedState),
             TunnelStateTransition::Blocked(block_reason),
@@ -32,16 +49,19 @@ impl TunnelState for BlockedState {
         use self::EventConsequence::*;
 
         match try_handle_event!(self, commands.poll()) {
+            Ok(TunnelCommand::AllowLan(allow_lan)) => {
+                Self::set_security_policy(shared_values, allow_lan);
+                SameState(self)
+            }
             Ok(TunnelCommand::Connect(parameters)) => {
                 NewState(ConnectingState::enter(shared_values, parameters))
             }
             Ok(TunnelCommand::Disconnect) | Err(_) => {
                 NewState(DisconnectedState::enter(shared_values, ()))
             }
-            Ok(TunnelCommand::Block(reason)) => {
-                NewState(BlockedState::enter(shared_values, reason))
+            Ok(TunnelCommand::Block(reason, allow_lan)) => {
+                NewState(BlockedState::enter(shared_values, (reason, allow_lan)))
             }
-            _ => SameState(self),
         }
     }
 }
diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index 0117958503..b9856dbedc 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -79,7 +79,10 @@ impl ConnectedState {
                             (
                                 self.close_handle,
                                 self.tunnel_close_event,
-                                AfterDisconnect::Block(BlockReason::SetSecurityPolicyError),
+                                AfterDisconnect::Block(
+                                    BlockReason::SetSecurityPolicyError,
+                                    allow_lan,
+                                ),
                             ),
                         ))
                     }
@@ -107,12 +110,12 @@ impl ConnectedState {
                     AfterDisconnect::Nothing,
                 ),
             )),
-            Ok(TunnelCommand::Block(reason)) => NewState(DisconnectingState::enter(
+            Ok(TunnelCommand::Block(reason, allow_lan)) => NewState(DisconnectingState::enter(
                 shared_values,
                 (
                     self.close_handle,
                     self.tunnel_close_event,
-                    AfterDisconnect::Block(reason),
+                    AfterDisconnect::Block(reason, allow_lan),
                 ),
             )),
         }
@@ -179,7 +182,10 @@ impl TunnelState for ConnectedState {
                     (
                         connected_state.close_handle,
                         connected_state.tunnel_close_event,
-                        AfterDisconnect::Block(BlockReason::SetSecurityPolicyError),
+                        AfterDisconnect::Block(
+                            BlockReason::SetSecurityPolicyError,
+                            connected_state.tunnel_parameters.allow_lan,
+                        ),
                     ),
                 )
             }
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 184832acaa..8b2867b4fb 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -187,7 +187,10 @@ impl ConnectingState {
                             (
                                 self.close_handle,
                                 self.tunnel_close_event,
-                                AfterDisconnect::Block(BlockReason::SetSecurityPolicyError),
+                                AfterDisconnect::Block(
+                                    BlockReason::SetSecurityPolicyError,
+                                    allow_lan,
+                                ),
                             ),
                         ))
                     }
@@ -215,12 +218,12 @@ impl ConnectingState {
                     AfterDisconnect::Nothing,
                 ),
             )),
-            Ok(TunnelCommand::Block(reason)) => NewState(DisconnectingState::enter(
+            Ok(TunnelCommand::Block(reason, allow_lan)) => NewState(DisconnectingState::enter(
                 shared_values,
                 (
                     self.close_handle,
                     self.tunnel_close_event,
-                    AfterDisconnect::Block(reason),
+                    AfterDisconnect::Block(reason, allow_lan),
                 ),
             )),
         }
@@ -276,12 +279,12 @@ impl TunnelState for ConnectingState {
         shared_values: &mut SharedTunnelStateValues,
         parameters: Self::Bootstrap,
     ) -> (TunnelStateWrapper, TunnelStateTransition) {
+        let allow_lan = parameters.allow_lan;
         if let Err(error) =
             Self::set_security_policy(shared_values, parameters.endpoint, parameters.allow_lan)
         {
             error!("{}", error.display_chain());
-
-            return BlockedState::enter(shared_values, BlockReason::StartTunnelError);
+            return BlockedState::enter(shared_values, (BlockReason::StartTunnelError, allow_lan));
         }
 
         match Self::new(parameters) {
@@ -292,8 +295,7 @@ impl TunnelState for ConnectingState {
             Err(error) => {
                 let chained_error = error.chain_err(|| "Failed to start tunnel");
                 error!("{}", chained_error.display_chain());
-
-                BlockedState::enter(shared_values, BlockReason::StartTunnelError)
+                BlockedState::enter(shared_values, (BlockReason::StartTunnelError, allow_lan))
             }
         }
     }
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index 80a7eb5ce0..1a22043f0a 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -47,8 +47,8 @@ impl TunnelState for DisconnectedState {
             Ok(TunnelCommand::Connect(parameters)) => {
                 NewState(ConnectingState::enter(shared_values, parameters))
             }
-            Ok(TunnelCommand::Block(reason)) => {
-                NewState(BlockedState::enter(shared_values, reason))
+            Ok(TunnelCommand::Block(reason, allow_lan)) => {
+                NewState(BlockedState::enter(shared_values, (reason, allow_lan)))
             }
             Ok(_) => SameState(self),
             Err(_) => Finished,
diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index 452b8d3ac1..a7f20ca5ea 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -33,14 +33,16 @@ impl DisconnectingState {
         self.after_disconnect = match after_disconnect {
             AfterDisconnect::Nothing => match event {
                 Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
-                Ok(TunnelCommand::Block(reason)) => Block(reason),
+                Ok(TunnelCommand::Block(reason, allow_lan)) => Block(reason, allow_lan),
                 _ => Nothing,
             },
-            AfterDisconnect::Block(reason) => match event {
+            AfterDisconnect::Block(reason, allow_lan) => match event {
                 Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
                 Ok(TunnelCommand::Disconnect) => Nothing,
-                Ok(TunnelCommand::Block(new_reason)) => Block(new_reason),
-                _ => AfterDisconnect::Block(reason),
+                Ok(TunnelCommand::Block(new_reason, new_allow_lan)) => {
+                    Block(new_reason, new_allow_lan)
+                }
+                _ => Block(reason, allow_lan),
             },
             AfterDisconnect::Reconnect(mut tunnel_parameters) => match event {
                 Ok(TunnelCommand::AllowLan(allow_lan)) => {
@@ -49,7 +51,7 @@ impl DisconnectingState {
                 }
                 Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
                 Ok(TunnelCommand::Disconnect) | Err(_) => Nothing,
-                Ok(TunnelCommand::Block(reason)) => Block(reason),
+                Ok(TunnelCommand::Block(reason, allow_lan)) => Block(reason, allow_lan),
             },
         };
 
@@ -74,7 +76,9 @@ impl DisconnectingState {
     ) -> (TunnelStateWrapper, TunnelStateTransition) {
         match self.after_disconnect {
             AfterDisconnect::Nothing => DisconnectedState::enter(shared_values, ()),
-            AfterDisconnect::Block(reason) => BlockedState::enter(shared_values, reason),
+            AfterDisconnect::Block(reason, allow_lan) => {
+                BlockedState::enter(shared_values, (reason, allow_lan))
+            }
             AfterDisconnect::Reconnect(tunnel_parameters) => {
                 ConnectingState::enter(shared_values, tunnel_parameters)
             }
@@ -121,6 +125,6 @@ impl TunnelState for DisconnectingState {
 /// Which state should be transitioned to after disconnection is complete.
 pub enum AfterDisconnect {
     Nothing,
-    Block(BlockReason),
+    Block(BlockReason, bool),
     Reconnect(TunnelParameters),
 }
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 433333efa0..aed46f3ab5 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -111,7 +111,7 @@ pub enum TunnelCommand {
     /// Close tunnel connection.
     Disconnect,
     /// Disconnect any open tunnel and block all network access
-    Block(BlockReason),
+    Block(BlockReason, bool),
 }
 
 /// Information necessary to open a tunnel.