Merge branch 'wg-key-rpc-remove'

4 files changed, 97 insertions, 6 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1d48bf630a..29dd172836 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -32,6 +32,7 @@ Line wrap the file at 100 chars.                                              Th
 - Increase WireGuard ping timeout from 7 to 15 seconds.
 - Updated `wireguard-go` to `v0.0.20200121`
 - Use traffic data from WireGuard to infer connectivity to improve stability of the connection.
+- Remove WireGuard keys from accounts when they are removed from the local account history.
 
 #### Linux
 - DNS management with static `/etc/resolv.conf` will now work even when no
diff --git a/mullvad-daemon/src/account_history.rs b/mullvad-daemon/src/account_history.rs
index a5fae4ce6d..ea8f23757a 100644
--- a/mullvad-daemon/src/account_history.rs
+++ b/mullvad-daemon/src/account_history.rs
@@ -1,3 +1,11 @@
+#[cfg(target_os = "android")]
+use futures::future::{Executor, Future};
+#[cfg(not(target_os = "android"))]
+use futures::{
+    future::{self, Executor, Future},
+    sync::oneshot,
+};
+use mullvad_rpc::{HttpHandle, WireguardKeyProxy};
 use mullvad_types::{account::AccountToken, wireguard::WireguardData};
 use std::{
     collections::VecDeque,
@@ -6,6 +14,7 @@ use std::{
     path::Path,
 };
 use talpid_types::ErrorExt;
+use tokio_core::reactor::Remote;
 
 pub type Result<T> = std::result::Result<T, Error>;
 
@@ -29,11 +38,17 @@ static ACCOUNT_HISTORY_LIMIT: usize = 3;
 pub struct AccountHistory {
     file: io::BufWriter<fs::File>,
     accounts: VecDeque<AccountEntry>,
+    rpc_handle: HttpHandle,
+    tokio_remote: Remote,
 }
 
 
 impl AccountHistory {
-    pub fn new(cache_dir: &Path) -> Result<AccountHistory> {
+    pub fn new(
+        cache_dir: &Path,
+        rpc_handle: HttpHandle,
+        tokio_remote: Remote,
+    ) -> Result<AccountHistory> {
         let mut options = fs::OpenOptions::new();
         #[cfg(unix)]
         {
@@ -73,7 +88,12 @@ impl AccountHistory {
             Ok(accounts) => accounts,
         };
         let file = io::BufWriter::new(reader.into_inner());
-        let mut history = AccountHistory { file, accounts };
+        let mut history = AccountHistory {
+            file,
+            accounts,
+            rpc_handle,
+            tokio_remote,
+        };
         if let Err(e) = history.save_to_disk() {
             log::error!("Failed to save account cache after opening it: {}", e);
         }
@@ -126,15 +146,34 @@ impl AccountHistory {
         Ok(())
     }
 
+    fn create_remove_wg_key_rpc(
+        &self,
+        account: &str,
+        wg_data: &WireguardData,
+    ) -> impl Future<Item = (), Error = ()> {
+        let mut rpc = WireguardKeyProxy::new(self.rpc_handle.clone());
+        rpc.remove_wg_key(String::from(account), wg_data.private_key.public_key())
+            .map(|removed| log::debug!("Key existed on account: {}", removed))
+            .map_err(|e| log::error!("Failed to remove WireGuard key: {}", e))
+    }
+
     /// Always inserts a new entry at the start of the list
     pub fn insert(&mut self, new_entry: AccountEntry) -> Result<()> {
         self.accounts
             .retain(|entry| entry.account != new_entry.account);
 
         self.accounts.push_front(new_entry);
+
         if self.accounts.len() > ACCOUNT_HISTORY_LIMIT {
-            let _ = self.accounts.pop_back();
+            let last_entry = self.accounts.pop_back().unwrap();
+            if let Some(wg_data) = last_entry.wireguard {
+                let fut = self.create_remove_wg_key_rpc(&last_entry.account, &wg_data);
+                if let Err(e) = self.tokio_remote.execute(fut) {
+                    log::error!("Failed to spawn future to remove WireGuard key: {:?}", e);
+                }
+            }
         }
+
         self.save_to_disk()
     }
 
@@ -148,13 +187,55 @@ impl AccountHistory {
 
     /// Remove account data
     pub fn remove_account(&mut self, account: &str) -> Result<()> {
-        self.accounts.retain(|entry| entry.account != account);
+        let entry = self.get(&String::from(account))?;
+        let entry = match entry {
+            Some(entry) => entry,
+            None => return Ok(()),
+        };
+
+        if let Some(wg_data) = entry.wireguard {
+            let fut = self.create_remove_wg_key_rpc(account, &wg_data);
+            if let Err(e) = self.tokio_remote.execute(fut) {
+                log::error!("Failed to spawn future to remove WireGuard key: {:?}", e);
+            }
+        }
+
+        let _ = self.accounts.pop_front();
         self.save_to_disk()
     }
 
     /// Remove account history
     #[cfg(not(target_os = "android"))]
     pub fn clear(&mut self) -> Result<()> {
+        let mut rpc = WireguardKeyProxy::new(self.rpc_handle.clone());
+
+        log::debug!("account_history::clear");
+
+        let mut removal_futures = Vec::with_capacity(ACCOUNT_HISTORY_LIMIT);
+
+        for entry in self.accounts.iter() {
+            if let Some(wg_data) = &entry.wireguard {
+                let fut = rpc
+                    .remove_wg_key(entry.account.clone(), wg_data.private_key.public_key())
+                    .map(|_| ())
+                    .map_err(|e| log::error!("Failed to remove WireGuard key: {}", e));
+                removal_futures.push(fut);
+            }
+        }
+
+        let joined_futs = future::join_all(removal_futures);
+        let (tx, rx) = oneshot::channel();
+
+        let execute_result = self.tokio_remote.execute(joined_futs.then(|result| {
+            let _ = tx.send(result);
+            Ok(())
+        }));
+        if let Err(e) = execute_result {
+            log::error!("Failed to spawn future to remove WireGuard keys: {:?}", e);
+        } else {
+            let _ = rx.wait();
+        }
+
         self.accounts = VecDeque::new();
         self.save_to_disk()
     }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 06f122fdf9..fe59f86d93 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -418,8 +418,12 @@ where
 
         let settings = settings::load();
 
-        let account_history =
-            account_history::AccountHistory::new(&cache_dir).map_err(Error::LoadAccountHistory)?;
+        let account_history = account_history::AccountHistory::new(
+            &cache_dir,
+            rpc_handle.clone(),
+            tokio_remote.clone(),
+        )
+        .map_err(Error::LoadAccountHistory)?;
 
         let tunnel_parameters_generator = MullvadTunnelParametersGenerator {
             tx: internal_event_tx.clone(),
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index d16029c16d..23dd2b6bea 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -149,4 +149,9 @@ jsonrpc_client!(pub struct WireguardKeyProxy {
         account_token: AccountToken,
         public_key: wireguard::PublicKey
     ) -> RpcRequest<bool>;
+    pub fn remove_wg_key(
+        &mut self,
+        account_token: AccountToken,
+        public_key: wireguard::PublicKey
+    ) -> RpcRequest<bool>;
 });