Introduce `Singlehop` and `Multihop` as distinct types

4 files changed, 111 insertions, 66 deletions

diff --git a/mullvad-relay-selector/src/error.rs b/mullvad-relay-selector/src/error.rs
index dea6f81c3e..0bd5fad5d4 100644
--- a/mullvad-relay-selector/src/error.rs
+++ b/mullvad-relay-selector/src/error.rs
@@ -3,7 +3,8 @@
 
 use mullvad_types::{relay_constraints::MissingCustomBridgeSettings, relay_list::Relay};
 
-use crate::{detailer, WireguardConfig};
+use crate::detailer;
+use crate::relay_selector::relays::WireguardConfig;
 
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
index 09708a059e..cbd7c94a36 100644
--- a/mullvad-relay-selector/src/lib.rs
+++ b/mullvad-relay-selector/src/lib.rs
@@ -8,8 +8,10 @@ mod relay_selector;
 
 // Re-exports
 pub use error::Error;
+pub use relay_selector::detailer;
+pub use relay_selector::query;
+pub use relay_selector::relays::WireguardConfig;
 pub use relay_selector::{
-    detailer, query, AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay,
-    RelaySelector, RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig,
-    WireguardConfig, RETRY_ORDER,
+    AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay, RelaySelector,
+    RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig, RETRY_ORDER,
 };
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 7f7517923c..3cdfa58f26 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -5,9 +5,11 @@ mod helpers;
 mod matcher;
 mod parsed_relays;
 pub mod query;
+pub mod relays;
 
 use matcher::{filter_matching_bridges, filter_matching_relay_list};
 use parsed_relays::ParsedRelays;
+use relays::{Multihop, Singlehop, WireguardConfig};
 
 use crate::detailer::{openvpn_endpoint, wireguard_endpoint};
 use crate::error::{EndpointErrorDetails, Error};
@@ -215,50 +217,6 @@ pub enum GetRelay {
     Custom(CustomTunnelEndpoint),
 }
 
-// TODO: Import `Either` and convert `Multihop` and `Singlehop` into concrete types.
-/// This struct defines the different Wireguard relays the the relay selector can end up selecting
-/// for an arbitrary Wireguard [`query`].
-///
-/// - [`WireguardConfig::Singlehop`]; A normal wireguard relay where VPN traffic enters and exits
-///   through this sole relay.
-/// - [`WireguardConfig::Multihop`]; Two wireguard relays to be used in a multihop circuit. VPN
-///   traffic will enter through `entry` and eventually come out from `exit` before the traffic will
-///   actually be routed to the broader internet.
-#[derive(Clone, Debug)]
-pub enum WireguardConfig {
-    /// Strongly prefer to instantiate this variant using [`WireguardConfig::singlehop`] as that
-    /// will assert that the relay is of the expected type.
-    Singlehop { exit: Relay },
-    /// Strongly prefer to instantiate this variant using [`WireguardConfig::multihop`] as that
-    /// will assert that the entry & exit relays are of the expected type.
-    Multihop { exit: Relay, entry: Relay },
-}
-
-impl WireguardConfig {
-    const fn singlehop(exit: Relay) -> Self {
-        // FIXME: This assert would be better to encode at the type level.
-        assert!(matches!(
-            exit.endpoint_data,
-            RelayEndpointData::Wireguard(_)
-        ));
-        Self::Singlehop { exit }
-    }
-
-    const fn multihop(exit: Relay, entry: Relay) -> Self {
-        // FIXME: This assert would be better to encode at the type level.
-        assert!(matches!(
-            exit.endpoint_data,
-            RelayEndpointData::Wireguard(_)
-        ));
-        // FIXME: This assert would be better to encode at the type level.
-        assert!(matches!(
-            entry.endpoint_data,
-            RelayEndpointData::Wireguard(_)
-        ));
-        Self::Multihop { exit, entry }
-    }
-}
-
 #[derive(Clone, Debug)]
 pub enum SelectedBridge {
     Normal { settings: CustomProxy, relay: Relay },
@@ -693,9 +651,17 @@ impl RelaySelector {
             Constraint::Only(TunnelType::Wireguard)
         );
         let inner = if !query.wireguard_constraints().multihop() {
-            Self::get_wireguard_singlehop_config(query, custom_lists, parsed_relays)?
+            WireguardConfig::from(Self::get_wireguard_singlehop_config(
+                query,
+                custom_lists,
+                parsed_relays,
+            )?)
         } else {
-            Self::get_wireguard_multihop_config(query, custom_lists, parsed_relays)?
+            WireguardConfig::from(Self::get_wireguard_multihop_config(
+                query,
+                custom_lists,
+                parsed_relays,
+            )?)
         };
         let endpoint = Self::get_wireguard_endpoint(query, parsed_relays, &inner)?;
         let obfuscator =
@@ -717,21 +683,23 @@ impl RelaySelector {
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
-    ) -> Result<WireguardConfig, Error> {
+    ) -> Result<Singlehop, Error> {
         let candidates = filter_matching_relay_list(query, parsed_relays, custom_lists);
 
-        if let Some(x) = Self::select_daita_multihop_config_if_necessary(
-            query,
-            custom_lists,
-            parsed_relays,
-            &candidates,
-        )? {
-            return Ok(x);
-        }
+        // // TODO: Move somewhere
+        // // Or keep here, but return either.
+        // if let Some(x) = Self::select_daita_multihop_config_if_necessary(
+        //     query,
+        //     custom_lists,
+        //     parsed_relays,
+        //     &candidates,
+        // )? {
+        //     return Ok(x);
+        // }
 
         helpers::pick_random_relay(&candidates)
             .cloned()
-            .map(WireguardConfig::singlehop)
+            .map(Singlehop::new)
             .ok_or(Error::NoRelay)
     }
 
@@ -741,7 +709,7 @@ impl RelaySelector {
         parsed_relays: &ParsedRelays,
         // TODO: What is this??
         candidates: &[Relay],
-    ) -> Result<Option<WireguardConfig>, Error> {
+    ) -> Result<Option<Multihop>, Error> {
         // are we using daita?
         let using_daita = || query.wireguard_constraints().daita == Constraint::Only(true);
 
@@ -767,6 +735,8 @@ impl RelaySelector {
         // if we found no matching relays because DAITA was enabled, and `use_multihop_if_necessary`
         // is enabled, try enabling multihop and connecting using an automatically selected
         // entry relay.
+        //
+        // TODO: I think this should be pushed up one level and not be in this function scope.
         if candidates.is_empty()
             && using_daita()
             && no_relay_because_daita()?
@@ -782,12 +752,12 @@ impl RelaySelector {
     ///
     /// # Returns
     /// * An `Err` if no entry/exit relay can be chosen
-    /// * `Ok(WireguardConfig::Multihop)` otherwise
+    /// * `Ok(Multihop)` otherwise
     fn get_wireguard_auto_multihop_config(
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
-    ) -> Result<WireguardConfig, Error> {
+    ) -> Result<Multihop, Error> {
         let mut exit_relay_query = query.clone();
 
         // DAITA should only be enabled for the entry relay
@@ -824,7 +794,7 @@ impl RelaySelector {
         let entry =
             helpers::pick_random_relay_excluding(&entry_candidates, exit).ok_or(Error::NoRelay)?;
 
-        Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
+        Ok(Multihop::new(entry.clone(), exit.clone()))
     }
 
     /// This function selects a valid entry and exit relay to be used in a multihop configuration.
@@ -838,7 +808,7 @@ impl RelaySelector {
         query: &RelayQuery,
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
-    ) -> Result<WireguardConfig, Error> {
+    ) -> Result<Multihop, Error> {
         // Here, we modify the original query just a bit.
         // The actual query for an exit relay is identical as for an exit relay, with the
         // exception that the location is different. It is simply the location as dictated by
@@ -878,7 +848,7 @@ impl RelaySelector {
         }
         .ok_or(Error::NoRelay)?;
 
-        Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
+        Ok(Multihop::new(entry.clone(), exit.clone()))
     }
 
     /// Constructs a [`MullvadEndpoint`] with details for how to connect to `relay`.
diff --git a/mullvad-relay-selector/src/relay_selector/relays.rs b/mullvad-relay-selector/src/relay_selector/relays.rs
new file mode 100644
index 0000000000..e8a1035c1d
--- /dev/null
+++ b/mullvad-relay-selector/src/relay_selector/relays.rs
@@ -0,0 +1,72 @@
+//! TODO: Document this module
+
+use mullvad_types::relay_list::{Relay, RelayEndpointData};
+
+// TODO: Import `Either` and convert `Multihop` and `Singlehop` into concrete types.
+/// This struct defines the different Wireguard relays the the relay selector can end up selecting
+/// for an arbitrary Wireguard [`query`].
+///
+/// - [`WireguardConfig::Singlehop`]; A normal wireguard relay where VPN traffic enters and exits
+///   through this sole relay.
+/// - [`WireguardConfig::Multihop`]; Two wireguard relays to be used in a multihop circuit. VPN
+///   traffic will enter through `entry` and eventually come out from `exit` before the traffic will
+///   actually be routed to the broader internet.
+#[derive(Clone, Debug)]
+pub enum WireguardConfig {
+    /// Strongly prefer to instantiate this variant using [`WireguardConfig::singlehop`] as that
+    /// will assert that the relay is of the expected type.
+    Singlehop { exit: Relay },
+    /// Strongly prefer to instantiate this variant using [`WireguardConfig::multihop`] as that
+    /// will assert that the entry & exit relays are of the expected type.
+    Multihop { exit: Relay, entry: Relay },
+}
+
+/// TODO: Document
+pub struct Singlehop(Relay);
+/// TODO: Document
+pub struct Multihop {
+    entry: Relay,
+    exit: Relay,
+}
+
+impl From<Singlehop> for WireguardConfig {
+    fn from(relay: Singlehop) -> Self {
+        Self::Singlehop { exit: relay.0 }
+    }
+}
+
+impl From<Multihop> for WireguardConfig {
+    fn from(relay: Multihop) -> Self {
+        WireguardConfig::Multihop {
+            exit: relay.exit,
+            entry: relay.entry,
+        }
+    }
+}
+
+impl Singlehop {
+    pub const fn new(exit: Relay) -> Self {
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            exit.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        Self(exit)
+    }
+}
+
+impl Multihop {
+    pub const fn new(entry: Relay, exit: Relay) -> Self {
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            exit.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            entry.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        Multihop { exit, entry }
+    }
+}