Add `mullvad proxy use`

Allow for settings a specific Access Method to use

7 files changed, 169 insertions, 30 deletions

diff --git a/mullvad-cli/src/cmds/proxy.rs b/mullvad-cli/src/cmds/proxy.rs
index 48ef302c19..c876a25d2e 100644
--- a/mullvad-cli/src/cmds/proxy.rs
+++ b/mullvad-cli/src/cmds/proxy.rs
@@ -21,9 +21,11 @@ pub enum Proxy {
     /// Remove an API proxy
     Remove(RemoveCustomCommands),
     /// Enable an API proxy
-    Enable(ToggleCustomCommands),
+    Enable(SelectItem),
     /// Disable an API proxy
-    Disable(ToggleCustomCommands),
+    Disable(SelectItem),
+    /// Force the use of a specific API proxy.
+    Use(SelectItem),
 }
 
 impl Proxy {
@@ -57,6 +59,10 @@ impl Proxy {
                 let enabled = false;
                 Self::toggle(index, enabled).await?;
             }
+            Proxy::Use(cmd) => {
+                let index = Self::zero_to_one_based_index(cmd.index)?;
+                Self::set(index).await?;
+            }
         };
         Ok(())
     }
@@ -193,6 +199,25 @@ impl Proxy {
         Ok(())
     }
 
+    /// Force the use of a specific access method when trying to reach the
+    /// Mullvad API. If this method fails, the daemon will resume the automatic
+    /// roll-over behavior (which is the default).
+    async fn set(index: usize) -> Result<()> {
+        let mut rpc = MullvadProxyClient::new().await?;
+        let access_method = rpc
+            .get_api_access_methods()
+            .await?
+            .get(index)
+            .ok_or(anyhow!(format!(
+                "Access method {} does not exist",
+                index + 1
+            )))?
+            .clone();
+
+        rpc.set_access_method(access_method).await?;
+        Ok(())
+    }
+
     fn zero_to_one_based_index(index: usize) -> Result<usize> {
         index
             .checked_sub(1)
@@ -224,9 +249,11 @@ pub enum AddCustomCommands {
     },
 }
 
+/// A minimal wrapper type allowing the user to supply a list index to some
+/// Access Method.
 #[derive(Args, Debug, Clone)]
-pub struct ToggleCustomCommands {
-    /// Which access method to enable/disable
+pub struct SelectItem {
+    /// Which access method to pick
     index: usize,
 }
 
diff --git a/mullvad-daemon/src/access_methods.rs b/mullvad-daemon/src/access_methods.rs
index 79afdae054..a0c168c026 100644
--- a/mullvad-daemon/src/access_methods.rs
+++ b/mullvad-daemon/src/access_methods.rs
@@ -86,22 +86,29 @@ where
             .map_err(Error::Settings)
     }
 
+    pub async fn set_api_access_method(
+        &mut self,
+        access_method: AccessMethod,
+    ) -> Result<(), Error> {
+        {
+            let mut connection_modes = self.connection_modes.lock().unwrap();
+            connection_modes.set_access_method(access_method);
+        }
+        // Force a rotation of Acces Methods.
+        let _ = self.api_handle.service().next_api_endpoint();
+        Ok(())
+    }
+
     /// If settings were changed due to an update, notify all listeners.
     fn notify_on_change(&mut self, settings_changed: MadeChanges) {
         if settings_changed {
             self.event_listener
                 .notify_settings(self.settings.to_settings());
-        };
 
-        // TODO: Could this be replaced by message passing? Yes plz.
-        let mut connection_modes = self.connection_modes.lock().unwrap();
-        *connection_modes = self
-            .settings
-            .api_access_methods
-            .api_access_methods
-            .clone()
-            .into_iter()
-            .map(|x| (x, 1))
-            .collect();
+            // TODO: Could this be replaced by message passing? Yes plz.
+            let mut connection_modes = self.connection_modes.lock().unwrap();
+            connection_modes
+                .update_access_methods(self.settings.api_access_methods.api_access_methods.clone())
+        };
     }
 }
diff --git a/mullvad-daemon/src/api.rs b/mullvad-daemon/src/api.rs
index e8568ac684..f96dea9484 100644
--- a/mullvad-daemon/src/api.rs
+++ b/mullvad-daemon/src/api.rs
@@ -44,6 +44,58 @@ pub struct ApiConnectionModeProvider {
     relay_selector: RelaySelector,
     retry_attempt: u32,
     current_task: Option<Pin<Box<dyn Future<Output = ApiConnectionMode> + Send>>>,
+    connection_modes: Arc<Mutex<ConnectionModesIterator>>,
+}
+
+/// An iterator which will always produce an [`AccessMethod`].
+///
+/// Safety: It is always safe to [`unwrap`] after calling [`next`] on a
+/// [`std::iter::Cycle`], so thereby it is safe to always call [`unwrap`] on a
+/// [`ConnectionModesIterator`]
+///
+/// [`unwrap`]: Option::unwrap
+/// [`next`]: std::iter::Iterator::next
+pub struct ConnectionModesIterator {
+    available_modes: Box<dyn Iterator<Item = AccessMethod> + Send>,
+    next: Option<AccessMethod>,
+}
+
+impl ConnectionModesIterator {
+    pub fn new(modes: Vec<AccessMethod>) -> ConnectionModesIterator {
+        Self {
+            next: None,
+            available_modes: Self::get_filtered_access_methods(modes),
+        }
+    }
+
+    /// Set the next [`AccessMethod`] to be returned from this iterator.
+    pub fn set_access_method(&mut self, next: AccessMethod) {
+        self.next = Some(next);
+    }
+    /// Update the collection of [`AccessMethod`] which this iterator will
+    /// return.
+    pub fn update_access_methods(&mut self, access_methods: Vec<AccessMethod>) {
+        self.available_modes = Self::get_filtered_access_methods(access_methods);
+    }
+
+    fn get_filtered_access_methods(
+        modes: Vec<AccessMethod>,
+    ) -> Box<dyn Iterator<Item = AccessMethod> + Send> {
+        Box::new(
+            modes
+                .into_iter()
+                .filter(|access_method| access_method.enabled())
+                .cycle(),
+        )
+    }
+}
+
+impl Iterator for ConnectionModesIterator {
+    type Item = AccessMethod;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next.take().or_else(|| self.available_modes.next())
+    }
 }
 
 impl Stream for ApiConnectionModeProvider {
@@ -83,23 +135,22 @@ impl Stream for ApiConnectionModeProvider {
 }
 
 impl ApiConnectionModeProvider {
-    pub(crate) fn new(cache_dir: PathBuf, relay_selector: RelaySelector) -> Self {
+    pub(crate) fn new(
+        cache_dir: PathBuf,
+        relay_selector: RelaySelector,
+        connection_modes: Vec<AccessMethod>,
+    ) -> Self {
         Self {
             cache_dir,
             relay_selector,
             retry_attempt: 0,
             current_task: None,
+            connection_modes: Arc::new(Mutex::new(ConnectionModesIterator::new(connection_modes))),
         }
     }
 
-    fn should_use_bridge(retry_attempt: u32) -> bool {
-        retry_attempt % 3 > 0
-    }
-
-    fn should_use_direct(retry_attempt: u32) -> bool {
-        // TODO: Change back before comitting!
-        false
-        // !Self::should_use_bridge(retry_attempt)
+    pub(crate) fn handle(&self) -> Arc<Mutex<ConnectionModesIterator>> {
+        self.connection_modes.clone()
     }
 
     /// Return a new connection mode to be used for the API connection.
@@ -112,10 +163,9 @@ impl ApiConnectionModeProvider {
         log::debug!("Rotating Access mode!");
         let access_method = {
             let mut access_methods_picker = self.connection_modes.lock().unwrap();
-            // Rotate through the cycle of access methods.
-            // Safety: It is always safe to unwrap after calling `next` on a [`std::iter::Cycle`]
             access_methods_picker.next().unwrap()
         };
+
         let connection_mode = self.from(&access_method);
         log::info!("New API connection mode selected: {}", connection_mode);
         connection_mode
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 8427b1ea6b..3c822fdf4e 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -65,7 +65,7 @@ use std::{
     mem,
     path::PathBuf,
     pin::Pin,
-    sync::{Arc, Weak},
+    sync::{Arc, Mutex, Weak},
     time::Duration,
 };
 #[cfg(any(target_os = "linux", windows))]
@@ -273,6 +273,8 @@ pub enum DaemonCommand {
     ReplaceApiAccessMethod(ResponseTx<(), Error>, ApiAccessMethodReplace),
     /// Toggle the status of an API access method
     ToggleApiAccessMethod(ResponseTx<(), Error>, ApiAccessMethodToggle),
+    /// Set the API access method to use
+    SetApiAccessMethod(ResponseTx<(), Error>, AccessMethod),
     /// Get information about the currently running and latest app versions
     GetVersionInfo(oneshot::Sender<Option<AppVersionInfo>>),
     /// Return whether the daemon is performing post-upgrade tasks
@@ -572,6 +574,7 @@ pub struct Daemon<L: EventListener> {
     account_history: account_history::AccountHistory,
     device_checker: device::TunnelStateChangeHandler,
     account_manager: device::AccountManagerHandle,
+    connection_modes: Arc<Mutex<api::ConnectionModesIterator>>,
     api_runtime: mullvad_api::Runtime,
     api_handle: mullvad_api::rest::MullvadRestHandle,
     version_updater_handle: version_check::VersionUpdaterHandle,
@@ -634,8 +637,15 @@ where
         let initial_selector_config = new_selector_config(&settings);
         let relay_selector = RelaySelector::new(initial_selector_config, &resource_dir, &cache_dir);
 
-        let proxy_provider =
-            api::ApiConnectionModeProvider::new(cache_dir.clone(), relay_selector.clone());
+        // TODO: Should ApiConnectionModeProvider be an Actor instead of sharing a datastructure-behind-locks with the daemon with the daemon?
+        let proxy_provider = api::ApiConnectionModeProvider::new(
+            cache_dir.clone(),
+            relay_selector.clone(),
+            settings.api_access_methods.api_access_methods.clone(),
+        );
+
+        let connection_modes = proxy_provider.handle();
+
         let api_handle = api_runtime
             .mullvad_rest_handle(proxy_provider, endpoint_updater.callback())
             .await;
@@ -772,6 +782,7 @@ where
             account_history,
             device_checker: device::TunnelStateChangeHandler::new(account_manager.clone()),
             account_manager,
+            connection_modes,
             api_runtime,
             api_handle,
             version_updater_handle,
@@ -1055,6 +1066,7 @@ where
                 self.on_replace_api_access_method(tx, method).await
             }
             ToggleApiAccessMethod(tx, method) => self.on_toggle_api_access_method(tx, method).await,
+            SetApiAccessMethod(tx, method) => self.on_set_api_access_method(tx, method).await,
             IsPerformingPostUpgrade(tx) => self.on_is_performing_post_upgrade(tx),
             GetCurrentVersion(tx) => self.on_get_current_version(tx),
             #[cfg(not(target_os = "android"))]
@@ -2275,7 +2287,15 @@ where
             .toggle_api_access_method(method)
             .await
             .map_err(Error::AccessMethodError);
-        Self::oneshot_send(tx, result, "edit_api_access_method response");
+        Self::oneshot_send(tx, result, "toggle_api_access_method response");
+    }
+
+    async fn on_set_api_access_method(&mut self, tx: ResponseTx<(), Error>, method: AccessMethod) {
+        let result = self
+            .set_api_access_method(method)
+            .await
+            .map_err(Error::AccessMethodError);
+        Self::oneshot_send(tx, result, "set_api_access_method response");
     }
 
     fn on_get_settings(&self, tx: oneshot::Sender<Settings>) {
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 91be9352ec..2b6481207a 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -639,6 +639,21 @@ impl ManagementService for ManagementServiceImpl {
             .map_err(map_daemon_error)
     }
 
+    async fn set_api_access_method(
+        &self,
+        request: Request<types::ApiAccessMethod>,
+    ) -> ServiceResult<()> {
+        log::debug!("set_api_access_method");
+        let access_method =
+            mullvad_types::api_access_method::AccessMethod::try_from(request.into_inner())?;
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::SetApiAccessMethod(tx, access_method))?;
+        self.wait_for_result(rx)
+            .await?
+            .map(Response::new)
+            .map_err(map_daemon_error)
+    }
+
     // Split tunneling
     //
 
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index ffc3f12120..a062e0b642 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -87,6 +87,9 @@ service ManagementService {
   rpc ToggleApiAccessMethod(ApiAccessMethodToggle) returns (google.protobuf.Empty) {
     // Can I return something useful here instead of Empty?
   }
+  rpc SetApiAccessMethod(ApiAccessMethod) returns (google.protobuf.Empty) {
+    // Can I return something useful here instead of Empty?
+  }
 
   // Split tunneling (Linux)
   rpc GetSplitTunnelProcesses(google.protobuf.Empty) returns (stream google.protobuf.Int32Value) {}
diff --git a/mullvad-management-interface/src/client.rs b/mullvad-management-interface/src/client.rs
index fb2eae28c3..900775b10a 100644
--- a/mullvad-management-interface/src/client.rs
+++ b/mullvad-management-interface/src/client.rs
@@ -513,6 +513,23 @@ impl MullvadProxyClient {
             .map(drop)
     }
 
+    /// Set the [`AccessMethod`] which [`ApiConnectionModeProvider`] should
+    /// pick.
+    ///
+    /// - `access_method`: If `Some(access_method)`, [`ApiConnectionModeProvider`] will skip
+    ///     ahead and return `access_method` when asked for a new access method.
+    ///     If `None`, [`ApiConnectionModeProvider`] will pick the next access
+    ///     method "randomly"
+    ///
+    /// [`ApiConnectionModeProvider`]: mullvad_daemon::api::ApiConnectionModeProvider
+    pub async fn set_access_method(&mut self, access_method: AccessMethod) -> Result<()> {
+        self.0
+            .set_api_access_method(types::ApiAccessMethod::from(access_method))
+            .await
+            .map_err(Error::Rpc)
+            .map(drop)
+    }
+
     #[cfg(target_os = "linux")]
     pub async fn get_split_tunnel_processes(&mut self) -> Result<Vec<i32>> {
         use futures::TryStreamExt;