From f46767e69a9b54584b4331aab5844c37300d961d Mon Sep 17 00:00:00 2001
From: Joakim Hulthe <joakim@hulthe.net>
Date: Tue, 13 Aug 2024 15:59:54 +0200
Subject: [PATCH] Pick DAITA use_anywhere relays based on distance to selected
 location

---
 .../src/relay_selector/mod.rs                 | 75 ++++++++++++++-----
 mullvad-types/src/location.rs                 |  5 +-
 2 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 753559f7d92d..5deb5fbefafa 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -756,12 +756,42 @@ impl RelaySelector {
         custom_lists: &CustomListsSettings,
         parsed_relays: &ParsedRelays,
     ) -> Result<WireguardConfig, Error> {
-        // Modify the query to enable multihop
-        let mut query = query.clone();
-        query.wireguard_constraints.use_multihop = Constraint::Only(true);
-        query.wireguard_constraints.entry_location = Constraint::Any; // TODO: smarter location selection
+        let mut exit_relay_query = query.clone();
+
+        // DAITA should only be enabled for the entry relay
+        exit_relay_query.wireguard_constraints.daita = Constraint::Only(false);
+
+        let exit_candidates =
+            filter_matching_relay_list(&exit_relay_query, parsed_relays, custom_lists);
+        let exit = helpers::pick_random_relay(&exit_candidates).ok_or(Error::NoRelay)?;
+
+        // generate a list of potential entry relays, disregarding any location constraint
+        let entry_query = RelayQuery {
+            location: Constraint::Any,
+            ..query.clone()
+        };
+        let mut entry_candidates =
+            filter_matching_relay_list(&entry_query, parsed_relays, custom_lists)
+                .into_iter()
+                .map(|entry| RelayWithDistance::new_with_distance_from(entry, &exit.location))
+                .collect_vec();
+
+        // sort entry relay candidates by distance, and pick one from those that are closest
+        entry_candidates.sort_unstable_by(|a, b| a.distance.total_cmp(&b.distance));
+        let smallest_distance = entry_candidates.first().map(|relay| relay.distance);
+        let smallest_distance = smallest_distance.unwrap_or_default();
+        let entry_candidates = entry_candidates
+            .into_iter()
+            // only consider the relay(s) with the smallest distance. note that the list is sorted.
+            // NOTE: we could relax this requirement, but since so few relays support DAITA
+            // (and this function is only used for daita) we might end up picking relays that are
+            // needlessly far away. Consider making this closure  configurable if needed.
+            .take_while(|relay| relay.distance <= smallest_distance)
+            .map(|relay_with_distance| relay_with_distance.relay)
+            .collect_vec();
+        let entry = pick_random_excluding(&entry_candidates, &exit).ok_or(Error::NoRelay)?;
 
-        Self::get_wireguard_multihop_config(&query, custom_lists, parsed_relays)
+        Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
     }
 
     /// This function selects a valid entry and exit relay to be used in a multihop configuration.
@@ -793,11 +823,6 @@ impl RelaySelector {
         let entry_candidates =
             filter_matching_relay_list(&entry_relay_query, parsed_relays, custom_lists);
 
-        fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
-            list.iter()
-                .filter(|&a| a != exclude)
-                .choose(&mut thread_rng())
-        }
         // We avoid picking the same relay for entry and exit by choosing one and excluding it when
         // choosing the other.
         let (exit, entry) = match (exit_candidates.as_slice(), entry_candidates.as_slice()) {
@@ -1043,19 +1068,10 @@ impl RelaySelector {
         const MIN_BRIDGE_COUNT: usize = 5;
         let location = location.into();
 
-        #[derive(Clone)]
-        struct RelayWithDistance {
-            relay: Relay,
-            distance: f64,
-        }
-
         // Filter out all candidate bridges.
         let matching_bridges: Vec<RelayWithDistance> = relays
             .into_iter()
-            .map(|relay| RelayWithDistance {
-                distance: relay.location.distance_from(&location),
-                relay,
-            })
+            .map(|relay| RelayWithDistance::new_with_distance_from(relay, location))
             .sorted_unstable_by_key(|relay| relay.distance as usize)
             .take(MIN_BRIDGE_COUNT)
             .collect();
@@ -1116,3 +1132,22 @@ impl RelaySelector {
         helpers::pick_random_relay(&candidates).cloned()
     }
 }
+
+fn pick_random_excluding<'a>(list: &'a [Relay], exclude: &'a Relay) -> Option<&'a Relay> {
+    list.iter()
+        .filter(|&a| a != exclude)
+        .choose(&mut thread_rng())
+}
+
+#[derive(Clone)]
+struct RelayWithDistance {
+    distance: f64,
+    relay: Relay,
+}
+
+impl RelayWithDistance {
+    fn new_with_distance_from(relay: Relay, from: impl Into<Coordinates>) -> Self {
+        let distance = relay.location.distance_from(from);
+        RelayWithDistance { relay, distance }
+    }
+}
diff --git a/mullvad-types/src/location.rs b/mullvad-types/src/location.rs
index 7807b65d6a91..d8a47176c190 100644
--- a/mullvad-types/src/location.rs
+++ b/mullvad-types/src/location.rs
@@ -19,7 +19,8 @@ pub struct Location {
 const RAIDUS_OF_EARTH: f64 = 6372.8;
 
 impl Location {
-    pub fn distance_from(&self, other: &Coordinates) -> f64 {
+    pub fn distance_from(&self, other: impl Into<Coordinates>) -> f64 {
+        let other: Coordinates = other.into();
         haversine_dist_deg(
             self.latitude,
             self.longitude,
@@ -33,7 +34,7 @@ impl Location {
     }
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub struct Coordinates {
     pub latitude: f64,
     pub longitude: f64,