Add provider relay constraint

4 files changed, 52 insertions, 18 deletions

diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index c8ec8445d0..a2414bae66 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -967,6 +967,13 @@ fn convert_relay_settings_update(
 
             Ok(RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
                 location,
+                provider: settings.provider.map(|provider_update| {
+                    if !provider_update.provider.is_empty() {
+                        Constraint::Only(provider_update.provider.clone())
+                    } else {
+                        Constraint::Any
+                    }
+                }),
                 tunnel_protocol,
                 wireguard_constraints: settings.wireguard_constraints.map(|constraints| {
                     WireguardConstraints {
@@ -1005,6 +1012,10 @@ fn convert_relay_settings(settings: &RelaySettings) -> types::RelaySettings {
         RelaySettings::Normal(constraints) => {
             relay_settings::Endpoint::Normal(types::NormalRelaySettings {
                 location: convert_location_constraint(&constraints.location),
+                provider: match &constraints.provider {
+                    Constraint::Any => "".to_string(),
+                    Constraint::Only(ref provider) => provider.clone(),
+                },
                 tunnel_type: match constraints.tunnel_protocol {
                     Constraint::Any => None,
                     Constraint::Only(TunnelType::Wireguard) => Some(types::TunnelType::Wireguard),
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 8640b9069e..075fa2468e 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -14,7 +14,7 @@ use mullvad_types::{
     location::Location,
     relay_constraints::{
         BridgeState, Constraint, InternalBridgeConstraints, LocationConstraint, Match,
-        OpenVpnConstraints, RelayConstraints, WireguardConstraints,
+        OpenVpnConstraints, Provider, RelayConstraints, WireguardConstraints,
     },
     relay_list::{OpenVpnEndpointData, Relay, RelayList, RelayTunnels, WireguardEndpointData},
 };
@@ -253,6 +253,7 @@ impl RelaySelector {
                 self.preferred_tunnel_constraints(
                     retry_attempt,
                     &original_constraints.location,
+                    &original_constraints.provider,
                     wg_key_exists,
                 )
             } else {
@@ -260,12 +261,9 @@ impl RelaySelector {
             };
 
 
-        let mut relay_constraints = RelayConstraints {
-            location: original_constraints.location.clone(),
-            tunnel_protocol: original_constraints.tunnel_protocol.clone(),
-            wireguard_constraints: original_constraints.wireguard_constraints,
-            ..Default::default()
-        };
+        let mut relay_constraints = original_constraints.clone();
+        relay_constraints.openvpn_constraints = Default::default();
+
         // Highest priority preference. Where we prefer OpenVPN using UDP. But without changing
         // any constraints that are explicitly specified.
         match original_constraints.tunnel_protocol {
@@ -380,6 +378,7 @@ impl RelaySelector {
         &self,
         retry_attempt: u32,
         location_constraint: &Constraint<LocationConstraint>,
+        provider_constraint: &Constraint<Provider>,
         wg_key_exists: bool,
     ) -> (Constraint<u16>, TransportProtocol, TunnelType) {
         #[cfg(not(target_os = "windows"))]
@@ -389,6 +388,7 @@ impl RelaySelector {
                     relay.active
                         && !relay.tunnels.wireguard.is_empty()
                         && location_constraint.matches(relay)
+                        && provider_constraint.matches_eq(&relay.provider)
                 });
             // If location does not support WireGuard, defer to preferred OpenVPN tunnel
             // constraints
@@ -475,6 +475,9 @@ impl RelaySelector {
         if !constraints.location.matches(relay) {
             return None;
         }
+        if !constraints.provider.matches_eq(&relay.provider) {
+            return None;
+        }
 
 
         let relay = match constraints.tunnel_protocol {
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 8616163f50..fae924238f 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -270,17 +270,23 @@ message TunnelTypeConstraint {
 
 message NormalRelaySettings {
 	RelayLocation location = 1;
-	TunnelTypeConstraint tunnel_type = 2;
-	WireguardConstraints wireguard_constraints = 3;
-	OpenvpnConstraints openvpn_constraints = 4;
+	string provider = 2;
+	TunnelTypeConstraint tunnel_type = 3;
+	WireguardConstraints wireguard_constraints = 4;
+	OpenvpnConstraints openvpn_constraints = 5;
 }
 
 // Constraints are only updated for fields that are provided
 message NormalRelaySettingsUpdate {
 	RelayLocation location = 1;
-	TunnelTypeUpdate tunnel_type = 2;
-	WireguardConstraints wireguard_constraints = 3;
-	OpenvpnConstraints openvpn_constraints = 4;
+	ProviderUpdate provider = 2;
+	TunnelTypeUpdate tunnel_type = 3;
+	WireguardConstraints wireguard_constraints = 4;
+	OpenvpnConstraints openvpn_constraints = 5;
+}
+
+message ProviderUpdate {
+	string provider = 1;
 }
 
 message TunnelTypeUpdate {
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 857b8747fb..cbd91d3def 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -179,11 +179,13 @@ impl RelaySettings {
 
 /// Limits the set of [`crate::relay_list::Relay`]s that a `RelaySelector` may select.
 #[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(default)]
 #[cfg_attr(not(target_os = "android"), derive(Default))]
 #[cfg_attr(target_os = "android", derive(IntoJava))]
 #[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
 pub struct RelayConstraints {
     pub location: Constraint<LocationConstraint>,
+    pub provider: Constraint<Provider>,
     #[cfg_attr(target_os = "android", jnix(skip))]
     pub tunnel_protocol: Constraint<TunnelType>,
     #[cfg_attr(target_os = "android", jnix(skip))]
@@ -196,10 +198,8 @@ pub struct RelayConstraints {
 impl Default for RelayConstraints {
     fn default() -> Self {
         RelayConstraints {
-            location: Constraint::Any,
             tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
-            wireguard_constraints: WireguardConstraints::default(),
-            openvpn_constraints: OpenVpnConstraints::default(),
+            ..Default::default()
         }
     }
 }
@@ -208,6 +208,7 @@ impl RelayConstraints {
     pub fn merge(&self, update: RelayConstraintsUpdate) -> Self {
         RelayConstraints {
             location: update.location.unwrap_or_else(|| self.location.clone()),
+            provider: update.provider.unwrap_or_else(|| self.provider.clone()),
             tunnel_protocol: update
                 .tunnel_protocol
                 .unwrap_or_else(|| self.tunnel_protocol.clone()),
@@ -243,8 +244,16 @@ impl fmt::Display for RelayConstraints {
         }
         write!(f, " in ")?;
         match self.location {
-            Constraint::Any => write!(f, "any location"),
-            Constraint::Only(ref location_constraint) => location_constraint.fmt(f),
+            Constraint::Any => write!(f, "any location")?,
+            Constraint::Only(ref location_constraint) => location_constraint.fmt(f)?,
+        }
+        write!(f, " using ")?;
+        match self.provider {
+            Constraint::Any => write!(f, "any provider"),
+            Constraint::Only(ref constraint) => {
+                write!(f, "provider ")?;
+                constraint.fmt(f)
+            }
         }
     }
 }
@@ -291,6 +300,10 @@ impl Match<Relay> for LocationConstraint {
     }
 }
 
+/// Limits the set of [`crate::relay_list::Relay`]s used by a `RelaySelector` based on
+/// provider.
+pub type Provider = String;
+
 impl fmt::Display for LocationConstraint {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         match self {
@@ -499,6 +512,7 @@ impl RelaySettingsUpdate {
 #[serde(default)]
 pub struct RelayConstraintsUpdate {
     pub location: Option<Constraint<LocationConstraint>>,
+    pub provider: Option<Constraint<Provider>>,
     #[cfg_attr(target_os = "android", jnix(default))]
     pub tunnel_protocol: Option<Constraint<TunnelType>>,
     #[cfg_attr(target_os = "android", jnix(default))]