WIP Implement relay selection algorithm anew

Author: MarkusPettersson98

mullvad-management-interface/src/types/conversions/relay_constraints.rs

@@ -44,12 +44,9 @@ impl TryFrom<&proto::WireguardConstraints>
                 .entry_location
                 .clone()
                 .and_then(|loc| {
-                    Constraint::<mullvad_types::relay_constraints::LocationConstraint>::try_from(
-                        loc,
-                    )
-                    .ok()
+                    mullvad_types::relay_constraints::LocationConstraint::try_from(loc).ok()
                 })
-                .unwrap_or(Constraint::Any),
+                .into(),
             entry_providers: try_providers_constraint_from_proto(&constraints.entry_providers)?,
             entry_ownership: try_ownership_constraint_from_i32(constraints.entry_ownership)?,
         })
@@ -89,8 +86,10 @@ impl TryFrom<proto::RelaySettings> for mullvad_types::relay_constraints::RelaySe
             proto::relay_settings::Endpoint::Normal(settings) => {
                 let location = settings
                     .location
-                    .and_then(|loc| Constraint::<mullvad_types::relay_constraints::LocationConstraint>::try_from(loc).ok())
-                    .unwrap_or(Constraint::Any);
+                    .and_then(|loc| {
+                        mullvad_types::relay_constraints::LocationConstraint::try_from(loc).ok()
+                    })
+                    .into();
                 let providers = try_providers_constraint_from_proto(&settings.providers)?;
                 let ownership = try_ownership_constraint_from_i32(settings.ownership)?;
 
@@ -277,26 +276,28 @@ impl From<mullvad_types::relay_constraints::LocationConstraint> for proto::Locat
     }
 }
 
-impl TryFrom<proto::LocationConstraint>
-    for Constraint<mullvad_types::relay_constraints::LocationConstraint>
-{
+impl TryFrom<proto::LocationConstraint> for mullvad_types::relay_constraints::LocationConstraint {
     type Error = FromProtobufTypeError;
 
     fn try_from(location: proto::LocationConstraint) -> Result<Self, Self::Error> {
         use mullvad_types::relay_constraints::LocationConstraint;
-        match location.r#type {
-            Some(proto::location_constraint::Type::Location(location)) => Ok(Constraint::Only(
+        let Some(typ) = location.r#type else {
+            return Err(FromProtobufTypeError::InvalidArgument(
+                "Type of location constraint was not provided",
+            ));
+        };
+        match typ {
+            proto::location_constraint::Type::Location(location) => Ok(
                 LocationConstraint::Location(GeographicLocationConstraint::try_from(location)?),
-            )),
-            Some(proto::location_constraint::Type::CustomList(list_id)) => {
+            ),
+            proto::location_constraint::Type::CustomList(list_id) => {
                 let location = LocationConstraint::CustomList {
                     list_id: Id::from_str(&list_id).map_err(|_| {
                         FromProtobufTypeError::InvalidArgument("Id could not be parsed to a uuid")
                     })?,
                 };
-                Ok(Constraint::Only(location))
+                Ok(location)
             }
-            None => Ok(Constraint::Any),
         }
     }
 }
@@ -492,6 +493,17 @@ impl TryFrom<proto::RelayOverride> for mullvad_types::relay_constraints::RelayOv
     }
 }
 
+// TODO: Rename
+pub fn try_providers_from_proto(
+    providers: Vec<proto::relay_selector::Provider>,
+) -> Result<Constraint<mullvad_types::relay_constraints::Providers>, FromProtobufTypeError> {
+    let providers: Vec<_> = providers
+        .into_iter()
+        .map(|provider| provider.name)
+        .collect();
+    try_providers_constraint_from_proto(&providers)
+}
+
 pub fn try_providers_constraint_from_proto(
     providers: &[String],
 ) -> Result<Constraint<mullvad_types::relay_constraints::Providers>, FromProtobufTypeError> {

mullvad-management-interface/src/types/conversions/relay_selector.rs

@@ -1,14 +1,48 @@
-use crate::types::relay_selector::*;
+use crate::types::{
+    FromProtobufTypeError,
+    relay_constraints::{try_ownership_constraint_from_i32, try_providers_from_proto},
+    relay_selector::*,
+};
 
 impl TryFrom<Predicate> for mullvad_relay_selector::Predicate {
-    type Error = String;
+    type Error = FromProtobufTypeError;
 
     fn try_from(predicate: Predicate) -> Result<Self, Self::Error> {
         let Some(context) = predicate.context else {
             todo!("Return early");
         };
         match context {
-            predicate::Context::Singlehop(_constraints) => Ok(Self::Singlehop),
+            predicate::Context::Singlehop(constraints) => {
+                let EntryConstraints {
+                    general_constraints,
+                    obfuscation_settings: _, // TODO: Consider this parameter
+                    daita_settings: _,       // TODO: Consider this parameter
+                    ip_version: _,           // TODO: Consider this parameter
+                } = constraints;
+                // TODO: It might be beneficial to consolidate this whole conversion into a single
+                // type.
+                let (location, providers, ownership) = {
+                    match general_constraints {
+                        None => Default::default(),
+                        Some(constraints) => {
+                            let location = constraints
+                                .location
+                                .map(mullvad_types::relay_constraints::LocationConstraint::try_from)
+                                .transpose()?
+                                .into();
+                            let providers = try_providers_from_proto(constraints.providers)?;
+                            let ownership =
+                                try_ownership_constraint_from_i32(constraints.ownership)?;
+                            (location, providers, ownership)
+                        }
+                    }
+                };
+                Ok(Self::Singlehop {
+                    location,
+                    providers,
+                    ownership,
+                })
+            }
             predicate::Context::Autohop(_constraints) => Ok(Self::Autohop),
             predicate::Context::Entry(_constraints) => Ok(Self::Entry),
             predicate::Context::Exit(_constraints) => Ok(Self::Exit),

mullvad-relay-selector/src/relay_selector/mod.rs

@@ -16,7 +16,7 @@ use crate::{
     query::{ObfuscationQuery, RelayQuery, RelayQueryExt, WireguardRelayQuery},
 };
 
-use itertools::Itertools;
+use itertools::{Either, Itertools};
 pub use mullvad_types::relay_list::Relay;
 use mullvad_types::{
     CustomTunnelEndpoint, Intersection,
@@ -25,7 +25,8 @@ use mullvad_types::{
     endpoint::MullvadEndpoint,
     location::Coordinates,
     relay_constraints::{
-        ObfuscationSettings, RelayConstraints, RelaySettings, WireguardConstraints,
+        LocationConstraint, ObfuscationSettings, Ownership, Providers, RelayConstraints,
+        RelaySettings, WireguardConstraints,
     },
     relay_list::{Bridge, BridgeList, RelayList, WireguardRelay},
     settings::Settings,
@@ -813,26 +814,158 @@ impl RelaySelector {
         Ok((endpoint, bridge))
     }
 
-    // NEW relay selector API.
+    // == NEW relay selector API. ==
     // Starting afresh, but this should be used in existing functions.
 
     /// As oppossed to the prior [`Self::get_relay_by_query`], this function is stateless with
     /// regards to any particular config / settings.
     pub fn partition_relays(
-        &self, // TODO: If relay list is an in-parameter, we don't need this to be an associated
-        // function.
-        _predicate: Predicate,
-        // relays: &'a [Relay], // Implicit argument for now.
+        // TODO: If relay list is an in-parameter, we don't need this to be an associated function.
+        &self,
+        predicate: Predicate,
     ) -> RelayPartitions {
-        let partitions: RelayPartitions = RelayPartitions::default();
-        partitions
+        // Implicit argument for now. Might as well be an explicit in-parameter.
+        // let relays = self.get_relays();
+        let relays: Vec<Relay> = vec![];
+        // The relay selection algorithm is embarrassingly parallel: https://en.wikipedia.org/wiki/Embarrassingly_parallel.
+        // We may explore the entire search space (`relays` x `criteria`) without any synchronisation
+        // between different branches.
+        let verdicts: Vec<(Relay, Verdict)> = match predicate {
+            Predicate::Singlehop {
+                location,
+                providers,
+                ownership,
+            } => {
+                /* pseudo-code (implemented mostly in `Criteria`).
+                 *
+                 * let relay := relays.pop()
+                 * let mut rejections := []
+                 *
+                 * if let Some(reason) critera(relay) {
+                 *   rejections.push(reason)
+                 * }
+                 * ..
+                 * if rejections.empty() {
+                 *   Accept(relay)
+                 * } else {
+                 *   Fail(relay, rejections)
+                 * }
+                 * */
+                let criteria = [
+                    Criteria::otherwise(|relay| relay.active, |_| Reject::Inactive),
+                    Criteria::otherwise(|relay| relay.hostname == "se", |_| Reject::Inactive),
+                ];
+                relays
+                    .into_iter()
+                    // This part of the algorithm maps each relay to a verdict: Either Accept or
+                    // Reject(Reason).
+                    .map(|relay| {
+                        let verdict = Criteria::eval(criteria.iter(), &relay);
+                        (relay, verdict)
+                    })
+                    .collect()
+            }
+            Predicate::Autohop => todo!("Implement partition_relays(Autohop)"),
+            Predicate::Entry => todo!("Implement partition_relays(Entry)"),
+            Predicate::Exit => todo!("Implement partition_relays(Exit)"),
+        };
+        // After this mapping, a single reduce is performed to partition the relays based on
+        // their assigned verdict.
+        verdicts
+            .into_iter()
+            .partition_map(|(relay, verdict)| match verdict {
+                Verdict::Accept => Either::Left(relay),
+                Verdict::Fail(rejected) => Either::Right((relay, rejected)),
+            })
+            .into()
+    }
+}
+
+/// A criteria is a function from a _single_ constraint and a relay to a [`Verdict`].
+///
+/// Multiple [`Criteria`] can be evaluated against a single relay at once by [`Criteria::eval`]. A
+/// final verdict is then compiled. If applicable, all reject reasons are accumulated and presented
+/// as a single [`Verdict::Fail`].
+struct Criteria<'a> {
+    f: Box<dyn Fn(&Relay) -> Verdict + 'a>,
+}
+
+impl<'a> Criteria<'a> {
+    /// Create a new [`Criteria`].
+    fn new(f: impl Fn(&Relay) -> Verdict + 'a) -> Self {
+        Criteria { f: Box::new(f) }
+    }
+
+    /// If the given criteria [`f`] evaulates to `false`, the second provided function `reason` is
+    /// run to provide a single [`Reject`] reason. `reason` gets access to the failing relay, which
+    /// means that `reason` may derivce additional information for why this particular relay was
+    /// rejected.
+    ///
+    /// This is a short-hand for how most common [`Criteria`]s will be formulated, and it allows the
+    /// caller to nicely separate the scrutinizing rejection logic from the logic extracting data to
+    /// provide together with the final rejection. In the happy case this carries minimal additional
+    /// runtime overhead compared to [`Criteria::new`], but upon a rejection two functions will run
+    /// instead of one. For more fine-grained control over this behavior, prefer [`Criteria::new`].
+    fn otherwise(f: impl Fn(&Relay) -> bool + 'a, reason: impl Fn(&Relay) -> Reject + 'a) -> Self {
+        Criteria::new(move |relay| {
+            if f(relay) {
+                Verdict::Accept
+            } else {
+                Verdict::Fail(vec![reason(relay)])
+            }
+        })
+    }
+
+    /// Evaluate a single [`Criteria`] for a single [`Relay`].
+    fn run(&self, relay: &Relay) -> Verdict {
+        (self.f)(relay)
+    }
+
+    /// Evaluate all criterias for a given relay, resulting in a single final verdict.
+    fn eval(criterias: impl Iterator<Item = &'a Criteria<'a>>, relay: &Relay) -> Verdict {
+        let mut rejections = vec![];
+        for criteria in criterias.into_iter() {
+            if let Verdict::Fail(reasons) = criteria.run(relay) {
+                for reason in reasons {
+                    rejections.push(reason);
+                }
+            };
+        }
+        match rejections.is_empty() {
+            true => Verdict::Accept,
+            false => Verdict::Fail(rejections),
+        }
+    }
+}
+
+/// If a relay is accepted or rejected .
+///
+/// # Note
+/// The associated relay is implied from the environment.
+#[derive(Debug)]
+enum Verdict {
+    Accept,
+    Fail(Vec<Reject>),
+}
+
+impl From<(Vec<Relay>, Vec<(Relay, Vec<Reject>)>)> for RelayPartitions {
+    /// Map the result of [`Itertools::partition_map`] to [`RelayPartitions`].
+    fn from(partitions: (Vec<Relay>, Vec<(Relay, Vec<Reject>)>)) -> Self {
+        Self {
+            matches: partitions.0,
+            discards: partitions.1,
+        }
     }
 }
 
 /// Specify the constraints that should be applied when selecting relays,
 /// along with a context that may affect the selection behavior.
 pub enum Predicate {
-    Singlehop,
+    Singlehop {
+        location: Constraint<LocationConstraint>,
+        providers: Constraint<Providers>,
+        ownership: Constraint<Ownership>,
+    },
     Autohop,
     // Multihop-only
     Entry,
@@ -853,6 +986,8 @@ pub enum Reject {
     // TODO: Add more reasons - at least all in `relay_selector.proto`.
 }
 
+// == End of new relay selector API. ==
+
 fn apply_ip_availability(
     runtime_ip_availability: IpAvailability,
     user_query: &mut RelayQuery,