Clean up error handling

3 files changed, 133 insertions, 144 deletions

diff --git a/mullvad-daemon/src/access_method.rs b/mullvad-daemon/src/access_method.rs
index 97dfa815f5..7d9d3dba95 100644
--- a/mullvad-daemon/src/access_method.rs
+++ b/mullvad-daemon/src/access_method.rs
@@ -1,8 +1,9 @@
 use crate::{
-    api,
+    api::{self, AccessModeSelectorHandle},
     settings::{self, MadeChanges},
     Daemon, EventListener,
 };
+use mullvad_api::rest::{self, MullvadRestHandle};
 use mullvad_types::{
     access_method::{self, AccessMethod, AccessMethodSetting},
     settings::Settings,
@@ -26,6 +27,8 @@ pub enum Error {
     /// [`AccessMethodSetting`]s & [`ApiConnectionMode`]s.
     #[error(display = "Error occured when handling connection settings & details")]
     ConnectionMode(#[error(source)] api::Error),
+    #[error(display = "API endpoint rotation failed")]
+    RestError(#[error(source)] rest::Error),
     /// Access methods settings error
     #[error(display = "Settings error")]
     Settings(#[error(source)] settings::Error),
@@ -202,7 +205,7 @@ where
             tokio::spawn(async move {
                 match handle.update_access_methods(new_access_methods).await {
                     Ok(_) => (),
-                    Err(crate::api::Error::NoAccessMethods) => {
+                    Err(api::Error::NoAccessMethods) | Err(_) => {
                         // `access_methods` was empty! This implies that the user
                         // disabled all access methods. If we ever get into this
                         // state, we should default to using the direct access
@@ -225,3 +228,92 @@ where
         }
     }
 }
+
+/// Try to reach the Mullvad API using a specific access method, returning
+/// an [`Error`] in the case where the test fails to reach the API.
+///
+/// Ephemerally sets a new access method (associated with `access_method`)
+/// to be used for subsequent API calls, before performing an API call and
+/// switching back to the previously active access method. The previous
+/// access method is *always* reset.
+pub async fn test_access_method(
+    new_access_method: AccessMethodSetting,
+    access_mode_selector: AccessModeSelectorHandle,
+    rest_handle: MullvadRestHandle,
+) -> Result<bool, Error> {
+    // Setup test
+    let previous_access_method = access_mode_selector
+        .get_access_method()
+        .await
+        .map_err(Error::ConnectionMode)?;
+
+    let method_under_test = new_access_method.clone();
+    access_mode_selector
+        .set_access_method(new_access_method)
+        .await
+        .map_err(Error::ConnectionMode)?;
+
+    // We need to perform a rotation of API endpoint after a set action
+    let rotation_handle = rest_handle.clone();
+    rotation_handle
+        .service()
+        .next_api_endpoint()
+        .await
+        .map_err(|err| {
+            log::error!("Failed to rotate API endpoint: {err}");
+            Error::RestError(err)
+        })?;
+
+    // Set up the reset
+    //
+    // In case the API call fails, the next API endpoint will
+    // automatically be selected, which means that we need to set up
+    // with the previous API endpoint beforehand.
+    access_mode_selector
+        .set_access_method(previous_access_method)
+        .await
+        .map_err(|err| {
+            log::error!(
+                "Could not reset to previous access
+            method after API reachability test was carried out. This should only
+            happen if the previous access method was removed in the meantime."
+            );
+            Error::ConnectionMode(err)
+        })?;
+
+    // Perform test
+    //
+    // Send a HEAD request to some Mullvad API endpoint. We issue a HEAD
+    // request because we are *only* concerned with if we get a reply from
+    // the API, and not with the actual data that the endpoint returns.
+    let result = mullvad_api::ApiProxy::new(rest_handle)
+        .api_addrs_available()
+        .await
+        .map_err(Error::RestError)?;
+
+    // We need to perform a rotation of API endpoint after a set action
+    // Note that this will be done automatically if the API call fails,
+    // so it only has to be done if the call succeeded ..
+    if result {
+        rotation_handle
+            .service()
+            .next_api_endpoint()
+            .await
+            .map_err(|err| {
+                log::error!("Failed to rotate API endpoint: {err}");
+                Error::RestError(err)
+            })?;
+    }
+
+    log::info!(
+        "The result of testing {method:?} is {result}",
+        method = method_under_test.access_method,
+        result = if result {
+            "success".to_string()
+        } else {
+            "failed".to_string()
+        }
+    );
+
+    Ok(result)
+}
diff --git a/mullvad-daemon/src/api.rs b/mullvad-daemon/src/api.rs
index 2814a1982a..2da307ff5f 100644
--- a/mullvad-daemon/src/api.rs
+++ b/mullvad-daemon/src/api.rs
@@ -35,15 +35,8 @@ pub enum Message {
 
 #[derive(err_derive::Error, Debug)]
 pub enum Error {
-    /// Oddly specific.
-    #[error(display = "Very Generic error.")]
-    Generic,
-    #[error(display = "{}", _0)]
-    Actor(#[error(source)] ActorError),
-}
-
-#[derive(err_derive::Error, Debug)]
-pub enum ActorError {
+    #[error(display = "No access methods were provided.")]
+    NoAccessMethods,
     #[error(display = "AccessModeSelector is not receiving any messages.")]
     SendFailed(#[error(source)] mpsc::TrySendError<Message>),
     #[error(display = "AccessModeSelector is not receiving any messages.")]
@@ -52,6 +45,9 @@ pub enum ActorError {
     NotRunning(#[error(source)] oneshot::Canceled),
 }
 
+type ResponseTx<T> = oneshot::Sender<Result<T>>;
+type Result<T> = std::result::Result<T, Error>;
+
 /// A channel for sending [`Message`] commands to a running
 /// [`AccessModeSelector`].
 #[derive(Clone)]
@@ -64,8 +60,8 @@ impl AccessModeSelectorHandle {
         let (tx, rx) = oneshot::channel();
         self.cmd_tx
             .unbounded_send(make_cmd(tx))
-            .map_err(ActorError::SendFailed)?;
-        rx.await.map_err(ActorError::NotRunning)?
+            .map_err(Error::SendFailed)?;
+        rx.await.map_err(Error::NotRunning)?
     }
 
     pub async fn get_access_method(&self) -> Result<AccessMethodSetting> {
@@ -105,7 +101,7 @@ impl AccessModeSelectorHandle {
     ///
     /// Practically converts the handle to a listener for when the
     /// currently valid connection modes changes.
-    pub fn as_stream(self) -> impl Stream<Item = ApiConnectionMode> {
+    pub fn into_stream(self) -> impl Stream<Item = ApiConnectionMode> {
         unfold(self, |handle| async move {
             match handle.next().await {
                 Ok(connection_mode) => Some((connection_mode, handle)),
@@ -146,7 +142,7 @@ impl AccessModeSelector {
 
         let connection_modes = match ConnectionModesIterator::new(connection_modes) {
             Ok(provider) => provider,
-            Err(Error::NoAccessMethods) => {
+            Err(Error::NoAccessMethods) | Err(_) => {
                 // No settings seem to have been found. Default to using the the
                 // direct access method.
                 let default = mullvad_types::access_method::Settings::direct();
@@ -187,9 +183,8 @@ impl AccessModeSelector {
     }
 
     fn reply<T>(&self, tx: ResponseTx<T>, value: T) -> Result<()> {
-        Ok(tx
-            .send(Ok(value))
-            .map_err(|_| ActorError::OneshotSendFailed)?)
+        tx.send(Ok(value)).map_err(|_| Error::OneshotSendFailed)?;
+        Ok(())
     }
 
     fn on_get_access_method(&mut self, tx: ResponseTx<AccessMethodSetting>) -> Result<()> {
@@ -240,10 +235,7 @@ impl AccessModeSelector {
             .unwrap_or(AccessMethod::from(BuiltInAccessMethod::Direct));
 
         let connection_mode = self.from(access_method);
-        log::info!(
-            "New API connection mode selected: {connection_mode}",
-            connection_mode = connection_mode
-        );
+        log::info!("New API connection mode selected: {connection_mode}");
         connection_mode
     }
 
@@ -252,12 +244,12 @@ impl AccessModeSelector {
         tx: ResponseTx<()>,
         values: Vec<AccessMethodSetting>,
     ) -> Result<()> {
-        self.update_access_methods(values);
+        self.update_access_methods(values)?;
         self.reply(tx, ())
     }
 
-    fn update_access_methods(&mut self, values: Vec<AccessMethodSetting>) {
-        self.connection_modes.update_access_methods(values);
+    fn update_access_methods(&mut self, values: Vec<AccessMethodSetting>) -> Result<()> {
+        self.connection_modes.update_access_methods(values)
     }
 
     /// Ad-hoc version of [`std::convert::From::from`], but since some
@@ -302,9 +294,6 @@ impl AccessModeSelector {
     }
 }
 
-type ResponseTx<T> = oneshot::Sender<Result<T>>;
-type Result<T> = std::result::Result<T, Error>;
-
 /// An iterator which will always produce an [`AccessMethod`].
 ///
 /// Safety: It is always safe to [`unwrap`] after calling [`next`] on a
@@ -319,14 +308,10 @@ pub struct ConnectionModesIterator {
     current: AccessMethodSetting,
 }
 
-#[derive(err_derive::Error, Debug)]
-pub enum Error {
-    #[error(display = "No access methods were provided.")]
-    NoAccessMethods,
-}
-
 impl ConnectionModesIterator {
-    pub fn new(access_methods: Vec<AccessMethodSetting>) -> Result<ConnectionModesIterator, Error> {
+    pub fn new(
+        access_methods: Vec<AccessMethodSetting>,
+    ) -> std::result::Result<ConnectionModesIterator, Error> {
         let mut iterator = Self::new_iterator(access_methods)?;
         Ok(Self {
             next: None,
@@ -344,7 +329,7 @@ impl ConnectionModesIterator {
     pub fn update_access_methods(
         &mut self,
         access_methods: Vec<AccessMethodSetting>,
-    ) -> Result<(), Error> {
+    ) -> std::result::Result<(), Error> {
         self.available_modes = Self::new_iterator(access_methods)?;
         Ok(())
     }
@@ -355,7 +340,7 @@ impl ConnectionModesIterator {
     /// returned.
     fn new_iterator(
         access_methods: Vec<AccessMethodSetting>,
-    ) -> Result<Box<dyn Iterator<Item = AccessMethodSetting> + Send>, Error> {
+    ) -> std::result::Result<Box<dyn Iterator<Item = AccessMethodSetting> + Send>, Error> {
         if access_methods.is_empty() {
             Err(Error::NoAccessMethods)
         } else {
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 9a2121260b..0288f9d8c4 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -693,7 +693,7 @@ where
 
         let api_handle = api_runtime
             .mullvad_rest_handle(
-                Box::pin(connection_modes_handler.clone().as_stream()),
+                Box::pin(connection_modes_handler.clone().into_stream()),
                 endpoint_updater.callback(),
             )
             .await;
@@ -1135,7 +1135,7 @@ where
             UpdateApiAccessMethod(tx, method) => self.on_update_api_access_method(tx, method).await,
             GetCurrentAccessMethod(tx) => self.on_get_current_api_access_method(tx),
             SetApiAccessMethod(tx, method) => self.on_set_api_access_method(tx, method).await,
-            TestApiAccessMethod(tx, method) => self.on_test_api_access_method(tx, method).await,
+            TestApiAccessMethod(tx, method) => self.on_test_api_access_method(tx, method),
             IsPerformingPostUpgrade(tx) => self.on_is_performing_post_upgrade(tx),
             GetCurrentVersion(tx) => self.on_get_current_version(tx),
             #[cfg(not(target_os = "android"))]
@@ -2369,14 +2369,7 @@ where
         });
     }
 
-    /// Try to reach the Mullvad API using a specific access method, returning
-    /// an [`Error`] in the case where the test fails to reach the API.
-    ///
-    /// Ephemerally sets a new access method (associated with `access_method`)
-    /// to be used for subsequent API calls, before performing an API call and
-    /// switching back to the previously active access method. The previous
-    /// access method is *always* reset.
-    async fn on_test_api_access_method(
+    fn on_test_api_access_method(
         &mut self,
         tx: ResponseTx<bool, Error>,
         access_method: mullvad_types::access_method::Id,
@@ -2387,105 +2380,24 @@ where
         // access method.
         let api_handle = self.api_handle.clone();
         let handle = self.connection_modes_handler.clone();
-        let access_method = self.get_api_access_method(access_method);
-        // TODO(markus): Clean up this error handling
-        let new_access_method = if let Ok(access_method) = access_method {
-            access_method
-        } else {
-            Self::oneshot_send(
-                tx,
-                access_method
-                    .map(|_| false)
-                    .map_err(Error::AccessMethodError),
-                "on_test_api_access_method response",
-            );
-            return;
-        };
-
-        let fut = async move {
-            // Setup test
-            let previous_access_method = handle
-                .get_access_method()
-                .await
-                .map_err(Error::ApiConnectionModeError)
-                // TODO(markus): Do not unwrap!
-                .unwrap();
-
-            let x = new_access_method.clone();
-            handle.set_access_method(new_access_method)
-                .await
-                .map_err(Error::ApiConnectionModeError)
-                // TODO(markus): Do not unwrap!
-                .unwrap();
-
-            // We need to perform a rotation of API endpoint after a set action
-            let rotation_handle = api_handle.clone();
-            rotation_handle
-                .service()
-                .next_api_endpoint()
-                .await
-                .map_err(|err| {
-                    log::error!("Failed to rotate API endpoint: {err}");
-                    err
-                })
-                // TODO(markus): Error handling
-                .unwrap();
-
-            // Set up the reset
-            //
-            // In case the API call fails, the next API endpoint will
-            // automatically be selected, which means that we need to set up
-            // with the previous API endpoint beforehand.
-            handle
-                .set_access_method(previous_access_method)
-                .await
-                .map_err(|err| {
-                    log::error!(
-                        "Could not reset to previous access
-            method after API reachability test was carried out. This should only
-            happen if the previous access method was removed in the meantime."
-                    );
-                    err
-                })
-                // TODO(markus): Do not unwrap!
-                .unwrap();
-
-            // Perform test
-            //
-            // Send a HEAD request to some Mullvad API endpoint. We issue a HEAD
-            // request because we are *only* concerned with if we get a reply from
-            // the API, and not with the actual data that the endpoint returns.
-            let result = mullvad_api::ApiProxy::new(api_handle)
-                .api_addrs_available()
-                .await
-                .map_err(Error::RestError);
+        let access_method_lookup = self
+            .get_api_access_method(access_method)
+            .map_err(Error::AccessMethodError);
 
-            // We need to perform a rotation of API endpoint after a set action
-            // Note that this will be done automatically if the API call fails,
-            // so it only has to be done if the call succeeded ..
-            if result.as_ref().is_ok_and(|&succeeded| succeeded) {
-                rotation_handle
-                .service()
-                .next_api_endpoint()
-                .await
-                .map_err(|err| {
-                    log::error!("Failed to rotate API endpoint: {err}");
-                    err
-                })
-                // TODO(markus): Error handling
-                .unwrap();
+        match access_method_lookup {
+            Ok(access_method) => {
+                tokio::spawn(async move {
+                    let result =
+                        access_method::test_access_method(access_method, handle, api_handle)
+                            .await
+                            .map_err(Error::AccessMethodError);
+                    Self::oneshot_send(tx, result, "on_test_api_access_method response");
+                });
             }
-
-            log::info!(
-                "The result of testing {method:?} is {result:?}",
-                method = x.access_method,
-                result = result
-            );
-
-            Self::oneshot_send(tx, result, "on_test_api_access_method response");
-        };
-
-        tokio::spawn(fut);
+            Err(err) => {
+                Self::oneshot_send(tx, Err(err), "on_test_api_access_method response");
+            }
+        }
     }
 
     fn on_get_settings(&self, tx: oneshot::Sender<Settings>) {