Refactor to avoid representing offline state

Serock3 · Serock3 · commit b1bd8bed41be · 2025-03-21T11:31:56.000+01:00
diff --git a/mullvad-daemon/src/tunnel.rs b/mullvad-daemon/src/tunnel.rs
@@ -8,7 +8,7 @@ use std::{
 
 use tokio::sync::Mutex;
 
-use mullvad_relay_selector::{GetRelay, RelaySelector, RuntimeIpAvailability, WireguardConfig};
+use mullvad_relay_selector::{GetRelay, RelaySelector, WireguardConfig};
 use mullvad_types::{
     endpoint::MullvadWireguardEndpoint, location::GeoIpLocation, relay_list::Relay,
     settings::TunnelOptions,
@@ -23,7 +23,7 @@ use talpid_types::net::{
 #[cfg(target_os = "android")]
 use talpid_types::net::{obfuscation::ObfuscatorConfig, wireguard, TunnelParameters};
 
-use talpid_types::{tunnel::ParameterGenerationError, ErrorExt};
+use talpid_types::{net::IpAvailbility, tunnel::ParameterGenerationError, ErrorExt};
 
 use crate::device::{AccountManagerHandle, Error as DeviceError, PrivateAccountAndDevice};
 
@@ -160,14 +160,12 @@ impl InnerParametersGenerator {
     async fn generate(
         &mut self,
         retry_attempt: u32,
-        ipv4: bool,
-        ipv6: bool,
+        ip_availability: IpAvailbility,
     ) -> Result<TunnelParameters, Error> {
         let data = self.device().await?;
-        let selected_relay = self.relay_selector.get_relay(
-            retry_attempt as usize,
-            RuntimeIpAvailability::new(ipv4, ipv6),
-        )?;
+        let selected_relay = self
+            .relay_selector
+            .get_relay(retry_attempt as usize, ip_availability)?;
 
         match selected_relay {
             #[cfg(not(target_os = "android"))]
@@ -301,14 +299,13 @@ impl TunnelParametersGenerator for ParametersGenerator {
     fn generate(
         &mut self,
         retry_attempt: u32,
-        ipv4: bool,
-        ipv6: bool,
+        ip_availbility: IpAvailbility,
     ) -> Pin<Box<dyn Future<Output = Result<TunnelParameters, ParameterGenerationError>>>> {
         let generator = self.0.clone();
         Box::pin(async move {
             let mut inner = generator.lock().await;
             inner
-                .generate(retry_attempt, ipv4, ipv6)
+                .generate(retry_attempt, ip_availbility)
                 .await
                 .inspect_err(|error| {
                     log::error!(
diff --git a/mullvad-relay-selector/src/lib.rs b/mullvad-relay-selector/src/lib.rs
@@ -11,6 +11,5 @@ pub use error::Error;
 pub use relay_selector::{
     detailer, matcher, matcher::filter_matching_relay_list, query, relays::WireguardConfig,
     AdditionalRelayConstraints, AdditionalWireguardConstraints, GetRelay, RelaySelector,
-    RuntimeIpAvailability, SelectedBridge, SelectedObfuscator, SelectorConfig, OPENVPN_RETRY_ORDER,
-    WIREGUARD_RETRY_ORDER,
+    SelectedBridge, SelectedObfuscator, SelectorConfig, OPENVPN_RETRY_ORDER, WIREGUARD_RETRY_ORDER,
 };
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -47,7 +47,7 @@ use talpid_types::{
     net::{
         obfuscation::ObfuscatorConfig,
         proxy::{CustomProxy, Shadowsocks},
-        Endpoint, IpVersion, TransportProtocol, TunnelType,
+        Endpoint, IpAvailbility, IpVersion, TransportProtocol, TunnelType,
     },
     ErrorExt,
 };
@@ -178,29 +178,6 @@ pub struct AdditionalWireguardConstraints {
     pub quantum_resistant: QuantumResistantState,
 }
 
-/// Whether IPv4 and IPv6 is available at runtime.
-///
-/// `None` means that no IP version is available, `Some(Any)` means that both are available,
-#[derive(Clone, Debug)]
-pub struct RuntimeIpAvailability(Option<Constraint<IpVersion>>);
-
-impl RuntimeIpAvailability {
-    pub fn new(ipv4: bool, ipv6: bool) -> RuntimeIpAvailability {
-        RuntimeIpAvailability(match (ipv4, ipv6) {
-            (true, true) => Some(Constraint::Any),
-            (false, true) => Some(Constraint::Only(IpVersion::V6)),
-            (true, false) => Some(Constraint::Only(IpVersion::V4)),
-            (false, false) => None,
-        })
-    }
-}
-
-impl Default for RuntimeIpAvailability {
-    fn default() -> Self {
-        RuntimeIpAvailability(Some(Constraint::Only(IpVersion::V4)))
-    }
-}
-
 /// This enum exists to separate the two types of [`SelectorConfig`] that exists.
 ///
 /// The first one is a "regular" config, where [`SelectorConfig::relay_settings`] is
@@ -550,7 +527,7 @@ impl RelaySelector {
     pub fn get_relay(
         &self,
         retry_attempt: usize,
-        runtime_ip_availability: RuntimeIpAvailability,
+        runtime_ip_availability: IpAvailbility,
     ) -> Result<GetRelay, Error> {
         let config_guard = self.config.lock().unwrap();
         let config = SpecializedSelectorConfig::from(&*config_guard);
@@ -584,7 +561,7 @@ impl RelaySelector {
         &self,
         retry_attempt: usize,
         retry_order: &[RelayQuery],
-        runtime_ip_availability: RuntimeIpAvailability,
+        runtime_ip_availability: IpAvailbility,
     ) -> Result<GetRelay, Error> {
         let config_guard = self.config.lock().unwrap();
         let config = SpecializedSelectorConfig::from(&*config_guard);
@@ -627,7 +604,7 @@ impl RelaySelector {
     fn pick_and_merge_query(
         retry_attempt: usize,
         retry_order: &[RelayQuery],
-        runtime_ip_availability: RuntimeIpAvailability,
+        runtime_ip_availability: IpAvailbility,
         user_config: &NormalSelectorConfig<'_>,
         parsed_relays: &RelayList,
     ) -> Result<RelayQuery, Error> {
@@ -1169,16 +1146,19 @@ impl RelaySelector {
 }
 
 fn apply_ip_availability(
-    runtime_ip_availability: RuntimeIpAvailability,
+    runtime_ip_availability: IpAvailbility,
     user_query: &mut RelayQuery,
 ) -> Result<(), Error> {
+    let ip_version = match runtime_ip_availability {
+        IpAvailbility::IpV4 => Constraint::Only(IpVersion::V4),
+        IpAvailbility::IpV6 => Constraint::Only(IpVersion::V6),
+        IpAvailbility::IpV4AndIpV6 => Constraint::Any,
+    };
     let wireguard_constraints = user_query
         .wireguard_constraints()
         .to_owned()
         .intersection(WireguardRelayQuery {
-            ip_version: runtime_ip_availability
-                .0
-                .ok_or(Error::IpVersionUnavailable)?,
+            ip_version,
             ..Default::default()
         })
         .ok_or(Error::IpVersionUnavailable)?;
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/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, RuntimeIpAvailability, SelectedObfuscator, SelectorConfig,
-    WireguardConfig, OPENVPN_RETRY_ORDER, WIREGUARD_RETRY_ORDER,
+    Error, GetRelay, RelaySelector, SelectedObfuscator, SelectorConfig, WireguardConfig,
+    OPENVPN_RETRY_ORDER, WIREGUARD_RETRY_ORDER,
 };
 use mullvad_types::{
     constraints::Constraint,
@@ -391,7 +391,7 @@ fn test_wireguard_retry_order() {
     let relay_selector = default_relay_selector();
     for (retry_attempt, query) in WIREGUARD_RETRY_ORDER.iter().enumerate() {
         let relay = relay_selector
-            .get_relay(retry_attempt, RuntimeIpAvailability::new(false, true))
+            .get_relay(retry_attempt, talpid_types::net::IpAvailbility::IpV6)
             .unwrap_or_else(|_| panic!("Retry attempt {retry_attempt} did not yield any relay"));
         // For each relay, cross-check that the it has the expected tunnel protocol
         let tunnel_type = tunnel_type(&unwrap_relay(relay.clone()));
@@ -450,7 +450,7 @@ fn test_openvpn_retry_order() {
 
     for (retry_attempt, query) in OPENVPN_RETRY_ORDER.iter().enumerate() {
         let relay = relay_selector
-            .get_relay(retry_attempt, RuntimeIpAvailability::new(false, true))
+            .get_relay(retry_attempt, talpid_types::net::IpAvailbility::IpV6)
             .unwrap_or_else(|_| panic!("Retry attempt {retry_attempt} did not yield any relay"));
         // For each relay, cross-check that the it has the expected tunnel protocol
         let tunnel_type = tunnel_type(&unwrap_relay(relay.clone()));
@@ -1159,7 +1159,7 @@ fn test_openvpn_auto_bridge() {
             .get_relay_with_custom_params(
                 retry_attempt,
                 &retry_order,
-                RuntimeIpAvailability::default(),
+                talpid_types::net::IpAvailbility::IpV4,
             )
             .unwrap();
         match relay {
@@ -1269,7 +1269,7 @@ fn test_include_in_country() {
     // If include_in_country is false for all relays, a relay must be selected anyway.
     let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relay_list.clone());
     assert!(relay_selector
-        .get_relay(0, RuntimeIpAvailability::default())
+        .get_relay(0, talpid_types::net::IpAvailbility::IpV4)
         .is_ok());
 
     // If include_in_country is true for some relay, it must always be selected.
@@ -1278,7 +1278,7 @@ fn test_include_in_country() {
     let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relay_list);
     let relay = unwrap_relay(
         relay_selector
-            .get_relay(0, RuntimeIpAvailability::default())
+            .get_relay(0, talpid_types::net::IpAvailbility::IpV4)
             .expect("expected match"),
     );
 
@@ -1615,7 +1615,7 @@ fn valid_user_setting_should_yield_relay() {
     let user_result = relay_selector.get_relay_by_query(user_query.clone());
     for retry_attempt in 0..WIREGUARD_RETRY_ORDER.len() {
         let post_unification_result =
-            relay_selector.get_relay(retry_attempt, RuntimeIpAvailability::default());
+            relay_selector.get_relay(retry_attempt, talpid_types::net::IpAvailbility::IpV4);
         if user_result.is_ok() {
             assert!(post_unification_result.is_ok(), "Expected Post-unification query to be valid because original query {:#?} yielded a connection configuration", user_query)
         }
@@ -1638,7 +1638,7 @@ fn test_shadowsocks_runtime_ipv4_unavailable() {
         ..SelectorConfig::default()
     };
     let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
-    let runtime_parameters = RuntimeIpAvailability::new(false, true);
+    let runtime_parameters = talpid_types::net::IpAvailbility::IpV6;
     let user_result = relay_selector.get_relay(0, runtime_parameters).unwrap();
     assert!(
         matches!(user_result, GetRelay::Wireguard {
@@ -1666,7 +1666,7 @@ fn test_runtime_ipv4_unavailable() {
         ..SelectorConfig::default()
     };
     let relay_selector = RelaySelector::from_list(config, RELAYS.clone());
-    let runtime_parameters = RuntimeIpAvailability::new(false, true);
+    let runtime_parameters = talpid_types::net::IpAvailbility::IpV6;
     let relay = relay_selector.get_relay(0, runtime_parameters).unwrap();
     match relay {
         GetRelay::Wireguard { endpoint, .. } => {
diff --git a/talpid-core/src/offline/linux.rs b/talpid-core/src/offline/linux.rs
@@ -82,7 +82,20 @@ async fn check_connectivity(handle: &RouteManagerHandle, fwmark: Option<u32>) ->
         route_exists(PUBLIC_INTERNET_ADDRESS_V4).await,
         route_exists(PUBLIC_INTERNET_ADDRESS_V6).await,
     ) {
-        (Ok(ipv4), Ok(ipv6)) => Connectivity::Status { ipv4, ipv6 },
+        (Ok(true), Ok(true)) => Connectivity::Online(talpid_types::net::IpAvailbility::IpV4AndIpV6),
+        (Ok(false), Ok(true)) => Connectivity::Online(talpid_types::net::IpAvailbility::IpV6),
+        (Ok(true), no_ipv6) => {
+            // Errors for IPv6 likely mean it's disabled, so assume it's unavailable
+            if let Err(err) = no_ipv6 {
+                log::trace!(
+                    "{}",
+                    err.display_chain_with_msg(
+                        "Failed to infer offline state for IPv6. Assuming it's unavailable"
+                    )
+                )
+            }
+            Connectivity::Online(talpid_types::net::IpAvailbility::IpV4)
+        }
         // If we fail to retrieve the IPv4 route, always assume we're connected
         (Err(err), _) => {
             log::error!(
@@ -91,15 +104,6 @@ async fn check_connectivity(handle: &RouteManagerHandle, fwmark: Option<u32>) ->
             );
             Connectivity::PresumeOnline
         }
-        // Errors for IPv6 likely mean it's disabled, so assume it's unavailable
-        (Ok(ipv4), Err(err)) => {
-            log::trace!(
-                "{}",
-                err.display_chain_with_msg(
-                    "Failed to infer offline state for IPv6. Assuming it's unavailable"
-                )
-            );
-            Connectivity::Status { ipv4, ipv6: false }
-        }
+        _ => Connectivity::Offline,
     }
 }
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -10,7 +10,7 @@ use talpid_routing::RouteManagerHandle;
 use talpid_tunnel::tun_provider::TunProvider;
 use talpid_tunnel::{EventHook, TunnelArgs, TunnelEvent, TunnelMetadata};
 use talpid_types::net::{
-    AllowedClients, AllowedEndpoint, AllowedTunnelTraffic, Connectivity, TunnelParameters,
+    AllowedClients, AllowedEndpoint, AllowedTunnelTraffic, IpAvailbility, TunnelParameters,
 };
 use talpid_types::tunnel::{ErrorStateCause, FirewallPolicyError};
 use talpid_types::ErrorExt;
@@ -75,22 +75,25 @@ impl ConnectingState {
                 });
         }
 
-        if shared_values.connectivity == Connectivity::Offline {
-            // FIXME: Temporary: Nudge route manager to update the default interface
-            #[cfg(target_os = "macos")]
-            {
-                log::debug!("Poking route manager to update default routes");
-                let _ = shared_values.route_manager.refresh_routes();
+        let ip_availability = match shared_values.connectivity {
+            talpid_types::net::Connectivity::Offline => {
+                // FIXME: Temporary: Nudge route manager to update the default interface
+                #[cfg(target_os = "macos")]
+                {
+                    log::debug!("Poking route manager to update default routes");
+                    let _ = shared_values.route_manager.refresh_routes();
+                }
+                return ErrorState::enter(shared_values, ErrorStateCause::IsOffline);
             }
-            return ErrorState::enter(shared_values, ErrorStateCause::IsOffline);
-        }
-        match shared_values
-            .runtime
-            .block_on(shared_values.tunnel_parameters_generator.generate(
-                retry_attempt,
-                shared_values.connectivity.has_ipv4(),
-                shared_values.connectivity.has_ipv6(),
-            )) {
+            talpid_types::net::Connectivity::PresumeOnline => IpAvailbility::IpV4,
+            talpid_types::net::Connectivity::Online(ip_availbility) => ip_availbility,
+        };
+
+        match shared_values.runtime.block_on(
+            shared_values
+                .tunnel_parameters_generator
+                .generate(retry_attempt, ip_availability),
+        ) {
             Err(err) => {
                 ErrorState::enter(shared_values, ErrorStateCause::TunnelParameterError(err))
             }
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -46,7 +46,7 @@ use std::{
 #[cfg(target_os = "android")]
 use talpid_types::{android::AndroidContext, ErrorExt};
 use talpid_types::{
-    net::{AllowedEndpoint, Connectivity, TunnelParameters},
+    net::{AllowedEndpoint, Connectivity, IpAvailbility, TunnelParameters},
     tunnel::{ErrorStateCause, ParameterGenerationError, TunnelStateTransition},
 };
 
@@ -452,8 +452,7 @@ pub trait TunnelParametersGenerator: Send + 'static {
     fn generate(
         &mut self,
         retry_attempt: u32,
-        ipv4: bool,
-        ipv6: bool,
+        ip_availbility: IpAvailbility,
     ) -> Pin<Box<dyn Future<Output = Result<TunnelParameters, ParameterGenerationError>>>>;
 }
 
diff --git a/talpid-types/src/lib.rs b/talpid-types/src/lib.rs
@@ -1,7 +1,6 @@
 #[cfg(target_os = "android")]
 pub mod android;
 pub mod net;
-pub use net::IpAvailability;
 pub mod tunnel;
 
 #[cfg(target_os = "linux")]
diff --git a/talpid-types/src/net/mod.rs b/talpid-types/src/net/mod.rs
