Rework account history

3 files changed, 187 insertions, 126 deletions

diff --git a/mullvad-daemon/src/account_history.rs b/mullvad-daemon/src/account_history.rs
index de57b60d3f..e79356d257 100644
--- a/mullvad-daemon/src/account_history.rs
+++ b/mullvad-daemon/src/account_history.rs
@@ -1,22 +1,18 @@
-use mullvad_types::account::AccountToken;
+use mullvad_types::{account::AccountToken, wireguard::WireguardData};
 use std::{
-    fs::File,
-    io,
-    path::{Path, PathBuf},
+    collections::VecDeque,
+    fs,
+    io::{self, Seek, Write},
+    path::Path,
 };
 
 error_chain! {
     errors {
-        ReadError(path: PathBuf) {
+        ReadError {
             description("Unable to read account history file")
-            display("Unable to read account history from {}", path.display())
         }
-        WriteError(path: PathBuf) {
+        WriteError {
             description("Unable to write account history file")
-            display("Unable to write account history to {}", path.display())
-        }
-        ParseError {
-            description("Malformed account history")
         }
     }
 }
@@ -24,84 +20,142 @@ error_chain! {
 static ACCOUNT_HISTORY_FILE: &str = "account-history.json";
 static ACCOUNT_HISTORY_LIMIT: usize = 3;
 
-#[derive(Debug, Clone, Deserialize, Serialize, Default)]
+/// A trivial MRU cache of account data
 pub struct AccountHistory {
-    accounts: Vec<AccountToken>,
-    #[serde(skip)]
-    cache_path: PathBuf,
+    file: io::BufWriter<fs::File>,
+    accounts: VecDeque<AccountEntry>,
 }
 
+
 impl AccountHistory {
-    /// Returns a new empty `AccountHistory` ready to load from, or save to, the given cache dir.
-    pub fn new(cache_dir: &Path) -> AccountHistory {
-        AccountHistory {
-            accounts: Vec::new(),
-            cache_path: cache_dir.join(ACCOUNT_HISTORY_FILE),
+    pub fn new(cache_dir: &Path) -> Result<AccountHistory> {
+        let mut options = fs::OpenOptions::new();
+        #[cfg(unix)]
+        {
+            use std::os::unix::fs::OpenOptionsExt;
+            options.mode(0o600);
         }
-    }
+        #[cfg(windows)]
+        {
+            use std::os::windows::fs::OpenOptionsExt;
+            // a share mode of zero ensures exclusive access to the file to *this* process
+            options.share_mode(0);
+        }
+        let path = cache_dir.join(ACCOUNT_HISTORY_FILE);
+        log::info!("Opening account history file in {}", path.display());
+        let mut reader = options
+            .write(true)
+            .read(true)
+            .create(true)
+            .open(path)
+            .map(io::BufReader::new)
+            .chain_err(|| ErrorKind::ReadError)?;
 
-    /// Loads account history from file. If no file is present this does nothing.
-    pub fn load(&mut self) -> Result<()> {
-        match File::open(&self.cache_path).map(io::BufReader::new) {
-            Ok(mut file) => {
-                log::info!(
-                    "Loading account history from {}",
-                    &self.cache_path.display()
-                );
-                self.accounts = Self::parse(&mut file)?.accounts;
-                Ok(())
+        let accounts: VecDeque<AccountEntry> = match serde_json::from_reader(&mut reader) {
+            Err(e) => {
+                log::error!("Failed to read account history - {}", e);
+                Self::try_old_format(&mut reader)?
+                    .into_iter()
+                    .map(|account| AccountEntry {
+                        account,
+                        wireguard: None,
+                    })
+                    .collect()
             }
-            Err(ref e) if e.kind() == io::ErrorKind::NotFound => {
-                log::info!("No account history file at {}", &self.cache_path.display());
-                Ok(())
-            }
-            Err(e) => Err(e).chain_err(|| ErrorKind::ReadError(self.cache_path.clone())),
-        }
+            Ok(accounts) => accounts,
+        };
+        let file = io::BufWriter::new(reader.into_inner());
+        Ok(AccountHistory { file, accounts })
+    }
+
+    fn try_old_format(reader: &mut io::BufReader<fs::File>) -> Result<Vec<AccountToken>> {
+        reader
+            .seek(io::SeekFrom::Start(0))
+            .chain_err(|| ErrorKind::ReadError)?;
+        Ok(serde_json::from_reader(reader).unwrap_or(vec![]))
     }
 
-    fn parse(file: &mut impl io::Read) -> Result<AccountHistory> {
-        serde_json::from_reader(file).chain_err(|| ErrorKind::ParseError)
+    /// Gets account data for a certain account id and bumps it's entry to the top of the list if
+    /// it isn't there already. Returns None if the account entry is not available.
+    pub fn get(&mut self, account: &AccountToken) -> Result<Option<AccountEntry>> {
+        let (idx, entry) = match self
+            .accounts
+            .iter()
+            .enumerate()
+            .find(|(_idx, entry)| &entry.account == account)
+        {
+            Some((idx, entry)) => (idx, entry.clone()),
+            None => {
+                return Ok(None);
+            }
+        };
+        // this account is already on top
+        if idx == 0 {
+            return Ok(Some(entry));
+        }
+        self.insert(entry.clone())?;
+        Ok(Some(entry))
     }
 
-    pub fn get_accounts(&self) -> &[AccountToken] {
-        &self.accounts
+    /// Bumps history of an account token. If the account token is not in history, it will be
+    /// added.
+    pub fn bump_history(&mut self, account: &AccountToken) -> Result<()> {
+        if let None = self.get(account)? {
+            let new_entry = AccountEntry {
+                account: account.to_string(),
+                wireguard: None,
+            };
+            self.insert(new_entry)?;
+        }
+        Ok(())
     }
 
-    /// Add account token to the account history removing duplicate entries
-    pub fn add_account_token(&mut self, account_token: AccountToken) -> Result<()> {
+    /// Always inserts a new entry at the start of the list
+    pub fn insert(&mut self, new_entry: AccountEntry) -> Result<()> {
         self.accounts
-            .retain(|existing_token| existing_token != &account_token);
-        self.accounts.push(account_token);
+            .retain(|entry| entry.account != new_entry.account);
 
-        let num_accounts = self.accounts.len();
-        if num_accounts > ACCOUNT_HISTORY_LIMIT {
-            self.accounts = self
-                .accounts
-                .split_off(num_accounts - ACCOUNT_HISTORY_LIMIT);
+        self.accounts.push_front(new_entry);
+        if self.accounts.len() > ACCOUNT_HISTORY_LIMIT {
+            let _ = self.accounts.pop_back();
         }
-
-        self.save()
+        self.save_to_disk()
     }
 
-    /// Remove account token from the account history
-    pub fn remove_account_token(&mut self, account_token: &AccountToken) -> Result<()> {
+    /// Retrieve account history.
+    pub fn get_account_history(&self) -> Vec<AccountToken> {
         self.accounts
-            .retain(|existing_token| existing_token != account_token);
-        self.save()
+            .iter()
+            .map(|entry| entry.account.clone())
+            .collect()
     }
 
-    /// Serializes the account history and saves it to the file it was loaded from.
-    fn save(&self) -> Result<()> {
-        log::debug!("Writing account history to {}", self.cache_path.display());
-        let mut file = File::create(&self.cache_path)
-            .map(io::BufWriter::new)
-            .chain_err(|| ErrorKind::WriteError(self.cache_path.clone()))?;
-
-        serde_json::to_writer_pretty(&mut file, self)
-            .chain_err(|| ErrorKind::WriteError(self.cache_path.clone()))?;
+    /// Remove account data
+    pub fn remove_account(&mut self, account: &str) -> Result<()> {
+        self.accounts.retain(|entry| entry.account != account);
+        self.save_to_disk()
+    }
 
-        file.get_mut()
+    fn save_to_disk(&mut self) -> Result<()> {
+        self.file
+            .get_mut()
+            .set_len(0)
+            .chain_err(|| ErrorKind::WriteError)?;
+        self.file
+            .seek(io::SeekFrom::Start(0))
+            .chain_err(|| ErrorKind::WriteError)?;
+        serde_json::to_writer_pretty(&mut self.file, &self.accounts)
+            .chain_err(|| ErrorKind::WriteError)?;
+        self.file.flush().chain_err(|| ErrorKind::WriteError)?;
+        self.file
+            .get_mut()
             .sync_all()
-            .chain_err(|| ErrorKind::WriteError(self.cache_path.clone()))
+            .chain_err(|| ErrorKind::WriteError)
     }
 }
+
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct AccountEntry {
+    pub account: AccountToken,
+    pub wireguard: Option<WireguardData>,
+}
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 0faf2b81ff..39c89adf53 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -69,6 +69,7 @@ error_chain! {
     }
     links {
         TunnelError(tunnel_state_machine::Error, tunnel_state_machine::ErrorKind);
+        AccountHistory(account_history::Error, account_history::ErrorKind);
     }
 }
 
@@ -156,6 +157,8 @@ pub struct Daemon {
     #[cfg(unix)]
     management_interface_socket_path: String,
     settings: Settings,
+    account_history: account_history::AccountHistory,
+    //wg_key_proxy: WireguardKeyProxy<HttpHandle>,
     accounts_proxy: AccountsProxy<HttpHandle>,
     version_proxy: AppVersionProxy<HttpHandle>,
     https_handle: mullvad_rpc::rest::RequestSender,
@@ -195,6 +198,9 @@ impl Daemon {
         let relay_selector =
             relays::RelaySelector::new(rpc_handle.clone(), &resource_dir, &cache_dir);
         let settings = Settings::load().chain_err(|| "Unable to read settings")?;
+        let account_history = account_history::AccountHistory::new(&cache_dir)
+            .chain_err(|| "Unable to read wireguard key cache")?;
+
 
         let (tx, rx) = mpsc::channel();
         let tunnel_parameters_generator = MullvadTunnelParametersGenerator { tx: tx.clone() };
@@ -209,7 +215,7 @@ impl Daemon {
         )?;
 
         let target_state = TargetState::Unsecured;
-        let management_interface_result = Self::start_management_interface(tx.clone(), cache_dir)?;
+        let management_interface_result = Self::start_management_interface(tx.clone())?;
 
         // Attempt to download a fresh relay list
         relay_selector.update();
@@ -226,6 +232,8 @@ impl Daemon {
             #[cfg(unix)]
             management_interface_socket_path: management_interface_result.1,
             settings,
+            account_history,
+            // wg_key_proxy: WireguardKeyProxy::new(rpc_handle.clone()),
             accounts_proxy: AccountsProxy::new(rpc_handle.clone()),
             version_proxy: AppVersionProxy::new(rpc_handle),
             https_handle,
@@ -240,10 +248,9 @@ impl Daemon {
     // Returns a handle that allows notifying all subscribers on events.
     fn start_management_interface(
         event_tx: mpsc::Sender<DaemonEvent>,
-        cache_dir: PathBuf,
     ) -> Result<(management_interface::EventBroadcaster, String)> {
         let multiplex_event_tx = IntoSender::from(event_tx.clone());
-        let server = Self::start_management_interface_server(multiplex_event_tx, cache_dir)?;
+        let server = Self::start_management_interface_server(multiplex_event_tx)?;
         let event_broadcaster = server.event_broadcaster();
         let socket_path = server.socket_path().to_owned();
         Self::spawn_management_interface_wait_thread(server, event_tx);
@@ -252,9 +259,8 @@ impl Daemon {
 
     fn start_management_interface_server(
         event_tx: IntoSender<ManagementCommand, DaemonEvent>,
-        cache_dir: PathBuf,
     ) -> Result<ManagementInterfaceServer> {
-        let server = ManagementInterfaceServer::start(event_tx, cache_dir)
+        let server = ManagementInterfaceServer::start(event_tx)
             .chain_err(|| ErrorKind::ManagementInterfaceError("Failed to start server"))?;
         info!(
             "Mullvad management interface listening on {}",
@@ -423,6 +429,10 @@ impl Daemon {
             GetRelayLocations(tx) => self.on_get_relay_locations(tx),
             UpdateRelayLocations => self.on_update_relay_locations(),
             SetAccount(tx, account_token) => self.on_set_account(tx, account_token),
+            GetAccountHistory(tx) => self.on_get_account_history(tx),
+            RemoveAccountFromHistory(tx, account_token) => {
+                self.on_remove_account_from_history(tx, account_token)
+            }
             UpdateRelaySettings(tx, update) => self.on_update_relay_settings(tx, update),
             SetAllowLan(tx, allow_lan) => self.on_set_allow_lan(tx, allow_lan),
             SetBlockWhenDisconnected(tx, block_when_disconnected) => {
@@ -537,8 +547,7 @@ impl Daemon {
     }
 
     fn on_set_account(&mut self, tx: oneshot::Sender<()>, account_token: Option<String>) {
-        let account_token_cleared = account_token.is_none();
-        let save_result = self.settings.set_account_token(account_token);
+        let save_result = self.settings.set_account_token(account_token.clone());
 
         match save_result.chain_err(|| "Unable to save settings") {
             Ok(account_changed) => {
@@ -546,12 +555,18 @@ impl Daemon {
                 if account_changed {
                     self.management_interface_broadcaster
                         .notify_settings(&self.settings);
-                    if account_token_cleared {
-                        info!("Disconnecting because account token was cleared");
-                        self.set_target_state(TargetState::Unsecured);
-                    } else {
-                        info!("Initiating tunnel restart because the account token changed");
-                        self.reconnect_tunnel();
+                    match account_token {
+                        Some(token) => {
+                            if let Err(e) = self.account_history.bump_history(&token) {
+                                log::error!("Failed to bump account history: {}", e);
+                            }
+                            info!("Initiating tunnel restart because the account token changed");
+                            self.reconnect_tunnel();
+                        }
+                        None => {
+                            info!("Disconnecting because account token was cleared");
+                            self.set_target_state(TargetState::Unsecured);
+                        }
                     }
                 }
             }
@@ -559,6 +574,24 @@ impl Daemon {
         }
     }
 
+    fn on_get_account_history(&mut self, tx: oneshot::Sender<Vec<AccountToken>>) {
+        Self::oneshot_send(
+            tx,
+            self.account_history.get_account_history(),
+            "get_account_history response",
+        );
+    }
+
+    fn on_remove_account_from_history(
+        &mut self,
+        tx: oneshot::Sender<()>,
+        account_token: AccountToken,
+    ) {
+        if let Ok(_) = self.account_history.remove_account(&account_token) {
+            Self::oneshot_send(tx, (), "remove_account_from_history response");
+        }
+    }
+
     fn on_get_version_info(
         &mut self,
         tx: oneshot::Sender<BoxFuture<AppVersionInfo, mullvad_rpc::Error>>,
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index aeeda483a7..7eecaf5056 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -1,4 +1,3 @@
-use crate::account_history::{AccountHistory, Error as AccountHistoryError};
 use error_chain::ChainedError;
 use jsonrpc_core::{
     futures::{
@@ -182,6 +181,10 @@ pub enum ManagementCommand {
         OneshotSender<BoxFuture<AccountData, mullvad_rpc::Error>>,
         AccountToken,
     ),
+    /// Request account history
+    GetAccountHistory(OneshotSender<Vec<AccountToken>>),
+    /// Request account history
+    RemoveAccountFromHistory(OneshotSender<()>, AccountToken),
     /// Get the list of countries and cities where there are relays.
     GetRelayLocations(OneshotSender<RelayList>),
     /// Trigger an asynchronous relay list update. This returns before the relay list is actually
@@ -233,14 +236,11 @@ pub struct ManagementInterfaceServer {
 }
 
 impl ManagementInterfaceServer {
-    pub fn start<T>(
-        tunnel_tx: IntoSender<ManagementCommand, T>,
-        cache_dir: PathBuf,
-    ) -> talpid_ipc::Result<Self>
+    pub fn start<T>(tunnel_tx: IntoSender<ManagementCommand, T>) -> talpid_ipc::Result<Self>
     where
         T: From<ManagementCommand> + 'static + Send,
     {
-        let rpc = ManagementInterface::new(tunnel_tx, cache_dir);
+        let rpc = ManagementInterface::new(tunnel_tx);
         let subscriptions = rpc.subscriptions.clone();
 
         let mut io = PubSubHandler::default();
@@ -309,15 +309,13 @@ impl EventBroadcaster {
 struct ManagementInterface<T: From<ManagementCommand> + 'static + Send> {
     subscriptions: Arc<ActiveSubscriptions>,
     tx: Mutex<IntoSender<ManagementCommand, T>>,
-    cache_dir: PathBuf,
 }
 
 impl<T: From<ManagementCommand> + 'static + Send> ManagementInterface<T> {
-    pub fn new(tx: IntoSender<ManagementCommand, T>, cache_dir: PathBuf) -> Self {
+    pub fn new(tx: IntoSender<ManagementCommand, T>) -> Self {
         ManagementInterface {
             subscriptions: Default::default(),
             tx: Mutex::new(tx),
-            cache_dir,
         }
     }
 
@@ -381,12 +379,6 @@ impl<T: From<ManagementCommand> + 'static + Send> ManagementInterface<T> {
             _ => Error::internal_error(),
         }
     }
-
-    fn load_history(&self) -> Result<AccountHistory, AccountHistoryError> {
-        let mut account_history = AccountHistory::new(&self.cache_dir);
-        account_history.load()?;
-        Ok(account_history)
-    }
 }
 
 impl<T: From<ManagementCommand> + 'static + Send> ManagementInterfaceApi
@@ -440,18 +432,6 @@ impl<T: From<ManagementCommand> + 'static + Send> ManagementInterfaceApi
         let future = self
             .send_command_to_daemon(ManagementCommand::SetAccount(tx, account_token.clone()))
             .and_then(|_| rx.map_err(|_| Error::internal_error()));
-
-        if let Some(new_account_token) = account_token {
-            if let Err(e) = self.load_history().and_then(|mut account_history| {
-                account_history.add_account_token(new_account_token)
-            }) {
-                log::error!(
-                    "Unable to add an account into the account history: {}",
-                    e.display_chain()
-                );
-            }
-        }
-
         Box::new(future)
     }
 
@@ -558,14 +538,11 @@ impl<T: From<ManagementCommand> + 'static + Send> ManagementInterfaceApi
 
     fn get_account_history(&self, _: Self::Metadata) -> BoxFuture<Vec<AccountToken>, Error> {
         log::debug!("get_account_history");
-        Box::new(future::result(
-            self.load_history()
-                .map(|history| history.get_accounts().to_vec())
-                .map_err(|error| {
-                    log::error!("Unable to get account history: {}", error.display_chain());
-                    Error::internal_error()
-                }),
-        ))
+        let (tx, rx) = sync::oneshot::channel();
+        let future = self
+            .send_command_to_daemon(ManagementCommand::GetAccountHistory(tx))
+            .and_then(|_| rx.map_err(|_| Error::internal_error()));
+        Box::new(future)
     }
 
     fn remove_account_from_history(
@@ -574,17 +551,14 @@ impl<T: From<ManagementCommand> + 'static + Send> ManagementInterfaceApi
         account_token: AccountToken,
     ) -> BoxFuture<(), Error> {
         log::debug!("remove_account_from_history");
-        Box::new(future::result(
-            self.load_history()
-                .and_then(|mut history| history.remove_account_token(&account_token))
-                .map_err(|error| {
-                    log::error!(
-                        "Unable to remove account from history: {}",
-                        error.display_chain()
-                    );
-                    Error::internal_error()
-                }),
-        ))
+        let (tx, rx) = sync::oneshot::channel();
+        let future = self
+            .send_command_to_daemon(ManagementCommand::RemoveAccountFromHistory(
+                tx,
+                account_token,
+            ))
+            .and_then(|_| rx.map_err(|_| Error::internal_error()));
+        Box::new(future)
     }
 
     fn set_openvpn_mssfix(&self, _: Self::Metadata, mssfix: Option<u16>) -> BoxFuture<(), Error> {