Add support for multihop entry filters in daemon

In the upcoming re-design of select location, separate sets of filters
can now be picked for the entry and the exit relays. This commit adds
support for that in the relay selector.

In order to not affect the current behavior of the desktop and Android
apps before the new UI is implemented, the entry filters are set to the
same as the exit filters when the relay settings are updated via gRPC.

10 files changed, 310 insertions, 9 deletions

diff --git a/mullvad-daemon/src/migrations/mod.rs b/mullvad-daemon/src/migrations/mod.rs
index ea7d4523ee..27a715b95d 100644
--- a/mullvad-daemon/src/migrations/mod.rs
+++ b/mullvad-daemon/src/migrations/mod.rs
@@ -48,6 +48,7 @@ mod device;
 mod v1;
 mod v10;
 mod v11;
+mod v12;
 mod v2;
 mod v3;
 mod v4;
@@ -210,8 +211,8 @@ async fn migrate_settings(
     )?;
 
     v10::migrate(settings)?;
-
     v11::migrate(settings)?;
+    v12::migrate(settings)?;
 
     Ok(migration_data)
 }
diff --git a/mullvad-daemon/src/migrations/snapshots/mullvad_daemon__migrations__v12__test__v12_to_v13_migration.snap b/mullvad-daemon/src/migrations/snapshots/mullvad_daemon__migrations__v12__test__v12_to_v13_migration.snap
new file mode 100644
index 0000000000..de7df7c9f0
--- /dev/null
+++ b/mullvad-daemon/src/migrations/snapshots/mullvad_daemon__migrations__v12__test__v12_to_v13_migration.snap
@@ -0,0 +1,57 @@
+---
+source: mullvad-daemon/src/migrations/v12.rs
+expression: "serde_json::to_string_pretty(&old_settings).unwrap()"
+---
+{
+  "relay_settings": {
+    "normal": {
+      "location": {
+        "only": {
+          "location": {
+            "country": "fr"
+          }
+        }
+      },
+      "openvpn_constraints": {
+        "port": "any"
+      },
+      "ownership": {
+        "only": "MullvadOwned"
+      },
+      "providers": {
+        "only": {
+          "providers": [
+            "Blix",
+            "Creanova"
+          ]
+        }
+      },
+      "tunnel_protocol": "wireguard",
+      "wireguard_constraints": {
+        "allowed_ips": "any",
+        "entry_location": {
+          "only": {
+            "location": {
+              "country": "se"
+            }
+          }
+        },
+        "entry_ownership": {
+          "only": "MullvadOwned"
+        },
+        "entry_providers": {
+          "only": {
+            "providers": [
+              "Blix",
+              "Creanova"
+            ]
+          }
+        },
+        "ip_version": "any",
+        "port": "any",
+        "use_multihop": true
+      }
+    }
+  },
+  "settings_version": 13
+}
diff --git a/mullvad-daemon/src/migrations/v12.rs b/mullvad-daemon/src/migrations/v12.rs
new file mode 100644
index 0000000000..b396cbfdfa
--- /dev/null
+++ b/mullvad-daemon/src/migrations/v12.rs
@@ -0,0 +1,102 @@
+use super::Result;
+use mullvad_types::settings::SettingsVersion;
+
+/// This version introduces 2 new fields to the [mullvad_constraints::WireguardConstraints] struct:
+/// pub entry_providers: Constraint<Providers>,
+/// pub entry_ownership: Constraint<Ownership>,
+/// When set, these filters apply to the entry relay when multihop is used.
+/// A migration is needed to transfer the current providers and ownership to these new fields
+/// so that the user's current filters don't change.
+pub fn migrate(settings: &mut serde_json::Value) -> Result<()> {
+    if !version_matches(settings) {
+        return Ok(());
+    }
+
+    log::info!("Migrating settings format to V13");
+
+    migrate_filters_to_new_entry_only_filters(settings);
+
+    settings["settings_version"] = serde_json::json!(SettingsVersion::V13);
+
+    Ok(())
+}
+
+fn version_matches(settings: &serde_json::Value) -> bool {
+    settings
+        .get("settings_version")
+        .map(|version| version == SettingsVersion::V12 as u64)
+        .unwrap_or(false)
+}
+
+fn migrate_filters_to_new_entry_only_filters(settings: &mut serde_json::Value) -> Option<()> {
+    let normal = settings.get_mut("relay_settings")?.get_mut("normal")?;
+    let providers = normal.get("providers")?.clone();
+    let ownership = normal.get("ownership")?.clone();
+
+    let wireguard_constraints = normal.get_mut("wireguard_constraints")?.as_object_mut()?;
+
+    wireguard_constraints.insert("entry_providers".to_string(), providers);
+    wireguard_constraints.insert("entry_ownership".to_string(), ownership);
+
+    Some(())
+}
+
+#[cfg(test)]
+mod test {
+    use super::{migrate, version_matches};
+
+    const V12_SETTINGS: &str = r#"
+{
+  "relay_settings": {
+    "normal": {
+      "location": {
+        "only": {
+          "location": {
+            "country": "fr"
+          }
+        }
+      },
+      "providers": {
+        "only": {
+          "providers": [
+            "Blix",
+            "Creanova"
+          ]
+        }
+      },
+      "ownership": {
+        "only": "MullvadOwned"
+      },
+      "tunnel_protocol": "wireguard",
+      "wireguard_constraints": {
+        "port": "any",
+        "ip_version": "any",
+        "allowed_ips": "any",
+        "use_multihop": true,
+        "entry_location": {
+          "only": {
+            "location": {
+              "country": "se"
+            }
+          }
+        }
+      },
+      "openvpn_constraints": {
+        "port": "any"
+      }
+    }
+  },
+  "settings_version": 12
+}
+"#;
+
+    #[test]
+    fn test_v12_to_v13_migration() {
+        let mut old_settings = serde_json::from_str(V12_SETTINGS).unwrap();
+
+        assert!(version_matches(&old_settings));
+
+        migrate(&mut old_settings).unwrap();
+        insta::assert_snapshot!(serde_json::to_string_pretty(&old_settings).unwrap());
+    }
+}
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 21e00573af..25e0321ce4 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -604,6 +604,8 @@ message WireguardConstraints {
   repeated string allowed_ips = 3;
   bool use_multihop = 4;
   LocationConstraint entry_location = 5;
+  repeated string entry_providers = 6;
+  Ownership entry_ownership = 7;
 }
 
 message CustomRelaySettings {
diff --git a/mullvad-management-interface/src/types/conversions/relay_constraints.rs b/mullvad-management-interface/src/types/conversions/relay_constraints.rs
index 6fa52b3fa5..c86107cd61 100644
--- a/mullvad-management-interface/src/types/conversions/relay_constraints.rs
+++ b/mullvad-management-interface/src/types/conversions/relay_constraints.rs
@@ -51,6 +51,8 @@ impl TryFrom<&proto::WireguardConstraints>
                     .ok()
                 })
                 .unwrap_or(Constraint::Any),
+            entry_providers: try_providers_constraint_from_proto(&constraints.entry_providers)?,
+            entry_ownership: try_ownership_constraint_from_i32(constraints.entry_ownership)?,
         })
     }
 }
@@ -117,11 +119,21 @@ impl TryFrom<proto::RelaySettings> for mullvad_types::relay_constraints::RelaySe
                             FromProtobufTypeError::InvalidArgument("missing openvpn constraints"),
                         )?,
                     )?;
-                let wireguard_constraints = mullvad_constraints::WireguardConstraints::try_from(
-                    &settings.wireguard_constraints.ok_or(
-                        FromProtobufTypeError::InvalidArgument("missing wireguard constraints"),
-                    )?,
-                )?;
+                let mut wireguard_constraints =
+                    mullvad_constraints::WireguardConstraints::try_from(
+                        &settings.wireguard_constraints.ok_or(
+                            FromProtobufTypeError::InvalidArgument("missing wireguard constraints"),
+                        )?,
+                    )?;
+
+                // TODO Remove this block when the frontends support setting multihop entry filters.
+                // This is needed in order to not change the current behavior (which
+                // is that the ownership and providers from `RelaySettings` apply to both the entry
+                // and exit multihop relays).
+                {
+                    wireguard_constraints.entry_ownership = ownership;
+                    wireguard_constraints.entry_providers = providers.clone();
+                }
 
                 Ok(mullvad_constraints::RelaySettings::Normal(
                     mullvad_constraints::RelayConstraints {
@@ -279,6 +291,12 @@ impl From<mullvad_types::relay_constraints::RelaySettings> for proto::RelaySetti
                             .entry_location
                             .option()
                             .map(proto::LocationConstraint::from),
+                        entry_providers: convert_providers_constraint(
+                            &constraints.wireguard_constraints.entry_providers,
+                        ),
+                        entry_ownership: convert_ownership_constraint(
+                            &constraints.wireguard_constraints.entry_ownership,
+                        ) as i32,
                     }),
 
                     openvpn_constraints: Some(proto::OpenvpnConstraints {
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index dc454a2cfd..d5087f43a3 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -350,6 +350,8 @@ impl<'a> TryFrom<NormalSelectorConfig<'a>> for RelayQuery {
                 allowed_ips,
                 use_multihop,
                 entry_location,
+                entry_providers,
+                entry_ownership,
             } = wireguard_constraints;
             let AdditionalWireguardConstraints {
                 daita,
@@ -362,6 +364,8 @@ impl<'a> TryFrom<NormalSelectorConfig<'a>> for RelayQuery {
                 allowed_ips,
                 use_multihop: Constraint::Only(use_multihop),
                 entry_location,
+                entry_providers,
+                entry_ownership,
                 obfuscation: ObfuscationQuery::from(obfuscation_settings),
                 daita: Constraint::Only(daita),
                 daita_use_multihop_if_necessary: Constraint::Only(daita_use_multihop_if_necessary),
@@ -829,10 +833,12 @@ impl RelaySelector {
     ) -> Result<Multihop, Error> {
         // Here, we modify the original query just a bit.
         // The actual query for an entry relay is identical as for an exit relay, with the
-        // exception that the location is different. It is simply the location as dictated by
-        // the query's multihop constraint.
+        // exception that the location is different and that the entry filters may be different.
+        // The location is dictated by the query's multihop constraint.
         let mut entry_relay_query = query.clone();
         entry_relay_query.set_location(query.wireguard_constraints().entry_location.clone())?;
+        entry_relay_query.set_providers(query.wireguard_constraints().entry_providers.clone());
+        entry_relay_query.set_ownership(query.wireguard_constraints().entry_ownership);
         // After we have our two queries (one for the exit relay & one for the entry relay),
         // we can query for all exit & entry candidates! All candidates are needed for the next
         // step.
diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
index 4e4813a9dd..93931cfa7d 100644
--- a/mullvad-relay-selector/src/relay_selector/query.rs
+++ b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -139,10 +139,18 @@ impl RelayQuery {
         &self.providers
     }
 
+    pub fn set_providers(&mut self, providers: Constraint<Providers>) {
+        self.providers = providers;
+    }
+
     pub fn ownership(&self) -> Constraint<Ownership> {
         self.ownership
     }
 
+    pub fn set_ownership(&mut self, ownership: Constraint<Ownership>) {
+        self.ownership = ownership;
+    }
+
     pub fn tunnel_protocol(&self) -> TunnelType {
         self.tunnel_protocol
     }
@@ -276,6 +284,8 @@ pub struct WireguardRelayQuery {
     pub allowed_ips: Constraint<AllowedIps>,
     pub use_multihop: Constraint<bool>,
     pub entry_location: Constraint<LocationConstraint>,
+    pub entry_providers: Constraint<Providers>,
+    pub entry_ownership: Constraint<Ownership>,
     pub obfuscation: ObfuscationQuery,
     pub daita: Constraint<bool>,
     pub daita_use_multihop_if_necessary: Constraint<bool>,
@@ -375,6 +385,8 @@ impl WireguardRelayQuery {
             allowed_ips: Constraint::Any,
             use_multihop: Constraint::Any,
             entry_location: Constraint::Any,
+            entry_providers: Constraint::Any,
+            entry_ownership: Constraint::Any,
             obfuscation: ObfuscationQuery::Auto,
             daita: Constraint::Any,
             daita_use_multihop_if_necessary: Constraint::Any,
@@ -389,6 +401,8 @@ impl WireguardRelayQuery {
             ip_version: self.ip_version,
             allowed_ips: self.allowed_ips,
             entry_location: self.entry_location,
+            entry_providers: self.entry_providers,
+            entry_ownership: self.entry_ownership,
             use_multihop: self.use_multihop.unwrap_or(false),
         }
     }
@@ -408,6 +422,8 @@ impl From<WireguardRelayQuery> for WireguardConstraints {
             ip_version: value.ip_version,
             allowed_ips: value.allowed_ips,
             entry_location: value.entry_location,
+            entry_providers: value.entry_providers,
+            entry_ownership: value.entry_ownership,
             use_multihop: value.use_multihop.unwrap_or(false),
         }
     }
@@ -765,6 +781,20 @@ pub mod builder {
             self.query.wireguard_constraints.entry_location = Constraint::Only(location.into());
             self
         }
+
+        /// Set the entry location in a multihop configuration. This requires
+        /// multihop to be enabled.
+        pub fn entry_providers(mut self, providers: Providers) -> Self {
+            self.query.wireguard_constraints.entry_providers = Constraint::Only(providers);
+            self
+        }
+
+        /// Set the entry location in a multihop configuration. This requires
+        /// multihop to be enabled.
+        pub fn entry_ownership(mut self, ownership: Ownership) -> Self {
+            self.query.wireguard_constraints.entry_ownership = Constraint::Only(ownership);
+            self
+        }
     }
 
     impl<Multihop, Daita, QuantumResistant>
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 7f35ff7084..37b97dcc4a 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -284,6 +284,18 @@ fn unwrap_entry_relay(get_result: GetRelay) -> Relay {
     }
 }
 
+fn unwrap_multihop_entry_exit_relays(get_result: GetRelay) -> (Relay, Relay) {
+    match get_result {
+        GetRelay::Wireguard {
+            inner: crate::WireguardConfig::Multihop { entry, exit },
+            ..
+        } => (entry, exit),
+        relay => {
+            panic!("Relay is not a Wireguard multihop relay: {relay:?}")
+        }
+    }
+}
+
 fn unwrap_endpoint(get_result: GetRelay) -> MullvadEndpoint {
     match get_result {
         GetRelay::Wireguard { endpoint, .. } => MullvadEndpoint::Wireguard(endpoint),
@@ -1192,6 +1204,41 @@ fn test_ownership() {
     }
 }
 
+/// Verify that any query which sets an explicit [`Ownership`] is respected by the relay selector
+/// and that it works to set separate entry and exit ownerships for a multihop.
+#[test]
+fn test_multihop_ownership() {
+    let relay_selector = default_relay_selector();
+
+    for _ in 0..100 {
+        // Construct an arbitrary query for owned relays.
+        let query = RelayQueryBuilder::wireguard()
+            .multihop()
+            .ownership(Ownership::MullvadOwned)
+            .entry_ownership(Ownership::Rented)
+            .build();
+        let relay = relay_selector.get_relay_by_query(query).unwrap();
+        // Check that the _exit_ relay is owned by Mullvad.
+        assert!(unwrap_relay(relay.clone()).owned);
+        // Check that the _entry_ relay is rented.
+        assert!(!unwrap_entry_relay(relay).owned);
+    }
+
+    for _ in 0..100 {
+        // Construct an arbitrary query for rented relays.
+        let query = RelayQueryBuilder::wireguard()
+            .multihop()
+            .ownership(Ownership::Rented)
+            .entry_ownership(Ownership::MullvadOwned)
+            .build();
+        let relay = relay_selector.get_relay_by_query(query).unwrap();
+        // Check that the _exit_ relay is rented.
+        assert!(!unwrap_relay(relay.clone()).owned);
+        // Check that the _entry_ relay is owned by Mullvad.
+        assert!(unwrap_entry_relay(relay).owned);
+    }
+}
+
 /// Verify that server and port selection varies between retry attempts.
 #[test]
 fn test_load_balancing() {
@@ -1258,6 +1305,40 @@ fn test_providers() {
     }
 }
 
+/// Construct a query for a relay with specific providers and verify that every chosen relay has
+/// the correct associated provider and that it works to select a separate set of providers for
+/// entry and exit relays when doing a multihop.
+#[test]
+fn test_multihop_providers() {
+    const EXPECTED_PROVIDERS: [&str; 2] = ["provider0", "provider2"];
+    const EXPECTED_ENTRY_PROVIDERS: [&str; 2] = ["provider1", "provider3"];
+    let providers = Providers::new(EXPECTED_PROVIDERS).unwrap();
+    let entry_providers = Providers::new(EXPECTED_ENTRY_PROVIDERS).unwrap();
+    let relay_selector = default_relay_selector();
+
+    for _attempt in 0..100 {
+        let query = RelayQueryBuilder::wireguard()
+            .multihop()
+            .providers(providers.clone())
+            .entry_providers(entry_providers.clone())
+            .build();
+        let relay = relay_selector.get_relay_by_query(query).unwrap();
+
+        let (entry, exit) = unwrap_multihop_entry_exit_relays(relay);
+
+        assert!(
+            EXPECTED_PROVIDERS.contains(&exit.provider.as_str()),
+            "cannot find exit provider {provider} in {EXPECTED_PROVIDERS:?}",
+            provider = exit.provider
+        );
+        assert!(
+            EXPECTED_ENTRY_PROVIDERS.contains(&entry.provider.as_str()),
+            "cannot find entry provider {provider} in {EXPECTED_ENTRY_PROVIDERS:?}",
+            provider = entry.provider
+        );
+    }
+}
+
 /// Verify that bridges are automatically used when bridge mode is set to automatic.
 #[test]
 fn test_openvpn_auto_bridge() {
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 013153524e..537ddb0152 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -410,6 +410,8 @@ pub struct WireguardConstraints {
     pub allowed_ips: Constraint<AllowedIps>,
     pub use_multihop: bool,
     pub entry_location: Constraint<LocationConstraint>,
+    pub entry_providers: Constraint<Providers>,
+    pub entry_ownership: Constraint<Ownership>,
 }
 
 pub use allowed_ip::AllowedIps;
diff --git a/mullvad-types/src/settings/mod.rs b/mullvad-types/src/settings/mod.rs
index 7b4f97f891..221d04ee47 100644
--- a/mullvad-types/src/settings/mod.rs
+++ b/mullvad-types/src/settings/mod.rs
@@ -20,7 +20,7 @@ mod dns;
 /// latest version that exists in `SettingsVersion`.
 /// This should be bumped when a new version is introduced along with a migration
 /// being added to `mullvad-daemon`.
-pub const CURRENT_SETTINGS_VERSION: SettingsVersion = SettingsVersion::V12;
+pub const CURRENT_SETTINGS_VERSION: SettingsVersion = SettingsVersion::V13;
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
 #[repr(u32)]
@@ -36,6 +36,7 @@ pub enum SettingsVersion {
     V10 = 10,
     V11 = 11,
     V12 = 12,
+    V13 = 13,
 }
 
 impl<'de> Deserialize<'de> for SettingsVersion {
@@ -55,6 +56,7 @@ impl<'de> Deserialize<'de> for SettingsVersion {
             v if v == SettingsVersion::V10 as u32 => Ok(SettingsVersion::V10),
             v if v == SettingsVersion::V11 as u32 => Ok(SettingsVersion::V11),
             v if v == SettingsVersion::V12 as u32 => Ok(SettingsVersion::V12),
+            v if v == SettingsVersion::V13 as u32 => Ok(SettingsVersion::V13),
             v => Err(serde::de::Error::custom(format!(
                 "{v} is not a valid SettingsVersion"
             ))),