Distinguise different types of endpoints in the GetRelay type

Distinguise different types of endpoints in the `GetRelay` type for Wireguard and OpenVPN, instead of bundling different endpoint-like types into the same enum and using specialized `unwrap_this_type` functions.
mullvad · Mar 12, 2024 · 34bc19f · 34bc19f
1 parent fe0eafc
commit 34bc19f
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 85 deletions.
diff --git a/mullvad-daemon/src/tunnel.rs b/mullvad-daemon/src/tunnel.rs
@@ -176,12 +176,7 @@ impl InnerParametersGenerator {
                     bridge: bridge_relay,
                 });
 
-                // TODO(markus): Remodel `GetRelay` to not require this unwrap.
-                Ok(self.create_openvpn_tunnel_parameters(
-                    *endpoint.unwrap_openvpn(),
-                    data,
-                    bridge_settings,
-                ))
+                Ok(self.create_openvpn_tunnel_parameters(endpoint, data, bridge_settings))
             }
             GetRelay::Wireguard {
                 endpoint,
@@ -204,7 +199,7 @@ impl InnerParametersGenerator {
                 });
 
                 Ok(self.create_wireguard_tunnel_parameters(
-                    endpoint.unwrap_wireguard().clone(),
+                    endpoint.clone(),
                     data,
                     obfuscator_config,
                 ))

diff --git a/mullvad-relay-selector/src/relay_selector/detailer.rs b/mullvad-relay-selector/src/relay_selector/detailer.rs
@@ -10,7 +10,7 @@ use std::net::{IpAddr, SocketAddr, SocketAddrV4};
 use ipnetwork::IpNetwork;
 use mullvad_types::{
     constraints::Constraint,
-    endpoint::{MullvadEndpoint, MullvadWireguardEndpoint},
+    endpoint::MullvadWireguardEndpoint,
     relay_constraints::TransportPort,
     relay_list::{OpenVpnEndpoint, OpenVpnEndpointData, Relay, WireguardEndpointData},
 };
@@ -25,7 +25,7 @@ use super::{
 
 /// Given a Wireguad relay (and optionally an entry relay if multihop is used) and the original
 /// query the relay selector used, fill in all connection details to produce a valid
-/// [`MullvadEndpoint`] by calling [`to_endpoint`].
+/// [`MullvadWireguardEndpoint`] by calling [`to_endpoint`].
 ///
 /// [`to_endpoint`]: WireguardDetailer::to_endpoint
 pub struct WireguardDetailer {
@@ -52,7 +52,7 @@ impl WireguardDetailer {
         }
     }
 
-    /// Constructs a `MullvadEndpoint` with details for a Wireguard circuit.
+    /// Constructs a [`MullvadWireguardEndpoint`] with details for a Wireguard circuit.
     ///
     /// If entry is `None`, `to_endpoint` configures a single-hop connection using the exit relay data.
     /// Otherwise, it constructs a multihop setup using both entry and exit to set up appropriate peer
@@ -62,15 +62,15 @@ impl WireguardDetailer {
     /// - A configured Mullvad endpoint for Wireguard, encapsulating either a single-hop or multi-hop connection setup.
     /// - Returns `None` if the desired port is not in a valid port range (see
     /// [`WireguradRelayQuery::port`]) or relay addresses cannot be resolved.
-    pub fn to_endpoint(&self) -> Option<MullvadEndpoint> {
+    pub fn to_endpoint(&self) -> Option<MullvadWireguardEndpoint> {
         match &self.config {
             WireguardConfig::Singlehop { exit } => self.to_singlehop_endpoint(exit),
             WireguardConfig::Multihop { exit, entry } => self.to_multihop_endpoint(exit, entry),
         }
     }
 
     /// Configure a single-hop connection using the exit relay data.
-    fn to_singlehop_endpoint(&self, exit: &Relay) -> Option<MullvadEndpoint> {
+    fn to_singlehop_endpoint(&self, exit: &Relay) -> Option<MullvadWireguardEndpoint> {
         let endpoint = {
             let host = self.get_address_for_wireguard_relay(exit)?;
             let port = self.get_port_for_wireguard_relay(&self.data)?;
@@ -83,20 +83,24 @@ impl WireguardDetailer {
             // This will be filled in later, not the relay selector's problem
             psk: None,
         };
-        Some(MullvadEndpoint::Wireguard(MullvadWireguardEndpoint {
+        Some(MullvadWireguardEndpoint {
             peer: peer_config,
             exit_peer: None,
             ipv4_gateway: self.data.ipv4_gateway,
             ipv6_gateway: self.data.ipv6_gateway,
-        }))
+        })
     }
 
     /// Configure a multihop connection using the entry & exit relay data.
     ///
     /// # Note
     /// In a multihop circuit, we need to provide an exit peer configuration in addition to the
     /// peer configuration.
-    fn to_multihop_endpoint(&self, exit: &Relay, entry: &Relay) -> Option<MullvadEndpoint> {
+    fn to_multihop_endpoint(
+        &self,
+        exit: &Relay,
+        entry: &Relay,
+    ) -> Option<MullvadWireguardEndpoint> {
         let exit_endpoint = {
             let ip = exit.ipv4_addr_in;
             // The port that the exit relay listens for incoming connections from entry
@@ -133,12 +137,12 @@ impl WireguardDetailer {
             psk: None,
         };
 
-        Some(MullvadEndpoint::Wireguard(MullvadWireguardEndpoint {
+        Some(MullvadWireguardEndpoint {
             peer: entry,
             exit_peer: Some(exit),
             ipv4_gateway: self.data.ipv4_gateway,
             ipv6_gateway: self.data.ipv6_gateway,
-        }))
+        })
     }
 
     /// Get the correct IP address for the given relay.
@@ -208,8 +212,8 @@ impl WireguardDetailer {
 }
 
 /// Given an OpenVPN relay and the original query the relay selector used,
-/// fill in all connection details to produce a valid [`MullvadEndpoint`]
-/// by calling [`to_endpoint`].
+/// fill in all connection details to produce a valid [`Endpoint`] by
+/// calling [`to_endpoint`].
 ///
 /// [`to_endpoint`]: OpenVpnDetailer::to_endpoint
 pub struct OpenVpnDetailer {
@@ -228,14 +232,14 @@ impl OpenVpnDetailer {
         }
     }
 
-    /// Map `self` to a [`MullvadEndpoint`].
+    /// Map `self` to a [`Endpoint`].
     ///
     /// If this endpoint is to be used in conjunction with a bridge, the resulting endpoint is
     /// guaranteed to use transport protocol `TCP`.
     ///
     /// This function can fail if no valid port + transport protocol combination is found.
     /// See [`OpenVpnEndpointData`] for more details.
-    pub fn to_endpoint(&self) -> Option<MullvadEndpoint> {
+    pub fn to_endpoint(&self) -> Option<Endpoint> {
         // If `bridge_mode` is true, this function may only return endpoints which use TCP, not UDP.
         if BridgeQuery::should_use_bridge(&self.openvpn_constraints.bridge_settings) {
             self.to_bridged_endpoint()
@@ -245,29 +249,23 @@ impl OpenVpnDetailer {
     }
 
     /// Configure a single-hop connection using the exit relay data.
-    fn to_singlehop_endpoint(&self) -> Option<MullvadEndpoint> {
+    fn to_singlehop_endpoint(&self) -> Option<Endpoint> {
         use rand::seq::IteratorRandom;
         let constraints_port = self.openvpn_constraints.port;
         self.data
             .ports
             .iter()
             .filter(|&endpoint| Self::compatible_port_combo(&constraints_port, endpoint))
             .choose(&mut rand::thread_rng())
-            .map(|endpoint| {
-                MullvadEndpoint::OpenVpn(Endpoint::new(
-                    self.exit.ipv4_addr_in,
-                    endpoint.port,
-                    endpoint.protocol,
-                ))
-            })
+            .map(|endpoint| Endpoint::new(self.exit.ipv4_addr_in, endpoint.port, endpoint.protocol))
     }
 
     /// Configure an endpoint that will be used together with a bridge.
     ///
     /// # Note
     /// In bridge mode, the only viable transport protocol is TCP. Otherwise, this function is
     /// identical to [`Self::to_singlehop_endpoint`].
-    fn to_bridged_endpoint(&self) -> Option<MullvadEndpoint> {
+    fn to_bridged_endpoint(&self) -> Option<Endpoint> {
         use rand::seq::IteratorRandom;
         let constraints_port = self.openvpn_constraints.port;
         self.data
@@ -276,13 +274,7 @@ impl OpenVpnDetailer {
             .filter(|endpoint| matches!(endpoint.protocol, TransportProtocol::Tcp))
             .filter(|endpoint| Self::compatible_port_combo(&constraints_port, endpoint))
             .choose(&mut rand::thread_rng())
-            .map(|endpoint| {
-                MullvadEndpoint::OpenVpn(Endpoint::new(
-                    self.exit.ipv4_addr_in,
-                    endpoint.port,
-                    endpoint.protocol,
-                ))
-            })
+            .map(|endpoint| Endpoint::new(self.exit.ipv4_addr_in, endpoint.port, endpoint.protocol))
     }
 
     /// Returns true if `port_constraint` can be used to connect to `endpoint`.

diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -19,7 +19,7 @@ use matcher::{BridgeMatcher, RelayMatcher, WireguardMatcher};
 use mullvad_types::{
     constraints::Constraint,
     custom_list::CustomListsSettings,
-    endpoint::{MullvadEndpoint, MullvadWireguardEndpoint},
+    endpoint::MullvadWireguardEndpoint,
     location::{Coordinates, Location},
     relay_constraints::{
         BridgeSettings, BridgeState, InternalBridgeConstraints, ObfuscationSettings,
@@ -31,7 +31,9 @@ use mullvad_types::{
     CustomTunnelEndpoint,
 };
 use talpid_types::{
-    net::{obfuscation::ObfuscatorConfig, proxy::CustomProxy, TransportProtocol, TunnelType},
+    net::{
+        obfuscation::ObfuscatorConfig, proxy::CustomProxy, Endpoint, TransportProtocol, TunnelType,
+    },
     ErrorExt,
 };
 
@@ -107,13 +109,13 @@ pub struct SelectorConfig {
 #[derive(Clone, Debug)]
 pub enum GetRelay {
     Wireguard {
-        endpoint: MullvadEndpoint,
+        endpoint: MullvadWireguardEndpoint,
         obfuscator: Option<SelectedObfuscator>,
         inner: WireguardConfig,
     },
     #[cfg(not(target_os = "android"))]
     OpenVpn {
-        endpoint: MullvadEndpoint,
+        endpoint: Endpoint,
         exit: Relay,
         bridge: Option<SelectedBridge>,
     },
@@ -442,14 +444,8 @@ impl RelaySelector {
             Self::get_wireguard_multihop_config(query, config, parsed_relays)?
         };
         let endpoint = Self::get_wireguard_endpoint(query, inner.clone(), parsed_relays)?;
-        let obfuscator = Self::get_wireguard_obfuscator(
-            query,
-            inner.clone(),
-            // Note: It should always be safe to call `unwrap_wireguard` here, unless there
-            // is a bug in `get_wireguard_endpoint`.
-            endpoint.unwrap_wireguard(),
-            parsed_relays,
-        );
+        let obfuscator =
+            Self::get_wireguard_obfuscator(query, inner.clone(), &endpoint, parsed_relays);
 
         Ok(GetRelay::Wireguard {
             endpoint,
@@ -538,7 +534,7 @@ impl RelaySelector {
         query: &RelayQuery,
         relay: WireguardConfig,
         parsed_relays: &ParsedRelays,
-    ) -> Result<MullvadEndpoint, Error> {
+    ) -> Result<MullvadWireguardEndpoint, Error> {
         WireguardDetailer::new(
             query.wireguard_constraints.clone(),
             relay,
@@ -588,15 +584,8 @@ impl RelaySelector {
     ) -> Result<GetRelay, Error> {
         let exit = Self::choose_relay(query, config, parsed_relays).ok_or(Error::NoRelay)?;
         let endpoint = Self::get_openvpn_endpoint(query, exit.clone(), parsed_relays)?;
-        let bridge = Self::get_openvpn_bridge(
-            query,
-            &exit,
-            // Note: It should always be safe to call `unwrap_openvpn` here, unless there
-            // is a bug in `OpenVpnDetailer::to_endpoint`.
-            &endpoint.unwrap_openvpn().protocol,
-            parsed_relays,
-            config,
-        )?;
+        let bridge =
+            Self::get_openvpn_bridge(query, &exit, &endpoint.protocol, parsed_relays, config)?;
 
         Ok(GetRelay::OpenVpn {
             endpoint,
@@ -611,7 +600,7 @@ impl RelaySelector {
         query: &RelayQuery,
         relay: Relay,
         parsed_relays: &ParsedRelays,
-    ) -> Result<MullvadEndpoint, Error> {
+    ) -> Result<Endpoint, Error> {
         OpenVpnDetailer::new(
             query.openvpn_constraints.clone(),
             relay,

diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
@@ -180,8 +180,8 @@ fn unwrap_relay(get_result: GetRelay) -> Relay {
 
 fn unwrap_endpoint(get_result: GetRelay) -> MullvadEndpoint {
     match get_result {
-        GetRelay::Wireguard { endpoint, .. } => endpoint,
-        GetRelay::OpenVpn { endpoint, .. } => endpoint,
+        GetRelay::Wireguard { endpoint, .. } => MullvadEndpoint::Wireguard(endpoint),
+        GetRelay::OpenVpn { endpoint, .. } => MullvadEndpoint::OpenVpn(endpoint),
         GetRelay::Custom(custom) => {
             panic!("Can not extract Mullvad endpoint from custom relay: {custom}")
         }
@@ -463,8 +463,7 @@ fn test_openvpn_constraints() {
             } else {
                 match relay.expect("Expected to find a relay") {
                     GetRelay::OpenVpn { endpoint, .. } =>  {
-                        assert!(matches!(endpoint, MullvadEndpoint::OpenVpn(_)));
-                        matches_constraints(endpoint.to_endpoint(), &query.openvpn_constraints);
+                        matches_constraints(endpoint, &query.openvpn_constraints);
                     },
                     wrong_relay => panic!("Relay selector should have picked an OpenVPN relay, instead chose {wrong_relay:?}")
                 };
@@ -501,18 +500,9 @@ fn test_selecting_any_relay_will_consider_multihop() {
 
     for _ in 0..100 {
         let relay = relay_selector.get_relay_by_query(query.clone()).unwrap();
-        match relay {
-            GetRelay::Wireguard {
-                endpoint,
-                inner: WireguardConfig::Multihop { .. },
-                ..
-            } => {
-                assert!(matches!(endpoint, MullvadEndpoint::Wireguard(_)));
-            }
-            wrong_relay => panic!(
-            "Relay selector should have picked a Wireguard relay, instead chose {wrong_relay:?}"
-        ),
-        }
+        assert!(matches!(relay, GetRelay::Wireguard { inner: WireguardConfig::Multihop { .. }, .. }),
+            "Relay selector should have picked a Wireguard relay with multihop, instead chose {relay:?}"
+        );
     }
 }
 
@@ -557,13 +547,8 @@ fn test_selecting_wireguard_endpoint_with_auto_obfuscation() {
     for _ in 0..100 {
         let relay = relay_selector.get_relay_by_query(query.clone()).unwrap();
         match relay {
-            GetRelay::Wireguard {
-                endpoint,
-                obfuscator,
-                ..
-            } => {
+            GetRelay::Wireguard { obfuscator, .. } => {
                 assert!(obfuscator.is_none());
-                assert!(matches!(endpoint, MullvadEndpoint::Wireguard { .. }));
             }
             wrong_relay => panic!(
             "Relay selector should have picked a Wireguard relay, instead chose {wrong_relay:?}"
@@ -871,7 +856,7 @@ fn openvpn_handle_bridge_settings() {
     // Assert that the resulting relay uses UDP.
     match relay {
         GetRelay::OpenVpn { endpoint, .. } => {
-            assert_eq!(endpoint.unwrap_openvpn().protocol, UDP);
+            assert_eq!(endpoint.protocol, UDP);
         }
         wrong_relay => panic!(
             "Relay selector should have picked an OpenVPN relay, instead chose {wrong_relay:?}"
@@ -895,7 +880,7 @@ fn openvpn_handle_bridge_settings() {
             endpoint, bridge, ..
         } => {
             assert!(bridge.is_some());
-            assert_eq!(endpoint.unwrap_openvpn().protocol, TCP);
+            assert_eq!(endpoint.protocol, TCP);
         }
         wrong_relay => panic!(
             "Relay selector should have picked an OpenVPN relay, instead chose {wrong_relay:?}"
@@ -928,7 +913,7 @@ fn openvpn_bridge_with_automatic_transport_protocol() {
         // auto.
         match relay {
             GetRelay::OpenVpn { endpoint, .. } => {
-                assert_eq!(endpoint.unwrap_openvpn().protocol, TCP);
+                assert_eq!(endpoint.protocol, TCP);
             }
             wrong_relay => panic!(
                 "Relay selector should have picked an OpenVPN relay, instead chose {wrong_relay:?}"