Use query for shadowsocks test

mullvad · Aug 26, 2024 · 76aa45e · 76aa45e
1 parent 53d1294
commit 76aa45e
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 42 deletions.
diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -33,8 +33,8 @@ use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
         BridgeConstraints, LocationConstraint, ObfuscationSettings, OpenVpnConstraints, Ownership,
-        Providers, RelayConstraints, RelaySettings, SelectedObfuscation, ShadowsocksSettings,
-        TransportPort, Udp2TcpObfuscationSettings, WireguardConstraints,
+        Providers, RelayConstraints, SelectedObfuscation, ShadowsocksSettings, TransportPort,
+        Udp2TcpObfuscationSettings, WireguardConstraints,
     },
     Intersection,
 };
@@ -107,31 +107,48 @@ impl RelayQuery {
             openvpn_constraints: OpenVpnRelayQuery::new(),
         }
     }
-}
 
-impl Default for RelayQuery {
-    fn default() -> Self {
-        Self::new()
+    /// The mapping from [`RelayQuery`] to [`RelayConstraints`]. Note that this does not contain
+    /// obfuscation or bridge settings.
+    pub fn into_relay_constraints(self) -> RelayConstraints {
+        RelayConstraints {
+            location: self.location,
+            providers: self.providers,
+            ownership: self.ownership,
+            tunnel_protocol: self.tunnel_protocol,
+            wireguard_constraints: WireguardConstraints::from(self.wireguard_constraints),
+            openvpn_constraints: OpenVpnConstraints::from(self.openvpn_constraints),
+        }
     }
-}
 
-impl From<RelayQuery> for RelayConstraints {
-    /// The mapping from [`RelayQuery`] to [`RelayConstraints`].
-    fn from(value: RelayQuery) -> Self {
-        RelayConstraints {
-            location: value.location,
-            providers: value.providers,
-            ownership: value.ownership,
-            tunnel_protocol: value.tunnel_protocol,
-            wireguard_constraints: WireguardConstraints::from(value.wireguard_constraints),
-            openvpn_constraints: OpenVpnConstraints::from(value.openvpn_constraints),
+    /// The mapping from [`RelayQuery`] to [`ObfuscationSettings`]
+    pub fn into_obfuscation_settings(self) -> ObfuscationSettings {
+        match self.wireguard_constraints.obfuscation {
+            ObfuscationQuery::Auto => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Auto,
+                ..Default::default()
+            },
+            ObfuscationQuery::Off => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Off,
+                ..Default::default()
+            },
+            ObfuscationQuery::Udp2tcp(settings) => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Udp2Tcp,
+                udp2tcp: settings,
+                ..Default::default()
+            },
+            ObfuscationQuery::Shadowsocks(settings) => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Shadowsocks,
+                shadowsocks: settings,
+                ..Default::default()
+            },
         }
     }
 }
 
-impl From<RelayQuery> for RelaySettings {
-    fn from(value: RelayQuery) -> Self {
-        RelayConstraints::from(value).into()
+impl Default for RelayQuery {
+    fn default() -> Self {
+        Self::new()
     }
 }
 
@@ -401,10 +418,6 @@ pub mod builder {
         pub fn build(self) -> RelayQuery {
             self.query
         }
-
-        pub fn into_constraint(self) -> RelayConstraints {
-            RelayConstraints::from(self.build())
-        }
     }
 
     impl RelayQueryBuilder<Any> {

diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
@@ -1516,7 +1516,8 @@ fn valid_user_setting_should_yield_relay() {
     let user_query = RelayQueryBuilder::new().location(location.clone()).build();
     let user_constraints = RelayQueryBuilder::new()
         .location(location.clone())
-        .into_constraint();
+        .build()
+        .into_relay_constraints();
 
     let config = SelectorConfig {
         relay_settings: user_constraints.into(),

diff --git a/mullvad-types/src/constraints/constraint.rs b/mullvad-types/src/constraints/constraint.rs
@@ -41,6 +41,15 @@ impl<T: fmt::Display> fmt::Display for Constraint<T> {
     }
 }
 
+impl<T: Default> Constraint<T> {
+    pub fn unwrap_or_default(self) -> T {
+        match self {
+            Constraint::Any => Default::default(),
+            Constraint::Only(value) => value,
+        }
+    }
+}
+
 impl<T> Constraint<T> {
     pub fn unwrap(self) -> T {
         match self {

diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
@@ -211,23 +211,14 @@ pub async fn test_wireguard_over_shadowsocks(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
+    let query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
+
     mullvad_client
-        .set_obfuscation_settings(relay_constraints::ObfuscationSettings {
-            selected_obfuscation: SelectedObfuscation::Shadowsocks,
-            shadowsocks: ShadowsocksSettings {
-                port: Constraint::Any,
-            },
-            ..Default::default()
-        })
+        .set_obfuscation_settings(query.clone().into_obfuscation_settings())
         .await
         .context("Failed to enable shadowsocks")?;
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
+    set_relay_settings(&mut mullvad_client, query.into_relay_constraints())
         .await
         .context("Failed to update relay settings")?;
 
@@ -344,7 +335,8 @@ pub async fn test_multihop(
     let relay_constraints = RelayQueryBuilder::new()
         .wireguard()
         .multihop()
-        .into_constraint();
+        .build()
+        .into_relay_constraints();
 
     set_relay_settings(
         &mut mullvad_client,
@@ -460,7 +452,7 @@ pub async fn test_daita(
 
     set_relay_settings(
         &mut mullvad_client,
-        RelayQueryBuilder::new().wireguard().build(),
+        RelayQueryBuilder::new().wireguard().build().into_relay_constraints(),
     )
     .await?;
 
@@ -630,7 +622,8 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
     let relay_constraints = RelayQueryBuilder::new()
         .wireguard()
         .multihop()
-        .into_constraint();
+        .build()
+        .into_relay_constraints();
 
     mullvad_client
         .set_relay_settings(RelaySettings::Normal(relay_constraints))
@@ -678,7 +671,8 @@ pub async fn test_quantum_resistant_multihop_shadowsocks_tunnel(
     let relay_constraints = RelayQueryBuilder::new()
         .wireguard()
         .multihop()
-        .into_constraint();
+        .build()
+        .into_relay_constraints();
 
     mullvad_client
         .set_relay_settings(RelaySettings::Normal(relay_constraints))