Add management interface for relay override

5 files changed, 123 insertions, 2 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 7f86e7b364..98a8febe2a 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -47,7 +47,9 @@ use mullvad_types::{
     custom_list::CustomList,
     device::{Device, DeviceEvent, DeviceEventCause, DeviceId, DeviceState, RemoveDeviceEvent},
     location::GeoIpLocation,
-    relay_constraints::{BridgeSettings, BridgeState, ObfuscationSettings, RelaySettings},
+    relay_constraints::{
+        BridgeSettings, BridgeState, ObfuscationSettings, RelayOverride, RelaySettings,
+    },
     relay_list::RelayList,
     settings::{DnsOptions, Settings},
     states::{TargetState, TunnelState},
@@ -253,6 +255,10 @@ pub enum DaemonCommand {
     SetQuantumResistantTunnel(ResponseTx<(), settings::Error>, QuantumResistantState),
     /// Set DNS options or servers to use
     SetDnsOptions(ResponseTx<(), settings::Error>, DnsOptions),
+    /// Set override options to use for a given relay
+    SetRelayOverride(ResponseTx<(), settings::Error>, RelayOverride),
+    /// Remove all relay override options
+    ClearAllRelayOverrides(ResponseTx<(), settings::Error>),
     /// Toggle macOS network check leak
     /// Set MTU for wireguard tunnels
     SetWireguardMtu(ResponseTx<(), settings::Error>, Option<u16>),
@@ -1077,6 +1083,10 @@ where
                     .await
             }
             SetDnsOptions(tx, dns_servers) => self.on_set_dns_options(tx, dns_servers).await,
+            SetRelayOverride(tx, relay_override) => {
+                self.on_set_relay_override(tx, relay_override).await
+            }
+            ClearAllRelayOverrides(tx) => self.on_clear_all_relay_overrides(tx).await,
             SetWireguardMtu(tx, mtu) => self.on_set_wireguard_mtu(tx, mtu).await,
             SetWireguardRotationInterval(tx, interval) => {
                 self.on_set_wireguard_rotation_interval(tx, interval).await
@@ -2163,6 +2173,56 @@ where
         }
     }
 
+    async fn on_set_relay_override(
+        &mut self,
+        tx: ResponseTx<(), settings::Error>,
+        relay_override: RelayOverride,
+    ) {
+        match self
+            .settings
+            .update(move |settings| settings.set_relay_override(relay_override))
+            .await
+        {
+            Ok(settings_changed) => {
+                Self::oneshot_send(tx, Ok(()), "set_relay_override response");
+                if settings_changed {
+                    self.event_listener
+                        .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
+                    self.reconnect_tunnel();
+                }
+            }
+            Err(e) => {
+                log::error!("{}", e.display_chain_with_msg("Unable to save settings"));
+                Self::oneshot_send(tx, Err(e), "set_relay_override response");
+            }
+        }
+    }
+
+    async fn on_clear_all_relay_overrides(&mut self, tx: ResponseTx<(), settings::Error>) {
+        match self
+            .settings
+            .update(move |settings| settings.relay_overrides.clear())
+            .await
+        {
+            Ok(settings_changed) => {
+                Self::oneshot_send(tx, Ok(()), "clear_all_relay_overrides response");
+                if settings_changed {
+                    self.event_listener
+                        .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
+                    self.reconnect_tunnel();
+                }
+            }
+            Err(e) => {
+                log::error!("{}", e.display_chain_with_msg("Unable to save settings"));
+                Self::oneshot_send(tx, Err(e), "clear_all_relay_overrides response");
+            }
+        }
+    }
+
     async fn on_set_wireguard_mtu(
         &mut self,
         tx: ResponseTx<(), settings::Error>,
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 79466d8878..e67a02117c 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -13,7 +13,9 @@ use mullvad_paths;
 use mullvad_types::settings::DnsOptions;
 use mullvad_types::{
     account::AccountToken,
-    relay_constraints::{BridgeSettings, BridgeState, ObfuscationSettings, RelaySettings},
+    relay_constraints::{
+        BridgeSettings, BridgeState, ObfuscationSettings, RelayOverride, RelaySettings,
+    },
     relay_list::RelayList,
     settings::Settings,
     states::{TargetState, TunnelState},
@@ -379,6 +381,31 @@ impl ManagementService for ManagementServiceImpl {
         Ok(Response::new(()))
     }
 
+    async fn set_relay_override(
+        &self,
+        request: Request<types::RelayOverride>,
+    ) -> ServiceResult<()> {
+        let relay_override =
+            RelayOverride::try_from(request.into_inner()).map_err(map_protobuf_type_err)?;
+        log::debug!("set_relay_override");
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::SetRelayOverride(tx, relay_override))?;
+        self.wait_for_result(rx)
+            .await?
+            .map(Response::new)
+            .map_err(map_settings_error)
+    }
+
+    async fn clear_all_relay_overrides(&self, _: Request<()>) -> ServiceResult<()> {
+        log::debug!("clear_all_relay_overrides");
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::ClearAllRelayOverrides(tx))?;
+        self.wait_for_result(rx)
+            .await?
+            .map(Response::new)
+            .map_err(map_settings_error)
+    }
+
     // Account management
     //
 
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 17854b6956..64d87d37ef 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -46,6 +46,8 @@ service ManagementService {
   rpc SetEnableIpv6(google.protobuf.BoolValue) returns (google.protobuf.Empty) {}
   rpc SetQuantumResistantTunnel(QuantumResistantState) returns (google.protobuf.Empty) {}
   rpc SetDnsOptions(DnsOptions) returns (google.protobuf.Empty) {}
+  rpc SetRelayOverride(RelayOverride) returns (google.protobuf.Empty) {}
+  rpc ClearAllRelayOverrides(google.protobuf.Empty) returns (google.protobuf.Empty) {}
 
   // Account management
   rpc CreateNewAccount(google.protobuf.Empty) returns (google.protobuf.StringValue) {}
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 219d0940ba..e036734896 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -1008,3 +1008,17 @@ pub struct RelayOverride {
     /// IPv6 address to use instead of the default
     pub ipv6_addr_in: Option<Ipv6Addr>,
 }
+
+impl RelayOverride {
+    pub fn empty(hostname: Hostname) -> RelayOverride {
+        RelayOverride {
+            hostname,
+            ipv4_addr_in: None,
+            ipv6_addr_in: None,
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self == &Self::empty(self.hostname.clone())
+    }
+}
diff --git a/mullvad-types/src/settings/mod.rs b/mullvad-types/src/settings/mod.rs
index 2ed2d79394..612e353487 100644
--- a/mullvad-types/src/settings/mod.rs
+++ b/mullvad-types/src/settings/mod.rs
@@ -175,6 +175,24 @@ impl Settings {
             self.relay_settings = new_settings;
         }
     }
+
+    pub fn set_relay_override(&mut self, relay_override: RelayOverride) {
+        let existing_override = self
+            .relay_overrides
+            .iter_mut()
+            .enumerate()
+            .find(|(_, elem)| elem.hostname == relay_override.hostname);
+        match existing_override {
+            None => self.relay_overrides.push(relay_override),
+            Some((index, elem)) => {
+                if relay_override.is_empty() {
+                    self.relay_overrides.swap_remove(index);
+                } else {
+                    *elem = relay_override;
+                }
+            }
+        }
+    }
 }
 
 /// TunnelOptions holds configuration data that applies to all kinds of tunnels.