Merge branch 'implement-support-for-separate-entryexit-filters-droid-2199'

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"
             ))),