Do not panic when reading public key of relay

mullvad-daemon/src/tunnel.rs

@@ -44,11 +44,8 @@ pub enum Error {
     #[error("Not logged in on a valid device")]
     NoAuthDetails,
 
-    #[error("No relay available")]
-    NoRelayAvailable,
-
-    #[error("No bridge available")]
-    NoBridgeAvailable,
+    #[error("Failed to select a matching relay")]
+    SelectRelay(#[from] mullvad_relay_selector::Error),
 
     #[error("Failed to resolve hostname for custom relay")]
     ResolveCustomHostname,
@@ -148,11 +145,7 @@ impl InnerParametersGenerator {
         let data = self.device().await?;
         let selected_relay = self
             .relay_selector
-            .get_relay(retry_attempt as usize, RuntimeParameters { ipv6 })
-            .map_err(|err| match err {
-                mullvad_relay_selector::Error::NoBridge => Error::NoBridgeAvailable,
-                _ => Error::NoRelayAvailable,
-            })?;
+            .get_relay(retry_attempt as usize, RuntimeParameters { ipv6 })?;
 
         match selected_relay {
             #[cfg(not(target_os = "android"))]
@@ -287,7 +280,9 @@ impl TunnelParametersGenerator for ParametersGenerator {
                 .generate(retry_attempt, ipv6)
                 .await
                 .map_err(|error| match error {
-                    Error::NoBridgeAvailable => ParameterGenerationError::NoMatchingBridgeRelay,
+                    Error::SelectRelay(mullvad_relay_selector::Error::NoBridge) => {
+                        ParameterGenerationError::NoMatchingBridgeRelay
+                    }
                     Error::ResolveCustomHostname => {
                         ParameterGenerationError::CustomTunnelHostResultionError
                     }

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

@@ -14,10 +14,14 @@ use mullvad_types::{
     constraints::Constraint,
     endpoint::MullvadWireguardEndpoint,
     relay_constraints::TransportPort,
-    relay_list::{OpenVpnEndpoint, OpenVpnEndpointData, Relay, WireguardEndpointData},
+    relay_list::{
+        OpenVpnEndpoint, OpenVpnEndpointData, Relay, RelayEndpointData, WireguardEndpointData,
+    },
 };
 use talpid_types::net::{
-    all_of_the_internet, wireguard::PeerConfig, Endpoint, IpVersion, TransportProtocol,
+    all_of_the_internet,
+    wireguard::{PeerConfig, PublicKey},
+    Endpoint, IpVersion, TransportProtocol,
 };
 
 use super::{
@@ -31,6 +35,8 @@ pub enum Error {
     NoOpenVpnEndpoint,
     #[error("No bridge endpoint could be derived")]
     NoBridgeEndpoint,
+    #[error("OpenVPN relays and bridges does not have a public key. Expected a Wireguard relay")]
+    MissingPublicKey,
     #[error("The selected relay does not support IPv6")]
     NoIPv6(Box<Relay>),
     #[error("Invalid port argument: port {0} is not in any valid Wireguard port range")]
@@ -70,7 +76,7 @@ fn wireguard_singlehop_endpoint(
         SocketAddr::new(host, port)
     };
     let peer_config = PeerConfig {
-        public_key: exit.endpoint_data.unwrap_wireguard_ref().public_key.clone(),
+        public_key: get_public_key(exit)?.clone(),
         endpoint,
         allowed_ips: all_of_the_internet(),
         // This will be filled in later, not the relay selector's problem
@@ -106,7 +112,7 @@ fn wireguard_multihop_endpoint(
         SocketAddr::from((ip, port))
     };
     let exit = PeerConfig {
-        public_key: exit.endpoint_data.unwrap_wireguard_ref().public_key.clone(),
+        public_key: get_public_key(exit)?.clone(),
         endpoint: exit_endpoint,
         // The exit peer should be able to route incoming VPN traffic to the rest of
         // the internet.
@@ -121,11 +127,7 @@ fn wireguard_multihop_endpoint(
         SocketAddr::from((host, port))
     };
     let entry = PeerConfig {
-        public_key: entry
-            .endpoint_data
-            .unwrap_wireguard_ref()
-            .public_key
-            .clone(),
+        public_key: get_public_key(entry)?.clone(),
         endpoint: entry_endpoint,
         // The entry peer should only be able to route incoming VPN traffic to the
         // exit peer.
@@ -177,6 +179,15 @@ fn get_port_for_wireguard_relay(
     }
 }
 
+/// Read the [`PublicKey`] of a relay. This will only succeed if [relay][`Relay`] is a
+/// [Wireguard][`RelayEndpointData::Wireguard`] relay.
+const fn get_public_key(relay: &Relay) -> Result<&PublicKey, Error> {
+    match &relay.endpoint_data {
+        RelayEndpointData::Wireguard(endpoint) => Ok(&endpoint.public_key),
+        RelayEndpointData::Openvpn | RelayEndpointData::Bridge => Err(Error::MissingPublicKey),
+    }
+}
+
 /// Selects a random port number from a list of provided port ranges.
 ///
 /// This function iterates over a list of port ranges, each represented as a tuple (u16, u16)

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

@@ -114,18 +114,24 @@ pub const fn filter_openvpn(relay: &Relay) -> bool {
 }
 
 /// Returns whether the relay matches the tunnel constraint `filter`
+#[cfg(not(target_os = "android"))]
 pub const fn filter_tunnel_type(filter: &Constraint<TunnelType>, relay: &Relay) -> bool {
     match filter {
         Constraint::Any => true,
         Constraint::Only(typ) => match typ {
-            // Do not keep OpenVPN relays on Android
-            TunnelType::OpenVpn if cfg!(target_os = "android") => false,
             TunnelType::OpenVpn => filter_openvpn(relay),
             TunnelType::Wireguard => filter_wireguard(relay),
         },
     }
 }
 
+/// Returns whether the relay matches the tunnel constraint `filter`
+#[cfg(target_os = "android")]
+pub const fn filter_tunnel_type(_: &Constraint<TunnelType>, relay: &Relay) -> bool {
+    // Only keep Wireguard relays on Android (i.e. filter out OpenVPN relays)
+    filter_wireguard(relay)
+}
+
 /// Returns whether the relay is a Wireguard relay.
 pub const fn filter_wireguard(relay: &Relay) -> bool {
     matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))

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

@@ -25,7 +25,7 @@ use mullvad_types::{
         OpenVpnConstraints, RelayConstraints, RelayOverride, RelaySettings, ResolvedBridgeSettings,
         SelectedObfuscation, WireguardConstraints,
     },
-    relay_list::{Relay, RelayList},
+    relay_list::{Relay, RelayEndpointData, RelayList},
     settings::Settings,
     CustomTunnelEndpoint,
 };
@@ -201,10 +201,39 @@ pub enum GetRelay {
 /// will actually be routed to the broader internet.
 #[derive(Clone, Debug)]
 pub enum WireguardConfig {
+    /// Strongly prefer to instantiate this variant using [`WireguardConfig::singlehop`] as that
+    /// will assert that the relay is of the expected type.
     Singlehop { exit: Relay },
+    /// Strongly prefer to instantiate this variant using [`WireguardConfig::multihop`] as that
+    /// will assert that the entry & exit relays are of the expected type.
     Multihop { exit: Relay, entry: Relay },
 }
 
+impl WireguardConfig {
+    const fn singlehop(exit: Relay) -> Self {
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            exit.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        Self::Singlehop { exit }
+    }
+
+    const fn multihop(exit: Relay, entry: Relay) -> Self {
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            exit.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(
+            entry.endpoint_data,
+            RelayEndpointData::Wireguard(_)
+        ));
+        Self::Multihop { exit, entry }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum SelectedBridge {
     Normal { settings: CustomProxy, relay: Relay },
@@ -454,7 +483,7 @@ impl RelaySelector {
             }
             SpecializedSelectorConfig::Normal(pure_config) => {
                 let parsed_relays = &self.parsed_relays.lock().unwrap();
-                Self::get_relay_inner(&query, parsed_relays, &pure_config)
+                Self::get_relay_inner(query, parsed_relays, &pure_config)
             }
         }
     }
@@ -522,7 +551,7 @@ impl RelaySelector {
                     runtime_params,
                     user_preferences,
                 );
-                Self::get_relay_inner(&query, parsed_relays, &normal_config)
+                Self::get_relay_inner(query, parsed_relays, &normal_config)
             }
         }
     }
@@ -568,7 +597,7 @@ impl RelaySelector {
     /// * An `Err` if no suitable bridge is found
     #[cfg(not(target_os = "android"))]
     fn get_relay_inner(
-        query: &RelayQuery,
+        query: RelayQuery,
         parsed_relays: &ParsedRelays,
         config: &NormalSelectorConfig<'_>,
     ) -> Result<GetRelay, Error> {
@@ -585,7 +614,7 @@ impl RelaySelector {
                     let mut new_query = query.clone();
                     new_query.tunnel_protocol = Constraint::Only(tunnel_type);
                     // If a suitable relay is found, short-circuit and return it
-                    if let Ok(relay) = Self::get_relay_inner(&new_query, parsed_relays, config) {
+                    if let Ok(relay) = Self::get_relay_inner(new_query, parsed_relays, config) {
                         return Ok(relay);
                     }
                 }
@@ -596,15 +625,24 @@ impl RelaySelector {
 
     #[cfg(target_os = "android")]
     fn get_relay_inner(
-        query: &RelayQuery,
+        mut query: RelayQuery,
         parsed_relays: &ParsedRelays,
         config: &NormalSelectorConfig<'_>,
     ) -> Result<GetRelay, Error> {
+        // FIXME: A bit of defensive programming - calling `get_wiregurad_relay` with a query that doesn't
+        // specify Wireguard as the desired tunnel type is not valid and will lead to unwanted
+        // behavior. This should be seen as a workaround, and it would be nicer to lift this
+        // invariant to be checked by the type system instead.
+        query.tunnel_protocol = Constraint::Only(TunnelType::Wireguard);
         Self::get_wireguard_relay(query, config, parsed_relays)
     }
 
     /// Derive a valid Wireguard relay configuration from `query`.
     ///
+    /// # Note
+    /// This function should *only* be called with a Wireguard query.
+    /// It will panic if the tunnel type is not specified as Wireguard.
+    ///
     /// # Returns
     /// * An `Err` if no exit relay can be chosen
     /// * An `Err` if no entry relay can be chosen (if multihop is enabled on `query`)
@@ -613,18 +651,22 @@ impl RelaySelector {
     ///
     /// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
     fn get_wireguard_relay(
-        query: &RelayQuery,
+        query: RelayQuery,
         config: &NormalSelectorConfig<'_>,
         parsed_relays: &ParsedRelays,
     ) -> Result<GetRelay, Error> {
+        assert_eq!(
+            query.tunnel_protocol,
+            Constraint::Only(TunnelType::Wireguard)
+        );
         let inner = if !query.wireguard_constraints.multihop() {
-            Self::get_wireguard_singlehop_config(query, config, parsed_relays)?
+            Self::get_wireguard_singlehop_config(&query, config, parsed_relays)?
         } else {
-            Self::get_wireguard_multihop_config(query, config, parsed_relays)?
+            Self::get_wireguard_multihop_config(&query, config, parsed_relays)?
         };
-        let endpoint = Self::get_wireguard_endpoint(query, parsed_relays, &inner)?;
+        let endpoint = Self::get_wireguard_endpoint(&query, parsed_relays, &inner)?;
         let obfuscator =
-            Self::get_wireguard_obfuscator(query, inner.clone(), &endpoint, parsed_relays)?;
+            Self::get_wireguard_obfuscator(&query, inner.clone(), &endpoint, parsed_relays)?;
 
         Ok(GetRelay::Wireguard {
             endpoint,
@@ -646,7 +688,8 @@ impl RelaySelector {
         let candidates =
             filter_matching_relay_list(query, parsed_relays.relays(), config.custom_lists);
         helpers::pick_random_relay(&candidates)
-            .map(|exit| WireguardConfig::Singlehop { exit: exit.clone() })
+            .cloned()
+            .map(WireguardConfig::singlehop)
             .ok_or(Error::NoRelay)
     }
 
@@ -700,10 +743,7 @@ impl RelaySelector {
         }
         .ok_or(Error::NoRelay)?;
 
-        Ok(WireguardConfig::Multihop {
-            exit: exit.clone(),
-            entry: entry.clone(),
-        })
+        Ok(WireguardConfig::multihop(exit.clone(), entry.clone()))
     }
 
     /// Constructs a [`MullvadEndpoint`] with details for how to connect to `relay`.
@@ -754,6 +794,10 @@ impl RelaySelector {
 
     /// Derive a valid OpenVPN relay configuration from `query`.
     ///
+    /// # Note
+    /// This function should *only* be called with an OpenVPN query.
+    /// It will panic if the tunnel type is not specified as OpenVPN.
+    ///
     /// # Returns
     /// * An `Err` if no exit relay can be chosen
     /// * An `Err` if no entry bridge can be chosen (if bridge mode is enabled on `query`)
@@ -763,16 +807,19 @@ impl RelaySelector {
     /// [`MullvadEndpoint`]: mullvad_types::endpoint::MullvadEndpoint
     #[cfg(not(target_os = "android"))]
     fn get_openvpn_relay(
-        query: &RelayQuery,
+        query: RelayQuery,
         config: &NormalSelectorConfig<'_>,
         parsed_relays: &ParsedRelays,
     ) -> Result<GetRelay, Error> {
+        assert_eq!(query.tunnel_protocol, Constraint::Only(TunnelType::OpenVpn));
         let exit =
-            Self::choose_openvpn_relay(query, config, parsed_relays).ok_or(Error::NoRelay)?;
-        let endpoint = Self::get_openvpn_endpoint(query, &exit, parsed_relays)?;
+            Self::choose_openvpn_relay(&query, config, parsed_relays).ok_or(Error::NoRelay)?;
+        let endpoint = Self::get_openvpn_endpoint(&query, &exit, parsed_relays)?;
         let bridge =
-            Self::get_openvpn_bridge(query, &exit, &endpoint.protocol, parsed_relays, config)?;
+            Self::get_openvpn_bridge(&query, &exit, &endpoint.protocol, parsed_relays, config)?;
 
+        // FIXME: This assert would be better to encode at the type level.
+        assert!(matches!(exit.endpoint_data, RelayEndpointData::Openvpn));
         Ok(GetRelay::OpenVpn {
             endpoint,
             exit,
@@ -964,6 +1011,7 @@ impl RelaySelector {
 
     /// # Returns
     /// A randomly selected relay that meets the specified constraints, or `None` if no suitable relay is found.
+    #[cfg(not(target_os = "android"))]
     fn choose_openvpn_relay(
         query: &RelayQuery,
         config: &NormalSelectorConfig<'_>,

mullvad-types/src/relay_constraints.rs

@@ -265,7 +265,7 @@ impl GeographicLocationConstraint {
         GeographicLocationConstraint::Hostname(country.into(), city.into(), hostname.into())
     }
 
-    // TODO(markus): Document
+    /// Check if `self` is _just_ a country. See [`GeographicLocationConstraint`] for more details.
     pub fn is_country(&self) -> bool {
         matches!(self, GeographicLocationConstraint::Country(_))
     }

mullvad-types/src/relay_list.rs

@@ -168,15 +168,6 @@ pub enum RelayEndpointData {
     Wireguard(WireguardRelayEndpointData),
 }
 
-impl RelayEndpointData {
-    pub fn unwrap_wireguard_ref(&self) -> &WireguardRelayEndpointData {
-        if let RelayEndpointData::Wireguard(wg) = &self {
-            return wg;
-        }
-        panic!("not a wireguard endpoint");
-    }
-}
-
 /// Data needed to connect to OpenVPN endpoints.
 #[derive(Debug, Default, Clone, Eq, PartialEq, Hash, Deserialize, Serialize)]
 pub struct OpenVpnEndpointData {