Merge branch 'fix-tests/relay-location-constraints'

2 files changed, 65 insertions, 88 deletions

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index f73e6f6e79..431e3f9c51 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -4,8 +4,8 @@ use futures::StreamExt;
 use mullvad_management_interface::{types, ManagementServiceClient};
 use mullvad_types::{
     relay_constraints::{
-        Constraint, GeographicLocationConstraint, LocationConstraint, OpenVpnConstraints,
-        RelayConstraintsUpdate, RelaySettingsUpdate, WireguardConstraints,
+        BridgeState, Constraint, GeographicLocationConstraint, LocationConstraint,
+        ObfuscationSettings, RelayConstraintsUpdate, RelaySettingsUpdate,
     },
     states::TunnelState,
 };
@@ -33,6 +33,22 @@ pub fn get_package_desc(name: &str) -> Result<Package, Error> {
     })
 }
 
+/// Reboot the guest virtual machine.
+///
+/// # macOS
+/// The tunnel must be reconfigured after the virtual machine is up,
+/// or macOS refuses to assign an IP. The reasons for this are poorly understood.
+pub async fn reboot(rpc: &mut ServiceClient) -> Result<(), Error> {
+    rpc.reboot().await?;
+
+    #[cfg(target_os = "macos")]
+    crate::vm::network::macos::configure_tunnel()
+        .await
+        .map_err(|error| Error::Other(format!("Failed to recreate custom wg tun: {error}")))?;
+
+    Ok(())
+}
+
 #[derive(Debug, Default)]
 pub struct ProbeResult {
     tcp: usize,
@@ -303,20 +319,24 @@ impl<T> Drop for AbortOnDrop<T> {
 }
 
 /// Disconnect and reset all relay, bridge, and obfuscation settings.
+///
+/// See [`mullvad_types::relay_constraints::RelayConstraintsUpdate`] for details, but in short:
+/// * Location constraint is [`Constraint::Any`]
+/// * Provider constraint is [`Constraint::Any`]
+/// * Ownership constraint is [`Constraint::Any`]
+/// * The default tunnel protocol is [`talpid_types::net::TunnelType::Wireguard`]
+/// * Wireguard settings are default (i.e. any port is used, no obfuscation ..)
+///   see [`mullvad_types::relay_constraints::WireguardConstraints`] for details.
+/// * OpenVPN settings are default (i.e. any port is used, no obfuscation ..)
+///   see [`mullvad_types::relay_constraints::OpenVpnConstraints`] for details.
 pub async fn reset_relay_settings(
     mullvad_client: &mut ManagementServiceClient,
 ) -> Result<(), Error> {
     disconnect_and_wait(mullvad_client).await?;
 
-    let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-        location: Some(Constraint::Only(LocationConstraint::Location(
-            GeographicLocationConstraint::Country("se".to_string()),
-        ))),
-        tunnel_protocol: Some(Constraint::Any),
-        openvpn_constraints: Some(OpenVpnConstraints::default()),
-        wireguard_constraints: Some(WireguardConstraints::default()),
-        ..Default::default()
-    });
+    let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate::default());
+    let bridge_state = BridgeState::Auto;
+    let obfuscation_settings = ObfuscationSettings::default();
 
     update_relay_settings(mullvad_client, relay_settings)
         .await
@@ -325,20 +345,16 @@ pub async fn reset_relay_settings(
         })?;
 
     mullvad_client
-        .set_bridge_state(types::BridgeState {
-            state: i32::from(types::bridge_state::State::Auto),
-        })
+        .set_bridge_state(types::BridgeState::from(bridge_state))
         .await
         .map_err(|error| Error::DaemonError(format!("Failed to reset bridge mode: {}", error)))?;
 
     mullvad_client
-        .set_obfuscation_settings(types::ObfuscationSettings {
-            selected_obfuscation: i32::from(types::obfuscation_settings::SelectedObfuscation::Off),
-            udp2tcp: Some(types::Udp2TcpObfuscationSettings { port: None }),
-        })
+        .set_obfuscation_settings(types::ObfuscationSettings::from(obfuscation_settings))
         .await
-        .map(|_| ())
-        .map_err(|error| Error::DaemonError(format!("Failed to reset obfuscation: {}", error)))
+        .map_err(|error| Error::DaemonError(format!("Failed to reset obfuscation: {}", error)))?;
+
+    Ok(())
 }
 
 pub async fn update_relay_settings(
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index be0a24dd14..99e30b087c 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -5,10 +5,11 @@ use std::net::IpAddr;
 use crate::network_monitor::{start_packet_monitor, MonitorOptions};
 use mullvad_management_interface::{types, ManagementServiceClient};
 use mullvad_types::relay_constraints::{
-    Constraint, LocationConstraint, OpenVpnConstraints, RelayConstraintsUpdate,
-    RelaySettingsUpdate, WireguardConstraints,
+    BridgeConstraints, BridgeSettings, BridgeState, Constraint, ObfuscationSettings,
+    OpenVpnConstraints, RelayConstraintsUpdate, RelaySettingsUpdate, SelectedObfuscation,
+    TransportPort, Udp2TcpObfuscationSettings, WireguardConstraints,
 };
-use mullvad_types::relay_constraints::{GeographicLocationConstraint, TransportPort};
+use mullvad_types::wireguard;
 use pnet_packet::ip::IpNextHeaderProtocols;
 use talpid_types::net::{TransportProtocol, TunnelType};
 use test_macro::test_function;
@@ -48,9 +49,6 @@ pub async fn test_openvpn_tunnel(
         log::info!("Connect to {protocol} OpenVPN endpoint");
 
         let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-            location: Some(Constraint::Only(LocationConstraint::Location(
-                GeographicLocationConstraint::Country("se".to_string()),
-            ))),
             tunnel_protocol: Some(Constraint::Only(TunnelType::OpenVpn)),
             openvpn_constraints: Some(OpenVpnConstraints { port: constraint }),
             ..Default::default()
@@ -91,9 +89,6 @@ pub async fn test_wireguard_tunnel(
         log::info!("Connect to WireGuard endpoint on port {port}");
 
         let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-            location: Some(Constraint::Only(LocationConstraint::Location(
-                GeographicLocationConstraint::Country("se".to_string()),
-            ))),
             tunnel_protocol: Some(Constraint::Only(TunnelType::Wireguard)),
             wireguard_constraints: Some(WireguardConstraints {
                 port: Constraint::Only(port),
@@ -139,21 +134,17 @@ pub async fn test_udp2tcp_tunnel(
     // TODO: ping a public IP on the fake network (not possible using real relay)
 
     mullvad_client
-        .set_obfuscation_settings(types::ObfuscationSettings {
-            selected_obfuscation: i32::from(
-                types::obfuscation_settings::SelectedObfuscation::Udp2tcp,
-            ),
-            udp2tcp: Some(types::Udp2TcpObfuscationSettings { port: None }),
-        })
+        .set_obfuscation_settings(types::ObfuscationSettings::from(ObfuscationSettings {
+            selected_obfuscation: SelectedObfuscation::Udp2Tcp,
+            udp2tcp: Udp2TcpObfuscationSettings {
+                port: Constraint::Any,
+            },
+        }))
         .await
         .expect("failed to enable udp2tcp");
 
     let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-        location: Some(Constraint::Only(LocationConstraint::Location(
-            GeographicLocationConstraint::Country("se".to_string()),
-        ))),
         tunnel_protocol: Some(Constraint::Only(TunnelType::Wireguard)),
-        wireguard_constraints: Some(WireguardConstraints::default()),
         ..Default::default()
     });
 
@@ -232,23 +223,20 @@ pub async fn test_bridge(
     log::info!("Updating bridge settings");
 
     mullvad_client
-        .set_bridge_state(types::BridgeState {
-            state: i32::from(types::bridge_state::State::On),
-        })
+        .set_bridge_state(types::BridgeState::from(BridgeState::On))
         .await
         .expect("failed to enable bridge mode");
 
     mullvad_client
-        .set_bridge_settings(types::BridgeSettings {
-            r#type: Some(types::bridge_settings::Type::Normal(
-                types::bridge_settings::BridgeConstraints {
-                    location: helpers::into_locationconstraint(&entry)
-                        .map(types::LocationConstraint::from),
-                    providers: vec![],
-                    ownership: i32::from(types::Ownership::Any),
-                },
-            )),
-        })
+        .set_bridge_settings(types::BridgeSettings::from(BridgeSettings::Normal(
+            BridgeConstraints {
+                location: helpers::into_constraint(&entry).expect(&format!(
+                    "Could not resolve location constraint: {:?}",
+                    entry
+                )),
+                ..Default::default()
+            },
+        )))
         .await
         .expect("failed to update bridge settings");
 
@@ -394,9 +382,6 @@ pub async fn test_wireguard_autoconnect(
     log::info!("Setting tunnel protocol to WireGuard");
 
     let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-        location: Some(Constraint::Only(LocationConstraint::Location(
-            GeographicLocationConstraint::Country("se".to_string()),
-        ))),
         tunnel_protocol: Some(Constraint::Only(TunnelType::Wireguard)),
         ..Default::default()
     });
@@ -410,7 +395,7 @@ pub async fn test_wireguard_autoconnect(
         .await
         .expect("failed to enable auto-connect");
 
-    reboot(&mut rpc).await?;
+    helpers::reboot(&mut rpc).await?;
     rpc.mullvad_daemon_wait_for_state(|state| state == ServiceStatus::Running)
         .await?;
 
@@ -439,9 +424,6 @@ pub async fn test_openvpn_autoconnect(
     log::info!("Setting tunnel protocol to OpenVPN");
 
     let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-        location: Some(Constraint::Only(LocationConstraint::Location(
-            GeographicLocationConstraint::Country("se".to_string()),
-        ))),
         tunnel_protocol: Some(Constraint::Only(TunnelType::OpenVpn)),
         ..Default::default()
     });
@@ -455,7 +437,7 @@ pub async fn test_openvpn_autoconnect(
         .await
         .expect("failed to enable auto-connect");
 
-    reboot(&mut rpc).await?;
+    helpers::reboot(&mut rpc).await?;
     rpc.mullvad_daemon_wait_for_state(|state| state == ServiceStatus::Running)
         .await?;
 
@@ -469,19 +451,6 @@ pub async fn test_openvpn_autoconnect(
     Ok(())
 }
 
-async fn reboot(rpc: &mut ServiceClient) -> Result<(), Error> {
-    rpc.reboot().await?;
-
-    // The tunnel must be reconfigured after the virtual machine is up,
-    // or macOS refuses to assign an IP. The reasons for this are poorly understood.
-    #[cfg(target_os = "macos")]
-    crate::vm::network::macos::configure_tunnel()
-        .await
-        .map_err(|error| Error::Other(format!("Failed to recreate custom wg tun: {error}")))?;
-
-    Ok(())
-}
-
 /// Test whether quantum-resistant tunnels can be set up.
 ///
 /// # Limitations
@@ -497,9 +466,9 @@ pub async fn test_quantum_resistant_tunnel(
     mut mullvad_client: ManagementServiceClient,
 ) -> Result<(), Error> {
     mullvad_client
-        .set_quantum_resistant_tunnel(types::QuantumResistantState {
-            state: i32::from(types::quantum_resistant_state::State::Off),
-        })
+        .set_quantum_resistant_tunnel(types::QuantumResistantState::from(
+            wireguard::QuantumResistantState::Off,
+        ))
         .await
         .expect("Failed to disable PQ tunnels");
 
@@ -513,21 +482,17 @@ pub async fn test_quantum_resistant_tunnel(
     log::info!("Setting tunnel protocol to WireGuard");
 
     let relay_settings = RelaySettingsUpdate::Normal(RelayConstraintsUpdate {
-        location: Some(Constraint::Only(LocationConstraint::Location(
-            GeographicLocationConstraint::Country("se".to_string()),
-        ))),
         tunnel_protocol: Some(Constraint::Only(TunnelType::Wireguard)),
         ..Default::default()
     });
-
     update_relay_settings(&mut mullvad_client, relay_settings)
         .await
         .expect("Failed to update relay settings");
 
     mullvad_client
-        .set_quantum_resistant_tunnel(types::QuantumResistantState {
-            state: i32::from(types::quantum_resistant_state::State::On),
-        })
+        .set_quantum_resistant_tunnel(types::QuantumResistantState::from(
+            wireguard::QuantumResistantState::On,
+        ))
         .await
         .expect("Failed to enable PQ tunnels");
 
@@ -581,13 +546,9 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: ManagementServiceClient,
 ) -> Result<(), Error> {
-    use mullvad_types::{
-        relay_constraints::{ObfuscationSettings, SelectedObfuscation, Udp2TcpObfuscationSettings},
-        wireguard::QuantumResistantState,
-    };
     mullvad_client
         .set_quantum_resistant_tunnel(types::QuantumResistantState::from(
-            QuantumResistantState::On,
+            wireguard::QuantumResistantState::On,
         ))
         .await
         .expect("Failed to enable PQ tunnels");