Handle security policy in tunnel state machine

5 files changed, 167 insertions, 39 deletions

diff --git a/mullvad-daemon/src/tunnel_state_machine/connected_state.rs b/mullvad-daemon/src/tunnel_state_machine/connected_state.rs
index 713a285159..332608f055 100644
--- a/mullvad-daemon/src/tunnel_state_machine/connected_state.rs
+++ b/mullvad-daemon/src/tunnel_state_machine/connected_state.rs
@@ -1,13 +1,14 @@
 use futures::sync::{mpsc, oneshot};
 use futures::{Async, Future, Stream};
 
+use talpid_core::firewall::{Firewall, SecurityPolicy};
 use talpid_core::tunnel::{CloseHandle, TunnelEvent, TunnelMetadata};
 use talpid_types::net::TunnelEndpoint;
 
 use super::{
-    AfterDisconnect, ConnectingState, DisconnectingState, EventConsequence,
+    AfterDisconnect, ConnectingState, DisconnectingState, EventConsequence, Result, ResultExt,
     SharedTunnelStateValues, StateEntryResult, TunnelCommand, TunnelParameters, TunnelState,
-    TunnelStateTransition, TunnelStateWrapper,
+    TunnelStateWrapper,
 };
 
 pub struct ConnectedStateBootstrap {
@@ -41,12 +42,22 @@ impl ConnectedState {
         }
     }
 
-    pub fn info(&self) -> TunnelStateTransition {
-        TunnelStateTransition::Connected(self.tunnel_endpoint, self.metadata.clone())
+    fn set_security_policy(&self, shared_values: &mut SharedTunnelStateValues) -> Result<()> {
+        let policy = SecurityPolicy::Connected {
+            relay_endpoint: self.tunnel_endpoint.to_endpoint(),
+            tunnel: self.metadata.clone(),
+            allow_lan: self.tunnel_parameters.allow_lan,
+        };
+
+        debug!("Set security policy: {:?}", policy);
+        shared_values
+            .firewall
+            .apply_policy(policy)
+            .chain_err(|| "Failed to apply security policy for connected state")
     }
 
     fn handle_commands(
-        self,
+        mut self,
         commands: &mut mpsc::UnboundedReceiver<TunnelCommand>,
         shared_values: &mut SharedTunnelStateValues,
     ) -> EventConsequence<Self> {
@@ -75,6 +86,24 @@ impl ConnectedState {
                     AfterDisconnect::Nothing,
                 ),
             )),
+            Ok(TunnelCommand::AllowLan(allow_lan)) => {
+                self.tunnel_parameters.allow_lan = allow_lan;
+
+                match self.set_security_policy(shared_values) {
+                    Ok(()) => SameState(self),
+                    Err(error) => {
+                        error!("{}", error.chain_err(|| "Failed to update security policy"));
+                        NewState(DisconnectingState::enter(
+                            shared_values,
+                            (
+                                self.close_handle,
+                                self.tunnel_close_event,
+                                AfterDisconnect::Nothing,
+                            ),
+                        ))
+                    }
+                }
+            }
         }
     }
 
@@ -120,8 +149,26 @@ impl ConnectedState {
 impl TunnelState for ConnectedState {
     type Bootstrap = ConnectedStateBootstrap;
 
-    fn enter(_: &mut SharedTunnelStateValues, bootstrap: Self::Bootstrap) -> StateEntryResult {
-        Ok(TunnelStateWrapper::from(ConnectedState::from(bootstrap)))
+    fn enter(
+        shared_values: &mut SharedTunnelStateValues,
+        bootstrap: Self::Bootstrap,
+    ) -> StateEntryResult {
+        let connected_state = ConnectedState::from(bootstrap);
+
+        match connected_state.set_security_policy(shared_values) {
+            Ok(()) => Ok(TunnelStateWrapper::from(connected_state)),
+            Err(error) => Err((
+                error,
+                DisconnectingState::enter(
+                    shared_values,
+                    (
+                        connected_state.close_handle,
+                        connected_state.tunnel_close_event,
+                        AfterDisconnect::Nothing,
+                    ),
+                ).expect("Failed to disconnect after failed transition to connected state"),
+            )),
+        }
     }
 
     fn handle_event(
diff --git a/mullvad-daemon/src/tunnel_state_machine/connecting_state.rs b/mullvad-daemon/src/tunnel_state_machine/connecting_state.rs
index e3590133a5..c134618cb3 100644
--- a/mullvad-daemon/src/tunnel_state_machine/connecting_state.rs
+++ b/mullvad-daemon/src/tunnel_state_machine/connecting_state.rs
@@ -8,14 +8,15 @@ use futures::sink::Wait;
 use futures::sync::{mpsc, oneshot};
 use futures::{Async, Future, Sink, Stream};
 
+use talpid_core::firewall::{Firewall, SecurityPolicy};
 use talpid_core::tunnel::{CloseHandle, TunnelEvent, TunnelMetadata, TunnelMonitor};
 use talpid_types::net::{TunnelEndpoint, TunnelEndpointData};
 
 use super::{
     AfterDisconnect, ConnectedState, ConnectedStateBootstrap, DisconnectedState,
     DisconnectingState, EventConsequence, Result, ResultExt, SharedTunnelStateValues,
-    StateEntryResult, TunnelCommand, TunnelParameters, TunnelState, TunnelStateTransition,
-    TunnelStateWrapper, OPENVPN_LOG_FILENAME, WIREGUARD_LOG_FILENAME,
+    StateEntryResult, TunnelCommand, TunnelParameters, TunnelState, TunnelStateWrapper,
+    OPENVPN_LOG_FILENAME, WIREGUARD_LOG_FILENAME,
 };
 use logging;
 
@@ -36,7 +37,12 @@ pub struct ConnectingState {
 }
 
 impl ConnectingState {
-    fn new(parameters: TunnelParameters) -> Result<Self> {
+    fn new(
+        shared_values: &mut SharedTunnelStateValues,
+        parameters: TunnelParameters,
+    ) -> Result<Self> {
+        Self::set_security_policy(shared_values, parameters.endpoint, parameters.allow_lan)?;
+
         let tunnel_endpoint = parameters.endpoint;
         let (tunnel_events, tunnel_close_event, close_handle) = Self::start_tunnel(&parameters)?;
 
@@ -49,6 +55,23 @@ impl ConnectingState {
         })
     }
 
+    fn set_security_policy(
+        shared_values: &mut SharedTunnelStateValues,
+        endpoint: TunnelEndpoint,
+        allow_lan: bool,
+    ) -> Result<()> {
+        let policy = SecurityPolicy::Connecting {
+            relay_endpoint: endpoint.to_endpoint(),
+            allow_lan,
+        };
+
+        debug!("Set security policy: {:?}", policy);
+        shared_values
+            .firewall
+            .apply_policy(policy)
+            .chain_err(|| "Failed to apply security policy for connecting state")
+    }
+
     fn start_tunnel(
         parameters: &TunnelParameters,
     ) -> Result<(
@@ -145,12 +168,8 @@ impl ConnectingState {
         }
     }
 
-    pub fn info(&self) -> TunnelStateTransition {
-        TunnelStateTransition::Connecting(self.tunnel_endpoint)
-    }
-
     fn handle_commands(
-        self,
+        mut self,
         commands: &mut mpsc::UnboundedReceiver<TunnelCommand>,
         shared_values: &mut SharedTunnelStateValues,
     ) -> EventConsequence<Self> {
@@ -179,6 +198,23 @@ impl ConnectingState {
                     AfterDisconnect::Nothing,
                 ),
             )),
+            Ok(TunnelCommand::AllowLan(allow_lan)) => {
+                self.tunnel_parameters.allow_lan = allow_lan;
+                match Self::set_security_policy(shared_values, self.tunnel_endpoint, allow_lan) {
+                    Ok(()) => SameState(self),
+                    Err(error) => {
+                        error!("{}", error.chain_err(|| "Failed to update security policy"));
+                        NewState(DisconnectingState::enter(
+                            shared_values,
+                            (
+                                self.close_handle,
+                                self.tunnel_close_event,
+                                AfterDisconnect::Nothing,
+                            ),
+                        ))
+                    }
+                }
+            }
         }
     }
 
@@ -232,7 +268,7 @@ impl TunnelState for ConnectingState {
         shared_values: &mut SharedTunnelStateValues,
         parameters: Self::Bootstrap,
     ) -> StateEntryResult {
-        Self::new(parameters)
+        Self::new(shared_values, parameters)
             .map(TunnelStateWrapper::from)
             .chain_err(|| "Failed to start tunnel")
             .map_err(|error| {
diff --git a/mullvad-daemon/src/tunnel_state_machine/disconnected_state.rs b/mullvad-daemon/src/tunnel_state_machine/disconnected_state.rs
index 1bc0d1f556..e01789094b 100644
--- a/mullvad-daemon/src/tunnel_state_machine/disconnected_state.rs
+++ b/mullvad-daemon/src/tunnel_state_machine/disconnected_state.rs
@@ -1,18 +1,33 @@
+use error_chain::ChainedError;
 use futures::sync::mpsc;
 use futures::Stream;
 
+use talpid_core::firewall::Firewall;
+
 use super::{
-    ConnectingState, EventConsequence, SharedTunnelStateValues, StateEntryResult, TunnelCommand,
-    TunnelState, TunnelStateWrapper,
+    ConnectingState, Error, EventConsequence, SharedTunnelStateValues, StateEntryResult,
+    TunnelCommand, TunnelState, TunnelStateWrapper,
 };
 
 /// No tunnel is running.
 pub struct DisconnectedState;
 
+impl DisconnectedState {
+    fn reset_security_policy(shared_values: &mut SharedTunnelStateValues) {
+        debug!("Reset security policy");
+        if let Err(error) = shared_values.firewall.reset_policy() {
+            let chained_error = Error::with_chain(error, "Failed to reset security policy");
+            error!("{}", chained_error.display_chain());
+        }
+    }
+}
+
 impl TunnelState for DisconnectedState {
     type Bootstrap = ();
 
-    fn enter(_: &mut SharedTunnelStateValues, _: Self::Bootstrap) -> StateEntryResult {
+    fn enter(shared_values: &mut SharedTunnelStateValues, _: Self::Bootstrap) -> StateEntryResult {
+        Self::reset_security_policy(shared_values);
+
         Ok(TunnelStateWrapper::from(DisconnectedState))
     }
 
@@ -27,7 +42,7 @@ impl TunnelState for DisconnectedState {
             Ok(TunnelCommand::Connect(parameters)) => {
                 NewState(ConnectingState::enter(shared_values, parameters))
             }
-            Ok(TunnelCommand::Disconnect) | Err(_) => SameState(self),
+            _ => SameState(self),
         }
     }
 }
diff --git a/mullvad-daemon/src/tunnel_state_machine/disconnecting_state.rs b/mullvad-daemon/src/tunnel_state_machine/disconnecting_state.rs
index a57a6e1e9c..9785264da1 100644
--- a/mullvad-daemon/src/tunnel_state_machine/disconnecting_state.rs
+++ b/mullvad-daemon/src/tunnel_state_machine/disconnecting_state.rs
@@ -23,9 +23,22 @@ impl DisconnectingState {
     ) -> EventConsequence<Self> {
         use self::AfterDisconnect::*;
 
-        self.after_disconnect = match try_handle_event!(self, commands.poll()) {
-            Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
-            Ok(TunnelCommand::Disconnect) | Err(_) => Nothing,
+        let event = try_handle_event!(self, commands.poll());
+        let after_disconnect = self.after_disconnect;
+
+        self.after_disconnect = match after_disconnect {
+            AfterDisconnect::Nothing => match event {
+                Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
+                _ => Nothing,
+            },
+            AfterDisconnect::Reconnect(mut tunnel_parameters) => match event {
+                Ok(TunnelCommand::Connect(parameters)) => Reconnect(parameters),
+                Ok(TunnelCommand::AllowLan(allow_lan)) => {
+                    tunnel_parameters.allow_lan = allow_lan;
+                    Reconnect(tunnel_parameters)
+                }
+                Ok(TunnelCommand::Disconnect) | Err(_) => Nothing,
+            },
         };
 
         EventConsequence::SameState(self)
diff --git a/mullvad-daemon/src/tunnel_state_machine/mod.rs b/mullvad-daemon/src/tunnel_state_machine/mod.rs
index 3ab7b87475..2d94bd066e 100644
--- a/mullvad-daemon/src/tunnel_state_machine/mod.rs
+++ b/mullvad-daemon/src/tunnel_state_machine/mod.rs
@@ -7,7 +7,7 @@ mod disconnected_state;
 mod disconnecting_state;
 
 use std::fmt::{Debug, Formatter, Result as FmtResult};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::mpsc as sync_mpsc;
 use std::thread;
 
@@ -17,8 +17,8 @@ use futures::{Async, Future, Poll, Stream};
 use tokio_core::reactor::Core;
 
 use mullvad_types::account::AccountToken;
+use talpid_core::firewall::{Firewall, FirewallProxy};
 use talpid_core::mpsc::IntoSender;
-use talpid_core::tunnel::TunnelMetadata;
 use talpid_types::net::{TunnelEndpoint, TunnelOptions};
 
 use self::connected_state::{ConnectedState, ConnectedStateBootstrap};
@@ -29,6 +29,9 @@ use super::{OPENVPN_LOG_FILENAME, WIREGUARD_LOG_FILENAME};
 
 error_chain! {
     errors {
+        FirewallError {
+            description("Firewall error")
+        }
         ReactorError {
             description("Failed to initialize tunnel state machine event loop executor")
         }
@@ -36,17 +39,19 @@ error_chain! {
 }
 
 /// Spawn the tunnel state machine thread, returning a channel for sending tunnel commands.
-pub fn spawn<T>(
+pub fn spawn<P, T>(
+    cache_dir: P,
     state_change_listener: IntoSender<TunnelStateTransition, T>,
 ) -> Result<mpsc::UnboundedSender<TunnelCommand>>
 where
+    P: AsRef<Path> + Send + 'static,
     T: From<TunnelStateTransition> + Send + 'static,
 {
     let (command_tx, command_rx) = mpsc::unbounded();
     let (startup_result_tx, startup_result_rx) = sync_mpsc::channel();
 
     thread::spawn(
-        move || match create_event_loop(command_rx, state_change_listener) {
+        move || match create_event_loop(cache_dir, command_rx, state_change_listener) {
             Ok((mut reactor, event_loop)) => {
                 startup_result_tx.send(Ok(())).expect(
                     "Tunnel state machine won't be started because the owner thread crashed",
@@ -72,15 +77,17 @@ where
         .map(|_| command_tx)
 }
 
-fn create_event_loop<T>(
+fn create_event_loop<P, T>(
+    cache_dir: P,
     commands: mpsc::UnboundedReceiver<TunnelCommand>,
     state_change_listener: IntoSender<TunnelStateTransition, T>,
 ) -> Result<(Core, impl Future<Item = (), Error = Error>)>
 where
+    P: AsRef<Path>,
     T: From<TunnelStateTransition> + Send + 'static,
 {
     let reactor = Core::new().chain_err(|| ErrorKind::ReactorError)?;
-    let state_machine = TunnelStateMachine::new(commands);
+    let state_machine = TunnelStateMachine::new(&cache_dir, commands)?;
 
     let future = state_machine.for_each(move |state_change_event| {
         state_change_listener
@@ -93,6 +100,8 @@ where
 
 /// Representation of external commands for the tunnel state machine.
 pub enum TunnelCommand {
+    /// Enable or disable LAN access in the firewall.
+    AllowLan(bool),
     /// Open tunnel connection.
     Connect(TunnelParameters),
     /// Close tunnel connection.
@@ -107,14 +116,15 @@ pub struct TunnelParameters {
     pub log_dir: Option<PathBuf>,
     pub resource_dir: PathBuf,
     pub account_token: AccountToken,
+    pub allow_lan: bool,
 }
 
 /// Event resulting from a transition to a new tunnel state.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum TunnelStateTransition {
     Disconnected,
-    Connecting(TunnelEndpoint),
-    Connected(TunnelEndpoint, TunnelMetadata),
+    Connecting,
+    Connected,
     Disconnecting,
 }
 
@@ -131,16 +141,21 @@ struct TunnelStateMachine {
 }
 
 impl TunnelStateMachine {
-    fn new(commands: mpsc::UnboundedReceiver<TunnelCommand>) -> Self {
-        let mut shared_values = SharedTunnelStateValues;
+    fn new<P: AsRef<Path>>(
+        cache_dir: P,
+        commands: mpsc::UnboundedReceiver<TunnelCommand>,
+    ) -> Result<Self> {
+        let firewall = FirewallProxy::new(cache_dir).chain_err(|| ErrorKind::FirewallError)?;
+        let mut shared_values = SharedTunnelStateValues { firewall };
+
         let initial_state = TunnelStateWrapper::enter(&mut shared_values, ())
             .expect("Failed to create initial tunnel state");
 
-        TunnelStateMachine {
+        Ok(TunnelStateMachine {
             current_state: Some(initial_state),
             commands,
             shared_values,
-        }
+        })
     }
 }
 
@@ -206,7 +221,9 @@ impl From<EventConsequence<TunnelStateWrapper>> for TunnelStateMachineAction {
 }
 
 /// Values that are common to all tunnel states.
-struct SharedTunnelStateValues;
+struct SharedTunnelStateValues {
+    firewall: FirewallProxy,
+}
 
 /// Asynchronous result of an attempt to progress a state.
 enum EventConsequence<T: TunnelState> {
@@ -291,8 +308,8 @@ impl TunnelStateWrapper {
     fn info(&self) -> TunnelStateTransition {
         match *self {
             TunnelStateWrapper::Disconnected(_) => TunnelStateTransition::Disconnected,
-            TunnelStateWrapper::Connecting(ref state) => state.info(),
-            TunnelStateWrapper::Connected(ref state) => state.info(),
+            TunnelStateWrapper::Connecting(_) => TunnelStateTransition::Connecting,
+            TunnelStateWrapper::Connected(_) => TunnelStateTransition::Connected,
             TunnelStateWrapper::Disconnecting(_) => TunnelStateTransition::Disconnecting,
         }
     }