Merge branch 'improve-keygen-on-login'

3 files changed, 75 insertions, 40 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index e8edcbcafa..63dca09237 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -82,6 +82,9 @@ type BoxFuture<T, E> = Box<dyn Future<Item = T, Error = E> + Send>;
 
 const TUNNEL_STATE_MACHINE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
 
+/// Timeout for first WireGuard key pushing
+const FIRST_KEY_PUSH_TIMEOUT: Duration = Duration::from_secs(5);
+
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
 pub enum Error {
@@ -1643,7 +1646,10 @@ where
                 .unwrap_or(true)
             {
                 log::info!("Automatically generating new wireguard key for account");
-                if let Err(e) = self.wireguard_key_manager.generate_key_async(account) {
+                if let Err(e) = self
+                    .wireguard_key_manager
+                    .generate_key_async(account, Some(FIRST_KEY_PUSH_TIMEOUT))
+                {
                     log::error!(
                         "{}",
                         e.display_chain_with_msg("Failed to start generating wireguard key")
diff --git a/mullvad-daemon/src/wireguard.rs b/mullvad-daemon/src/wireguard.rs
index 3e928827be..1b5913699f 100644
--- a/mullvad-daemon/src/wireguard.rs
+++ b/mullvad-daemon/src/wireguard.rs
@@ -112,7 +112,7 @@ impl KeyManager {
         self.reset();
         let private_key = PrivateKey::new_from_random();
 
-        self.run_future_sync(self.push_future_generator(account, private_key)())
+        self.run_future_sync(self.push_future_generator(account, private_key, None)())
             .map_err(Self::map_rpc_error)
     }
 
@@ -168,38 +168,60 @@ impl KeyManager {
 
 
     /// Generate a new private key asynchronously. The new keys will be sent to the daemon channel.
-    pub fn generate_key_async(&mut self, account: AccountToken) -> Result<()> {
+    pub fn generate_key_async(
+        &mut self,
+        account: AccountToken,
+        timeout: Option<Duration>,
+    ) -> Result<()> {
         self.reset();
         let private_key = PrivateKey::new_from_random();
-        let future_generator = self.push_future_generator(account.clone(), private_key);
+
+        let error_tx = self.daemon_tx.clone();
+        let error_account = account.clone();
+
+        let mut inner_future_generator =
+            self.push_future_generator(account.clone(), private_key, timeout);
+
+        let future_generator = move || {
+            let fut = inner_future_generator();
+            let error_tx = error_tx.clone();
+            let error_account = error_account.clone();
+            fut.map_err(move |err| {
+                let should_retry = match &err {
+                    RestError::ApiError(_status, code)
+                        if code == mullvad_rpc::KEY_LIMIT_REACHED =>
+                    {
+                        false
+                    }
+                    _ => true,
+                };
+
+                let _ = error_tx.send(InternalDaemonEvent::WgKeyEvent((
+                    error_account,
+                    Err(Self::map_rpc_error(err)),
+                )));
+
+                should_retry
+            })
+        };
+
 
         let retry_strategy = ExponentialBackoff::from_millis(300)
             .max_delay(Duration::from_secs(60 * 60))
             .map(jitter);
 
-        let should_retry = |err: &RestError| -> bool {
-            match err {
-                RestError::ApiError(_status, code) if code == mullvad_rpc::KEY_LIMIT_REACHED => {
-                    false
-                }
-                _ => true,
-            }
-        };
 
         let upload_future =
-            RetryIf::spawn(retry_strategy, future_generator, should_retry).map_err(move |err| {
+            RetryIf::spawn(retry_strategy, future_generator, |should_retry: &bool| {
+                *should_retry
+            })
+            .map_err(move |err| {
                 match err {
                     // This should really be unreachable, since the retry strategy is infinite.
-                    tokio_retry::Error::OperationError(e) => {
-                        log::error!(
-                            "{}",
-                            e.display_chain_with_msg("Failed to generate wireguard key:")
-                        );
-                        Self::map_rpc_error(e)
-                    }
+                    tokio_retry::Error::OperationError(_) => {}
                     tokio_retry::Error::TimerError(timer_error) => {
                         log::error!("Tokio timer error {}", timer_error);
-                        Error::ExectuionError
+                        ()
                     }
                 }
             });
@@ -215,9 +237,7 @@ impl KeyManager {
                         Ok(wireguard_data),
                     )));
                 }
-                Err(CancelErr::Inner(e)) => {
-                    let _ = daemon_tx.send(InternalDaemonEvent::WgKeyEvent((account, Err(e))));
-                }
+                Err(CancelErr::Inner(_)) => {}
                 Err(CancelErr::Cancelled) => {
                     log::error!("Key generation cancelled");
                 }
@@ -240,6 +260,7 @@ impl KeyManager {
         &self,
         account: AccountToken,
         private_key: PrivateKey,
+        timeout: Option<Duration>,
     ) -> Box<dyn FnMut() -> Box<dyn Future<Item = WireguardData, Error = RestError> + Send> + Send>
     {
         let mut rpc = mullvad_rpc::WireguardKeyProxy::new(self.http_handle.clone());
@@ -248,17 +269,19 @@ impl KeyManager {
         let push_future =
             move || -> Box<dyn Future<Item = WireguardData, Error = RestError> + Send> {
                 let key = private_key.clone();
-                Box::new(rpc.push_wg_key(account.clone(), public_key.clone()).map(
-                    move |addresses| WireguardData {
-                        private_key: key,
-                        addresses,
-                        created: Utc::now(),
-                    },
-                ))
+                Box::new(
+                    rpc.push_wg_key(account.clone(), public_key.clone(), timeout)
+                        .map(move |addresses| WireguardData {
+                            private_key: key,
+                            addresses,
+                            created: Utc::now(),
+                        }),
+                )
             };
         Box::new(push_future)
     }
 
+
     fn replace_key_rpc(
         http_handle: MullvadRestHandle,
         account: AccountToken,
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index bbe7a1f7bd..aaea9aad4e 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -337,6 +337,7 @@ impl WireguardKeyProxy {
         &mut self,
         account_token: AccountToken,
         public_key: wireguard::PublicKey,
+        timeout: Option<std::time::Duration>,
     ) -> impl Future01<Item = mullvad_types::wireguard::AssociatedAddresses, Error = rest::Error>
     {
         #[derive(serde::Serialize)]
@@ -347,17 +348,22 @@ impl WireguardKeyProxy {
         let service = self.handle.service.clone();
         let body = PublishRequest { pubkey: public_key };
 
-        let request = rest::post_request_with_json(
-            &self.handle.factory,
-            service,
-            &"/v1/wireguard-keys",
-            &body,
-            Some(account_token),
-            StatusCode::CREATED,
-        );
+        let request = self.handle.factory.post_json(&"/v1/wireguard-keys", &body);
+
+        let future = async move {
+            let mut request = request?;
+            if let Some(timeout) = timeout {
+                request.set_timeout(timeout);
+            }
+            request.set_auth(Some(account_token))?;
+            let response = service.request(request).await?;
+            rest::parse_rest_response(response, StatusCode::CREATED).await
+        };
+
+
         self.handle
             .service
-            .compat_spawn(async move { rest::deserialize_body(request.await?).await })
+            .compat_spawn(async move { rest::deserialize_body(future.await?).await })
     }
 
     pub fn replace_wg_key(