Changed ip availability to an enum

Author: Pururun

mullvad-daemon/src/tunnel.rs

@@ -166,7 +166,7 @@ impl InnerParametersGenerator {
         let data = self.device().await?;
         let selected_relay = self
             .relay_selector
-            .get_relay(retry_attempt as usize, RuntimeParameters { ipv4, ipv6 })?;
+            .get_relay(retry_attempt as usize, RuntimeParameters::new(ipv4, ipv6))?;
 
         match selected_relay {
             #[cfg(not(target_os = "android"))]

mullvad-relay-selector/src/lib.rs

@@ -10,7 +10,7 @@ mod relay_selector;
 pub use error::Error;
 pub use relay_selector::{
     detailer, matcher, matcher::filter_matching_relay_list, query, relays::WireguardConfig,
-    AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay, RelaySelector,
-    RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig, OPENVPN_RETRY_ORDER,
-    WIREGUARD_RETRY_ORDER,
+    AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay, IpAvailability,
+    RelaySelector, RuntimeParameters, SelectedBridge, SelectedObfuscator, SelectorConfig,
+    OPENVPN_RETRY_ORDER, WIREGUARD_RETRY_ORDER,
 };

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

@@ -181,32 +181,64 @@ pub struct AdditionalWireguardConstraints {
 /// Values which affect the choice of relay but are only known at runtime.
 #[derive(Clone, Debug)]
 pub struct RuntimeParameters {
-    /// Whether IPv4 is available
-    pub ipv4: bool,
-    /// Whether IPv6 is available
-    pub ipv6: bool,
+    /// Whether IPv4, IPv6 or both is available
+    pub ip_availability: IpAvailability,
 }
 
 impl RuntimeParameters {
     /// Return whether a given [query][`RelayQuery`] is valid given the current runtime parameters
     pub fn compatible(&self, query: &RelayQuery) -> bool {
         match query.wireguard_constraints().ip_version {
-            Constraint::Any => self.ipv4 || self.ipv6,
-            Constraint::Only(talpid_types::net::IpVersion::V4) => self.ipv4,
-            Constraint::Only(talpid_types::net::IpVersion::V6) => self.ipv6,
+            Constraint::Any => true,
+            Constraint::Only(talpid_types::net::IpVersion::V4) => self.ip_availability.has_ipv4(),
+            Constraint::Only(talpid_types::net::IpVersion::V6) => self.ip_availability.has_ipv6(),
+        }
+    }
+
+    pub fn new(ipv4: bool, ipv6: bool) -> RuntimeParameters {
+        if ipv4 && ipv6 {
+            RuntimeParameters {
+                ip_availability: IpAvailability::All,
+            }
+        } else if !ipv6 {
+            RuntimeParameters {
+                ip_availability: IpAvailability::Ipv4,
+            }
+        } else if !ipv4 {
+            RuntimeParameters {
+                ip_availability: IpAvailability::Ipv6,
+            }
+        } else {
+            panic!("Device is offline!")
         }
     }
 }
 
 impl Default for RuntimeParameters {
     fn default() -> Self {
         RuntimeParameters {
-            ipv4: true,
-            ipv6: false,
+            ip_availability: IpAvailability::Ipv4,
         }
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum IpAvailability {
+    Ipv4,
+    Ipv6,
+    All,
+}
+
+impl IpAvailability {
+    fn has_ipv4(&self) -> bool {
+        self.clone() == IpAvailability::Ipv4 || self.clone() == IpAvailability::All
+    }
+
+    fn has_ipv6(&self) -> bool {
+        self.clone() == IpAvailability::Ipv6 || self.clone() == IpAvailability::All
+    }
+}
+
 /// This enum exists to separate the two types of [`SelectorConfig`] that exists.
 ///
 /// The first one is a "regular" config, where [`SelectorConfig::relay_settings`] is
@@ -1184,10 +1216,10 @@ impl RelaySelector {
 fn resolve_valid_ip_version(query: &RelayQuery, runtime_params: &RuntimeParameters) -> RelayQuery {
     let mut wireguard_constraints = query.wireguard_constraints().clone();
     if wireguard_constraints.ip_version.is_any() {
-        if runtime_params.ipv4 && !runtime_params.ipv6 {
+        if runtime_params.ip_availability == IpAvailability::Ipv4 {
             wireguard_constraints.ip_version = Constraint::Only(talpid_types::net::IpVersion::V4)
         }
-        if runtime_params.ipv6 && !runtime_params.ipv4 {
+        if runtime_params.ip_availability == IpAvailability::Ipv6 {
             wireguard_constraints.ip_version = Constraint::Only(talpid_types::net::IpVersion::V6)
         }
     }

mullvad-relay-selector/tests/relay_selector.rs

@@ -15,8 +15,8 @@ use talpid_types::net::{
 
 use mullvad_relay_selector::{
     query::{builder::RelayQueryBuilder, BridgeQuery, ObfuscationQuery, OpenVpnRelayQuery},
-    Error, GetRelay, RelaySelector, RuntimeParameters, SelectedObfuscator, SelectorConfig,
-    WireguardConfig, OPENVPN_RETRY_ORDER, WIREGUARD_RETRY_ORDER,
+    Error, GetRelay, IpAvailability, RelaySelector, RuntimeParameters, SelectedObfuscator,
+    SelectorConfig, WireguardConfig, OPENVPN_RETRY_ORDER, WIREGUARD_RETRY_ORDER,
 };
 use mullvad_types::{
     constraints::Constraint,
@@ -393,8 +393,7 @@ fn test_wireguard_retry_order() {
             .get_relay(
                 retry_attempt,
                 RuntimeParameters {
-                    ipv4: true,
-                    ipv6: true,
+                    ip_availability: IpAvailability::All,
                 },
             )
             .unwrap_or_else(|_| panic!("Retry attempt {retry_attempt} did not yield any relay"));
@@ -457,8 +456,7 @@ fn test_openvpn_retry_order() {
             .get_relay(
                 retry_attempt,
                 RuntimeParameters {
-                    ipv4: true,
-                    ipv6: true,
+                    ip_availability: IpAvailability::All,
                 },
             )
             .unwrap_or_else(|_| panic!("Retry attempt {retry_attempt} did not yield any relay"));
@@ -1645,8 +1643,7 @@ fn test_shadowsocks_runtime_ipv4_unavailable() {
     };
     let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
     let runtime_parameters = RuntimeParameters {
-        ipv4: false,
-        ipv6: true,
+        ip_availability: IpAvailability::Ipv6,
     };
     let user_result = relay_selector.get_relay(0, runtime_parameters).unwrap();
     assert!(
@@ -1676,8 +1673,7 @@ fn test_runtime_ipv4_unavailable() {
     };
     let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
     let runtime_parameters = RuntimeParameters {
-        ipv4: false,
-        ipv6: true,
+        ip_availability: IpAvailability::Ipv6,
     };
     let relay = relay_selector.get_relay(0, runtime_parameters).unwrap();
     match relay {