Do not panic when reading public key of relay

Avoid panic when reading public key of relay due to unwrapping an
unexpected relay type. The unwrap happened if an OpenVPN relay was
selected on Android, which should not happen in the first place.

6 files changed, 103 insertions, 52 deletions

diff --git a/mullvad-daemon/src/tunnel.rs b/mullvad-daemon/src/tunnel.rs
index f230a3c4b2..31f4686df7 100644
--- a/mullvad-daemon/src/tunnel.rs
+++ b/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
                     }
diff --git a/mullvad-relay-selector/src/relay_selector/detailer.rs b/mullvad-relay-selector/src/relay_selector/detailer.rs
index b0e6e3b0ea..374d92da09 100644
--- a/mullvad-relay-selector/src/relay_selector/detailer.rs
+++ b/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)
diff --git a/mullvad-relay-selector/src/relay_selector/matcher.rs b/mullvad-relay-selector/src/relay_selector/matcher.rs
index c287e30fe6..f06e04224f 100644
--- a/mullvad-relay-selector/src/relay_selector/matcher.rs
+++ b/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(_))
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index 4967267e79..424ca6385c 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/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<'_>,
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 11f27156b3..6d6620677e 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/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(_))
     }
diff --git a/mullvad-types/src/relay_list.rs b/mullvad-types/src/relay_list.rs
index 5711b812c1..77cc621728 100644
--- a/mullvad-types/src/relay_list.rs
+++ b/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 {