Add location midpoint function

2 files changed, 89 insertions, 3 deletions

diff --git a/mullvad-daemon/src/relays/mod.rs b/mullvad-daemon/src/relays/mod.rs
index 43bd2a697a..d1c34adb70 100644
--- a/mullvad-daemon/src/relays/mod.rs
+++ b/mullvad-daemon/src/relays/mod.rs
@@ -6,7 +6,7 @@ use ipnetwork::IpNetwork;
 use mullvad_rpc::{availability::ApiAvailabilityHandle, rest::MullvadRestHandle};
 use mullvad_types::{
     endpoint::{MullvadEndpoint, MullvadWireguardEndpoint},
-    location::Location,
+    location::{Coordinates, Location},
     relay_constraints::{
         BridgeState, Constraint, InternalBridgeConstraints, LocationConstraint, Match,
         OpenVpnConstraints, Providers, RelayConstraints, Set, TransportPort, WireguardConstraints,
@@ -302,6 +302,33 @@ impl RelaySelector {
         }
     }
 
+    /// Returns the average location of relays that match the given constraints.
+    /// This returns none if the location is `any` or if no relays match the constraints.
+    pub fn get_relay_midpoint(&self, relay_constraints: &RelayConstraints) -> Option<Coordinates> {
+        if relay_constraints.location.is_any() {
+            return None;
+        }
+
+        let matcher = RelayMatcher::from(relay_constraints.clone());
+        let matching_locations: Vec<Location> = self
+            .parsed_relays
+            .lock()
+            .relays()
+            .iter()
+            .filter(|relay| relay.active)
+            .filter_map(|relay| {
+                matcher
+                    .filter_matching_relay(relay)
+                    .and_then(|relay| relay.location)
+            })
+            .collect();
+
+        if matching_locations.is_empty() {
+            return None;
+        }
+        Some(Coordinates::midpoint(&matching_locations))
+    }
+
     /// Returns an OpenVpn endpoint, should only ever be used when the user has specified the tunnel
     /// protocol as only OpenVPN.
     fn get_openvpn_endpoint(
@@ -683,7 +710,7 @@ impl RelaySelector {
 
         if let Some(location) = location {
             matching_relays.sort_by_cached_key(|relay| {
-                (relay.location.as_ref().unwrap().distance_from(&location) * 1000.0) as i64
+                (relay.location.as_ref().unwrap().distance_from(location) * 1000.0) as i64
             });
         }
         matching_relays.get(0).and_then(|relay| {
diff --git a/mullvad-types/src/location.rs b/mullvad-types/src/location.rs
index b9c460fa87..9bf46cb586 100644
--- a/mullvad-types/src/location.rs
+++ b/mullvad-types/src/location.rs
@@ -21,7 +21,7 @@ pub struct Location {
 const RAIDUS_OF_EARTH: f64 = 6372.8;
 
 impl Location {
-    pub fn distance_from(&self, other: &Location) -> f64 {
+    pub fn distance_from(&self, other: &Coordinates) -> f64 {
         haversine_dist_deg(
             self.latitude,
             self.longitude,
@@ -31,6 +31,65 @@ impl Location {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct Coordinates {
+    pub latitude: f64,
+    pub longitude: f64,
+}
+
+impl From<&Location> for Coordinates {
+    fn from(location: &Location) -> Self {
+        Self {
+            latitude: location.latitude,
+            longitude: location.longitude,
+        }
+    }
+}
+
+impl From<Location> for Coordinates {
+    fn from(location: Location) -> Self {
+        Coordinates::from(&location)
+    }
+}
+
+impl Coordinates {
+    /// Computes the approximate midpoint of a set of locations.
+    ///
+    /// This works by calculating the mean Cartesian coordinates, and converting them
+    /// back to spherical coordinates. This is approximate, because the semi-minor (polar)
+    /// axis is assumed to equal the semi-major (equatorial) axis.
+    ///
+    /// https://en.wikipedia.org/wiki/Spherical_coordinate_system#Cartesian_coordinates
+    pub fn midpoint(locations: &[Location]) -> Self {
+        let mut x = 0f64;
+        let mut y = 0f64;
+        let mut z = 0f64;
+
+        for location in locations {
+            let cos_lat = location.latitude.to_radians().cos();
+            let sin_lat = location.latitude.to_radians().sin();
+            let cos_lon = location.longitude.to_radians().cos();
+            let sin_lon = location.longitude.to_radians().sin();
+            x += cos_lat * cos_lon;
+            y += cos_lat * sin_lon;
+            z += sin_lat;
+        }
+        let inv_total_weight = 1f64 / (locations.len() as f64);
+        x *= inv_total_weight;
+        y *= inv_total_weight;
+        z *= inv_total_weight;
+
+        let longitude = y.atan2(x);
+        let hypotenuse = (x * x + y * y).sqrt();
+        let latitude = z.atan2(hypotenuse);
+
+        Coordinates {
+            latitude: latitude.to_degrees(),
+            longitude: longitude.to_degrees(),
+        }
+    }
+}
+
 /// Takes input as latitude and longitude degrees.
 fn haversine_dist_deg(lat: f64, lon: f64, other_lat: f64, other_lon: f64) -> f64 {
     haversine_dist_rad(