Merge branch 'timer-fix'

2 files changed, 77 insertions, 93 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index b1c75d86a2..06f122fdf9 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -480,7 +480,7 @@ where
                     .get_tunnel_options()
                     .wireguard
                     .automatic_rotation
-                    .map(|hours| 60 * hours),
+                    .map(|hours| Duration::from_secs(60u64 * 60u64 * hours as u64)),
             );
         }
 
@@ -1378,7 +1378,7 @@ where
                         self.wireguard_key_manager.set_rotation_interval(
                             &mut self.account_history,
                             token,
-                            interval.map(|hours| 60 * hours),
+                            interval.map(|hours| Duration::from_secs(60u64 * 60u64 * hours as u64)),
                         );
                     }
 
@@ -1462,7 +1462,7 @@ where
                             .get_tunnel_options()
                             .wireguard
                             .automatic_rotation
-                            .map(|hours| 60 * hours),
+                            .map(|hours| Duration::from_secs(60u64 * 60u64 * hours as u64)),
                     );
 
                     Ok(keygen_event)
diff --git a/mullvad-daemon/src/wireguard.rs b/mullvad-daemon/src/wireguard.rs
index 4e29f410b9..7bd98cee52 100644
--- a/mullvad-daemon/src/wireguard.rs
+++ b/mullvad-daemon/src/wireguard.rs
@@ -8,7 +8,7 @@ use futures::{
 use jsonrpc_client_core::Error as JsonRpcError;
 use mullvad_types::account::AccountToken;
 pub use mullvad_types::wireguard::*;
-use std::{cmp, time::Duration};
+use std::time::Duration;
 pub use talpid_types::net::wireguard::{
     ConnectionConfig, PrivateKey, TunnelConfig, TunnelParameters,
 };
@@ -22,12 +22,13 @@ use tokio_timer;
 
 const TOO_MANY_KEYS_ERROR_CODE: i64 = -703;
 
-/// Default automatic key rotation (in minutes)
-const DEFAULT_AUTOMATIC_KEY_ROTATION: u32 = 7 * 24 * 60;
-/// How long to wait before reattempting to rotate keys on failure (secs)
-const AUTOMATIC_ROTATION_RETRY_DELAY: u64 = 5;
-/// Minimum interval used by automatic rotation (secs)
-const MINIMUM_ROTATION_INTERVAL: u64 = 5;
+/// Default automatic key rotation
+const DEFAULT_AUTOMATIC_KEY_ROTATION: Duration = Duration::from_secs(7 * 24 * 60 * 60);
+/// How long to wait before reattempting to rotate keys on failure
+const AUTOMATIC_ROTATION_RETRY_DELAY: Duration = Duration::from_secs(5);
+/// How often to check whether the key has expired.
+/// A short interval is used in case the computer is ever suspended.
+const KEY_CHECK_INTERVAL: Duration = Duration::from_secs(60);
 
 
 #[derive(err_derive::Error, Debug)]
@@ -53,8 +54,7 @@ pub struct KeyManager {
     current_job: Option<CancelHandle>,
 
     abort_scheduler_tx: Option<CancelHandle>,
-    // unit: minutes
-    auto_rotation_interval: u32,
+    auto_rotation_interval: Duration,
 }
 
 impl KeyManager {
@@ -69,7 +69,7 @@ impl KeyManager {
             tokio_remote,
             current_job: None,
             abort_scheduler_tx: None,
-            auto_rotation_interval: 0,
+            auto_rotation_interval: Duration::new(0, 0),
         }
     }
 
@@ -96,19 +96,19 @@ impl KeyManager {
         };
     }
 
-    /// Update automatic key rotation interval (given in minutes)
+    /// Update automatic key rotation interval
     /// Passing `None` for the interval will use the default value.
-    /// A value of `0` disables automatic key rotation.
+    /// A duration of `0` disables automatic key rotation.
     pub fn set_rotation_interval(
         &mut self,
         account_history: &mut AccountHistory,
         account_token: AccountToken,
-        auto_rotation_interval_mins: Option<u32>,
+        auto_rotation_interval: Option<Duration>,
     ) {
         log::debug!("set_rotation_interval");
 
         self.auto_rotation_interval =
-            auto_rotation_interval_mins.unwrap_or(DEFAULT_AUTOMATIC_KEY_ROTATION);
+            auto_rotation_interval.unwrap_or(DEFAULT_AUTOMATIC_KEY_ROTATION);
 
         self.reset_rotation(account_history, account_token);
     }
@@ -283,31 +283,26 @@ impl KeyManager {
             _ => Error::RpcError(err),
         }
     }
-    fn create_key_expiration_timer(
-        public_key: PublicKey,
-        rotation_interval_secs: u64,
-    ) -> impl Future<Item = (), Error = Error> + Send {
-        let last_update = public_key.created.clone();
-        let key_age = Duration::from_secs(
-            (Utc::now().signed_duration_since(last_update)).num_seconds() as u64,
-        );
 
-        let interval_duration = Duration::from_secs(rotation_interval_secs);
-        let remaining_time = interval_duration
-            .checked_sub(key_age)
-            .unwrap_or(Duration::from_secs(0));
-        let key_expiry = cmp::max(
-            Duration::from_secs(MINIMUM_ROTATION_INTERVAL),
-            remaining_time,
-        );
-
-        log::info!("Next key rotation (time left): {:?}", key_expiry);
-
-        tokio_timer::wheel()
-            .max_timeout(Duration::new(std::u64::MAX, 0))
-            .build()
-            .sleep(key_expiry)
-            .map_err(|e| Error::RotationScheduleError(e))
+    fn create_rotation_check(
+        key: PublicKey,
+        rotation_interval_secs: u64,
+    ) -> Box<dyn Future<Item = (), Error = Error> + Send> {
+        Box::new(
+            tokio_timer::wheel()
+                .build()
+                .sleep(KEY_CHECK_INTERVAL)
+                .map_err(|e| Error::RotationScheduleError(e))
+                .and_then(move |_| {
+                    let key_age =
+                        (Utc::now().signed_duration_since(key.created)).num_seconds() as u64;
+                    if key_age >= rotation_interval_secs {
+                        Box::new(futures::future::ok(()))
+                    } else {
+                        Self::create_rotation_check(key, rotation_interval_secs)
+                    }
+                }),
+        )
     }
 
     fn next_automatic_rotation(
@@ -316,50 +311,45 @@ impl KeyManager {
         public_key: PublicKey,
         rotation_interval_secs: u64,
         account_token: AccountToken,
-    ) -> impl Future<Item = PublicKey, Error = Error> + Send {
+    ) -> Box<dyn Future<Item = PublicKey, Error = Error> + Send> {
         let expiration_timer =
-            Self::create_key_expiration_timer(public_key.clone(), rotation_interval_secs);
-
-
+            Self::create_rotation_check(public_key.clone(), rotation_interval_secs);
         let account_token_copy = account_token.clone();
 
-        expiration_timer
-            .and_then(move |_| {
-                log::info!("Replacing WireGuard key");
+        Box::new(
+            expiration_timer
+                .and_then(move |_| {
+                    log::info!("Replacing WireGuard key");
 
-                let private_key = PrivateKey::new_from_random()
-                    .map_err(Error::GenerationError)
-                    .into_future();
-                private_key.and_then(move |private_key| {
-                    Self::replace_key_rpc(
-                        http_handle.clone(),
-                        account_token.clone(),
-                        public_key.clone(),
-                        private_key,
-                    )
-                    .map_err(Self::map_rpc_error)
+                    let private_key = PrivateKey::new_from_random()
+                        .map_err(Error::GenerationError)
+                        .into_future();
+                    private_key.and_then(move |private_key| {
+                        Self::replace_key_rpc(http_handle, account_token, public_key, private_key)
+                            .map_err(Self::map_rpc_error)
+                    })
                 })
-            })
-            .then(move |rpc_result| {
-                match rpc_result {
-                    Ok(data) => {
-                        // Update account data
-                        let _ = daemon_tx.unbounded_send(InternalDaemonEvent::WgKeyEvent((
-                            account_token_copy,
-                            Ok(data.clone()),
-                        )));
-                        Ok(data.get_public_key())
-                    }
-                    Err(Error::TooManyKeys) => {
-                        let _ = daemon_tx.unbounded_send(InternalDaemonEvent::WgKeyEvent((
-                            account_token_copy,
-                            Err(Error::TooManyKeys),
-                        )));
-                        Err(Error::TooManyKeys)
+                .then(move |rpc_result| {
+                    match rpc_result {
+                        Ok(data) => {
+                            // Update account data
+                            let _ = daemon_tx.unbounded_send(InternalDaemonEvent::WgKeyEvent((
+                                account_token_copy,
+                                Ok(data.clone()),
+                            )));
+                            Ok(data.get_public_key())
+                        }
+                        Err(Error::TooManyKeys) => {
+                            let _ = daemon_tx.unbounded_send(InternalDaemonEvent::WgKeyEvent((
+                                account_token_copy,
+                                Err(Error::TooManyKeys),
+                            )));
+                            Err(Error::TooManyKeys)
+                        }
+                        Err(unknown_err) => Err(unknown_err),
                     }
-                    Err(unknown_err) => Err(unknown_err),
-                }
-            })
+                }),
+        )
     }
 
     fn create_automatic_rotation(
@@ -381,35 +371,29 @@ impl KeyManager {
 
         let create_repeat_future = move |result: Result<PublicKey>| match result {
             Ok(next_public_key) => Self::create_automatic_rotation(
-                daemon_tx.clone(),
-                http_handle.clone(),
+                daemon_tx,
+                http_handle,
                 next_public_key,
                 rotation_interval_secs,
-                account_token.clone(),
+                account_token,
             ),
             Err(Error::TooManyKeys) => Box::new(futures::future::ok(())),
             Err(e) => {
                 log::error!(
                     "Key rotation failed: {}. Retrying in {} seconds",
                     e,
-                    AUTOMATIC_ROTATION_RETRY_DELAY,
+                    AUTOMATIC_ROTATION_RETRY_DELAY.as_secs(),
                 );
 
-                let next_public_key = public_key.clone();
-
-                let daemon_tx = daemon_tx.clone();
-                let http_handle = http_handle.clone();
-                let account_token = account_token.clone();
-
                 Box::new(
                     tokio_timer::wheel()
                         .build()
-                        .sleep(Duration::from_secs(AUTOMATIC_ROTATION_RETRY_DELAY))
+                        .sleep(AUTOMATIC_ROTATION_RETRY_DELAY)
                         .then(move |_| {
                             Self::create_automatic_rotation(
                                 daemon_tx,
                                 http_handle,
-                                next_public_key,
+                                public_key,
                                 rotation_interval_secs,
                                 account_token,
                             )
@@ -424,7 +408,7 @@ impl KeyManager {
     fn run_automatic_rotation(&mut self, account_token: AccountToken, public_key: PublicKey) {
         self.stop_automatic_rotation();
 
-        if let 0 = self.auto_rotation_interval {
+        if self.auto_rotation_interval == Duration::new(0, 0) {
             // disabled
             return;
         }
@@ -434,7 +418,7 @@ impl KeyManager {
             self.daemon_tx.clone(),
             self.http_handle.clone(),
             public_key,
-            60u64 * (self.auto_rotation_interval as u64),
+            self.auto_rotation_interval.as_secs(),
             account_token,
         );
         let (fut, cancel_handle) = Cancellable::new(fut);