Merge branch 'restart-tunnel'

2 files changed, 51 insertions, 25 deletions

diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index 3710adc087..dd35239d91 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -31,6 +31,7 @@ mod settings;
 mod shutdown;
 
 use error_chain::ChainedError;
+use jsonrpc_core::futures::sync;
 use management_interface::{ManagementInterfaceServer, TunnelCommand};
 use mullvad_types::states::{DaemonState, SecurityState, TargetState};
 use std::io;
@@ -259,36 +260,59 @@ impl Daemon {
     fn handle_management_interface_event(&mut self, event: TunnelCommand) -> Result<()> {
         use TunnelCommand::*;
         match event {
-            SetTargetState(state) => {
-                if !self.shutdown {
-                    self.set_target_state(state)?;
-                } else {
-                    warn!("Ignoring target state change request due to shutdown");
-                }
-            }
-            GetState(tx) => {
-                if let Err(_) = tx.send(self.last_broadcasted_state) {
-                    warn!("Unable to send current state to management interface client",);
-                }
-            }
-            SetAccount(tx, account_token) => {
-                let save_result = self.settings.set_account_token(account_token);
-                match save_result.chain_err(|| "Unable to save settings") {
-                    Ok(()) => if let Err(_) = tx.send(()) {
-                        warn!("Unable to send response to management interface client");
-                    },
-                    Err(e) => error!("{}", e.display()),
+            SetTargetState(state) => self.on_set_target_state(state),
+            GetState(tx) => Ok(self.on_get_state(tx)),
+            SetAccount(tx, account_token) => self.on_set_account(tx, account_token),
+            GetAccount(tx) => Ok(self.on_get_account(tx)),
+        }
+    }
+
+    fn on_set_target_state(&mut self, new_target_state: TargetState) -> Result<()> {
+        if !self.shutdown {
+            self.set_target_state(new_target_state)
+        } else {
+            warn!("Ignoring target state change request due to shutdown");
+            Ok(())
+        }
+    }
+
+    fn on_get_state(&self, tx: sync::oneshot::Sender<DaemonState>) {
+        if let Err(_) = tx.send(self.last_broadcasted_state) {
+            warn!("Unable to send current state to management interface client",);
+        }
+    }
+
+    fn on_set_account(&mut self,
+                      tx: sync::oneshot::Sender<()>,
+                      account_token: Option<String>)
+                      -> Result<()> {
+
+        let save_result = self.settings.set_account_token(account_token);
+
+        match save_result.chain_err(|| "Unable to save settings") {
+            Ok(account_changed) => {
+                if let Err(_) = tx.send(()) {
+                    warn!("Unable to send response to management interface client");
                 }
-            }
-            GetAccount(tx) => {
-                if let Err(_) = tx.send(self.settings.get_account_token()) {
-                    warn!("Unable to send current account to management interface client");
+
+                let tunnel_needs_restart = self.state == TunnelState::Connecting ||
+                                           self.state == TunnelState::Connected;
+                if account_changed && tunnel_needs_restart {
+                    info!("Initiating tunnel restart because the account token changed");
+                    self.kill_tunnel()?;
                 }
             }
+            Err(e) => error!("{}", e.display()),
         }
         Ok(())
     }
 
+    fn on_get_account(&self, tx: sync::oneshot::Sender<Option<String>>) {
+        if let Err(_) = tx.send(self.settings.get_account_token()) {
+            warn!("Unable to send current account to management interface client");
+        }
+    }
+
     fn handle_management_interface_exited(&self, result: talpid_ipc::Result<()>) -> Result<()> {
         let error = ErrorKind::ManagementInterfaceError("Server exited unexpectedly");
         match result {
diff --git a/mullvad-daemon/src/settings.rs b/mullvad-daemon/src/settings.rs
index 229ccd9709..d2d652adf0 100644
--- a/mullvad-daemon/src/settings.rs
+++ b/mullvad-daemon/src/settings.rs
@@ -89,7 +89,8 @@ impl Settings {
     }
 
     /// Changes account number to the one given. Also saves the new settings to disk.
-    pub fn set_account_token(&mut self, account_token: Option<String>) -> Result<()> {
+    /// The boolean in the Result indicates if the account token changed or not
+    pub fn set_account_token(&mut self, account_token: Option<String>) -> Result<bool> {
         if account_token != self.account_token {
             info!(
                 "Changing account token from {} to {}",
@@ -98,8 +99,9 @@ impl Settings {
             );
             self.account_token = account_token;
             self.save()
+                .map(|_| true)
         } else {
-            Ok(())
+            Ok(false)
         }
     }