Set relay to current country on first start

Sets the default relay selection to the current country (as determined
by am.i.mullvad.net). If the current country does not have any relays
the country with the closest relay is choosen instead.

In non-release builds of the Android app we do not bundle a relay list
in the APK, and the relay list is fetched when the user logs in.

So one of the following can happen:

1. Geolocation request returns, we have a relay list.
2. Geolocation request returns, we do not yet have a relay list.
3. Relay list request returns, we have a geolocation.
4. Relay list request returns, we do not have a geolocation.

In 1. and 3. we can update the default location. In 2. we have to wait
until the relay list is fetched from the api until we can update the
default location. 4. is unlikely to happen but could happen if
am.i.mullvad is down.

10 files changed, 322 insertions, 2 deletions

diff --git a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/DefaultLocationTest.kt b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/DefaultLocationTest.kt
new file mode 100644
index 0000000000..935e1a28a5
--- /dev/null
+++ b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/DefaultLocationTest.kt
@@ -0,0 +1,44 @@
+package net.mullvad.mullvadvpn.test.e2e
+
+import java.io.File
+import net.mullvad.mullvadvpn.test.common.page.ConnectPage
+import net.mullvad.mullvadvpn.test.common.page.LoginPage
+import net.mullvad.mullvadvpn.test.common.page.on
+import net.mullvad.mullvadvpn.test.e2e.misc.AccountTestRule
+import org.json.JSONObject
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
+
+class DefaultLocationTest : EndToEndTest() {
+
+    @RegisterExtension @JvmField val accountTestRule = AccountTestRule()
+
+    @Test
+    fun testUpdateDefaultLocationFlag() {
+        app.launchAndEnsureOnLoginPage()
+
+        // The update_default_location flag should be set to true when first starting the app
+        assert(readUpdateDefaultLocationKeyFromSettings())
+
+        on<LoginPage> {
+            enterAccountNumber(accountTestRule.validAccountNumber)
+            clickLoginButton()
+        }
+
+        on<ConnectPage>()
+
+        // After we have logged in the daemon will have set the new default location so the
+        // flag should be false.
+        assertFalse(readUpdateDefaultLocationKeyFromSettings())
+    }
+
+    private fun readUpdateDefaultLocationKeyFromSettings(): Boolean {
+        val settings = File(targetApplication.filesDir, "settings.json")
+        if (!settings.isFile()) error("settings.json does not exist")
+
+        val text = settings.readText()
+        val json = JSONObject(text)
+        return json.getBoolean("update_default_location")
+    }
+}
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 9cdb61e942..62c4e443ea 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -45,7 +45,9 @@ use mullvad_encrypted_dns_proxy::state::EncryptedDnsProxyState;
 use mullvad_relay_selector::{RelaySelector, SelectorConfig};
 #[cfg(target_os = "android")]
 use mullvad_types::account::{PlayPurchase, PlayPurchasePaymentToken};
-use mullvad_types::relay_constraints::GeographicLocationConstraint;
+use mullvad_types::relay_constraints::{
+    GeographicLocationConstraint, LocationConstraint, RelayConstraints, WireguardConstraints,
+};
 #[cfg(any(windows, target_os = "android", target_os = "macos"))]
 use mullvad_types::settings::SplitApp;
 #[cfg(daita)]
@@ -413,6 +415,9 @@ pub enum DaemonCommand {
     ExportJsonSettings(ResponseTx<String, settings::patch::Error>),
     /// Request the current feature indicators.
     GetFeatureIndicators(oneshot::Sender<FeatureIndicators>),
+    // Updates the default (initial) country selection that the user will see when starting the
+    // app for the first time based on their current geolocation.
+    UpdateDefaultLocationCountry(ResponseTx<(), settings::Error>),
 
     // Debug features
     DisableRelay {
@@ -915,8 +920,13 @@ impl Daemon {
         api::forward_offline_state(api_availability.clone(), offline_state_rx);
 
         let relay_list_listener = management_interface.notifier().clone();
+        let internal_event_tx_clone = internal_event_tx.clone();
         let on_relay_list_update = move |relay_list: &RelayList| {
             relay_list_listener.notify_relay_list(relay_list.clone());
+            let (tx, _) = oneshot::channel();
+            let _ = internal_event_tx_clone.send(InternalDaemonEvent::Command(
+                DaemonCommand::UpdateDefaultLocationCountry(tx),
+            ));
         };
 
         let mut relay_list_updater = RelayListUpdater::spawn(
@@ -1303,6 +1313,13 @@ impl Daemon {
             _ => return,
         };
 
+        if self.settings.update_default_location {
+            let (tx, _) = oneshot::channel();
+            let _ = self.tx.send(InternalDaemonEvent::Command(
+                DaemonCommand::UpdateDefaultLocationCountry(tx),
+            ));
+        }
+
         self.management_interface
             .notifier()
             .notify_new_state(self.tunnel_state.clone());
@@ -1400,6 +1417,7 @@ impl Daemon {
             SubmitVoucher(tx, voucher) => self.on_submit_voucher(tx, voucher),
             GetRelayLocations(tx) => self.on_get_relay_locations(tx),
             UpdateRelayLocations => self.on_update_relay_locations().await,
+            UpdateDefaultLocationCountry(tx) => self.on_update_default_location(tx).await,
             LoginAccount(tx, account_number) => self.on_login_account(tx, account_number),
             LogoutAccount(tx) => self.on_logout_account(tx),
             GetDevice(tx) => self.on_get_device(tx),
@@ -1863,9 +1881,43 @@ impl Daemon {
         self.relay_list_updater.update().await;
     }
 
+    async fn on_update_default_location(&mut self, tx: ResponseTx<(), settings::Error>) {
+        log::info!(
+            "should_update_default_country: {}",
+            &self.settings.update_default_location
+        );
+
+        if self.settings.update_default_location
+            && let Some(location) = self.tunnel_state.get_location()
+            && let Some(country_code) = self.relay_selector.access_relays(|relays| {
+                relays
+                    .lookup_country_code_by_name(&location.country)
+                    .or_else(|| relays.get_nearest_country_with_relay(location))
+            })
+        {
+            let relay_settings = RelaySettings::Normal(RelayConstraints {
+                location: Constraint::Only(LocationConstraint::Location(
+                    GeographicLocationConstraint::Country(country_code.clone()),
+                )),
+                wireguard_constraints: WireguardConstraints {
+                    entry_location: Constraint::Only(LocationConstraint::Location(
+                        GeographicLocationConstraint::Country(country_code),
+                    )),
+                    ..Default::default()
+                },
+                ..Default::default()
+            });
+
+            self.on_set_relay_settings(tx, relay_settings).await;
+        } else {
+            let _ = tx.send(Ok(()));
+        }
+    }
+
     fn on_login_account(&mut self, tx: ResponseTx<(), Error>, account_number: String) {
         let account_manager = self.account_manager.clone();
         let availability = self.api_runtime.availability_handle();
+
         tokio::spawn(async move {
             let result = async {
                 account_manager
@@ -2324,6 +2376,19 @@ impl Daemon {
                 Self::oneshot_send(tx, Err(e), "set_relay_settings response");
             }
         }
+        if self.settings.update_default_location {
+            if let Err(e) = self
+                .settings
+                .update(move |settings| settings.update_default_location = false)
+                .await
+                .map_err(Error::SettingsError)
+            {
+                log::error!(
+                    "{}",
+                    e.display_chain_with_msg("Unable to save has_updated_default_country")
+                );
+            }
+        }
     }
 
     async fn on_set_allow_lan(&mut self, tx: ResponseTx<(), settings::Error>, allow_lan: bool) {
diff --git a/mullvad-daemon/src/settings/mod.rs b/mullvad-daemon/src/settings/mod.rs
index 8a0c98478f..6e286bd060 100644
--- a/mullvad-daemon/src/settings/mod.rs
+++ b/mullvad-daemon/src/settings/mod.rs
@@ -274,6 +274,13 @@ impl SettingsPersister {
         if crate::version::is_beta_version() {
             settings.show_beta_releases = true;
         }
+        // We only want to set this flag to true if the settings file hasn't been
+        // created yet so that we don't affect existing users' relay settings.
+        #[cfg(target_os = "android")]
+        {
+            settings.update_default_location = true;
+        }
+
         settings
     }
 
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 4d405bcde2..e9e13874ee 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -528,6 +528,7 @@ message Settings {
   ApiAccessMethodSettings api_access_methods = 12;
   repeated RelayOverride relay_overrides = 13;
   optional Recents recents = 14;
+  bool update_default_location = 15;
 }
 
 message RelayOverride {
diff --git a/mullvad-management-interface/src/types/conversions/settings.rs b/mullvad-management-interface/src/types/conversions/settings.rs
index 7b2e2b30d5..e9f68524ce 100644
--- a/mullvad-management-interface/src/types/conversions/settings.rs
+++ b/mullvad-management-interface/src/types/conversions/settings.rs
@@ -57,6 +57,7 @@ impl From<&mullvad_types::settings::Settings> for proto::Settings {
                 .map(proto::RelayOverride::from)
                 .collect(),
             recents: settings.recents.clone().map(proto::Recents::from),
+            update_default_location: settings.update_default_location,
         }
     }
 }
@@ -203,6 +204,7 @@ impl TryFrom<proto::Settings> for mullvad_types::settings::Settings {
                 api_access_methods_settings,
             )?,
             recents: Some(vec![]),
+            update_default_location: settings.update_default_location,
         })
     }
 }
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index c780e5b5bb..8c5f729d56 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -465,6 +465,11 @@ impl RelaySelector {
         parsed_relays.original_list().clone()
     }
 
+    pub fn access_relays<T>(&mut self, access_fn: impl Fn(&RelayList) -> T) -> T {
+        let parsed_relays = self.parsed_relays.lock().unwrap();
+        access_fn(parsed_relays.original_list())
+    }
+
     pub fn etag(&self) -> Option<String> {
         self.parsed_relays.lock().unwrap().etag()
     }
diff --git a/mullvad-types/src/location.rs b/mullvad-types/src/location.rs
index 52f237043b..ed7a128108 100644
--- a/mullvad-types/src/location.rs
+++ b/mullvad-types/src/location.rs
@@ -55,6 +55,21 @@ impl From<Location> for Coordinates {
     }
 }
 
+impl From<&GeoIpLocation> for Coordinates {
+    fn from(location: &GeoIpLocation) -> Self {
+        Self {
+            latitude: location.latitude,
+            longitude: location.longitude,
+        }
+    }
+}
+
+impl From<GeoIpLocation> for Coordinates {
+    fn from(location: GeoIpLocation) -> Self {
+        Coordinates::from(&location)
+    }
+}
+
 impl Coordinates {
     /// Computes the approximate midpoint of a set of locations.
     ///
diff --git a/mullvad-types/src/relay_list.rs b/mullvad-types/src/relay_list.rs
index c6f1693c14..b8f528bba2 100644
--- a/mullvad-types/src/relay_list.rs
+++ b/mullvad-types/src/relay_list.rs
@@ -1,4 +1,4 @@
-use crate::location::{CityCode, CountryCode, Location};
+use crate::location::{CityCode, Coordinates, CountryCode, Location};
 use serde::{Deserialize, Serialize};
 use std::{
     collections::HashSet,
@@ -30,6 +30,41 @@ impl RelayList {
             .find(|country| country.code == country_code)
     }
 
+    pub fn lookup_country_code_by_name(&self, country_name: &str) -> Option<CountryCode> {
+        self.countries
+            .iter()
+            .find(|country| country.name == country_name)
+            .map(|country| country.code.clone())
+    }
+
+    /// Returns the closest (geographical distance) country that has a relay for a given location.
+    pub fn get_nearest_country_with_relay(
+        &self,
+        location: impl Into<Coordinates>,
+    ) -> Option<CountryCode> {
+        if self.countries.is_empty() {
+            return None;
+        }
+
+        let location = location.into();
+
+        let mut min_dist = f64::MAX;
+        let mut min_dist_country = &self.countries[0];
+
+        for country in &self.countries {
+            for city in &country.cities {
+                for relay in &city.relays {
+                    let distance = relay.location.distance_from(location);
+                    if distance < min_dist {
+                        min_dist = distance;
+                        min_dist_country = country;
+                    }
+                }
+            }
+        }
+        Some(min_dist_country.code.clone())
+    }
+
     /// Return a flat iterator of all [`Relay`]s
     pub fn relays(&self) -> impl Iterator<Item = &Relay> + Clone + '_ {
         self.countries
@@ -384,3 +419,136 @@ impl ShadowsocksEndpointData {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use talpid_types::net::wireguard::PublicKey;
+
+    #[test]
+    fn test_get_nearest_country_with_relay() {
+        let location_sweden = Location {
+            country: "Sweden".to_string(),
+            country_code: "se".to_string(),
+            city: "Gothenburg".to_string(),
+            city_code: "got".to_string(),
+            latitude: 57.71,
+            longitude: 11.97,
+        };
+
+        let location_japan = Location {
+            country: "Japan".to_string(),
+            country_code: "jp".to_string(),
+            city: "Osaka".to_string(),
+            city_code: "osa".to_string(),
+            latitude: 34.67231,
+            longitude: 135.484802,
+        };
+
+        let location_south_korea = Location {
+            country: "South Korea".to_string(),
+            country_code: "sk".to_string(),
+            city: "Seoul".to_string(),
+            city_code: "seo".to_string(),
+            latitude: 37.532600,
+            longitude: 127.024612,
+        };
+
+        let location_germany = Location {
+            country: "Germany".to_string(),
+            country_code: "ger".to_string(),
+            city: "Berlin".to_string(),
+            city_code: "ber".to_string(),
+            latitude: 52.5200080,
+            longitude: 13.404954,
+        };
+
+        let countries = vec![
+            RelayListCountry {
+                name: "Sweden".to_string(),
+                code: "se".to_string(),
+                cities: vec![RelayListCity {
+                    name: "Gothenburg".to_string(),
+                    code: "got".to_string(),
+                    latitude: 57.70887,
+                    longitude: 11.97456,
+                    relays: vec![Relay {
+                        hostname: "se9-wireguard".to_string(),
+                        ipv4_addr_in: "185.213.154.68".parse().unwrap(),
+                        ipv6_addr_in: Some("2a03:1b20:5:f011::a09f".parse().unwrap()),
+                        overridden_ipv4: false,
+                        overridden_ipv6: false,
+                        include_in_country: true,
+                        active: true,
+                        owned: true,
+                        provider: "provider0".to_string(),
+                        weight: 1,
+                        endpoint_data: RelayEndpointData::Wireguard(
+                            WireguardRelayEndpointData::new(
+                                PublicKey::from_base64(
+                                    "BLNHNoGO88LjV/wDBa7CUUwUzPq/fO2UwcGLy56hKy4=",
+                                )
+                                .unwrap(),
+                            ),
+                        ),
+                        location: location_sweden.clone(),
+                    }],
+                }],
+            },
+            //  34.672314, 135.484802.
+            RelayListCountry {
+                name: "Japan".to_string(),
+                code: "jp".to_string(),
+                cities: vec![RelayListCity {
+                    name: "Osaka".to_string(),
+                    code: "osa".to_string(),
+                    latitude: 34.672314,
+                    longitude: 135.484802,
+                    relays: vec![Relay {
+                        hostname: "jp9-wireguard".to_string(),
+                        ipv4_addr_in: "194.114.136.3".parse().unwrap(),
+                        ipv6_addr_in: Some("2404:1b20:5:f011::a09f".parse().unwrap()),
+                        overridden_ipv4: false,
+                        overridden_ipv6: false,
+                        include_in_country: true,
+                        active: true,
+                        owned: true,
+                        provider: "provider0".to_string(),
+                        weight: 1,
+                        endpoint_data: RelayEndpointData::Wireguard(
+                            WireguardRelayEndpointData::new(
+                                PublicKey::from_base64(
+                                    "BLNHNoGO88LjV/wDBa7CUUwUzPq/fO2UwcGLy56hKy4=",
+                                )
+                                .unwrap(),
+                            ),
+                        ),
+                        location: location_japan.clone(),
+                    }],
+                }],
+            },
+        ];
+
+        let relay_list = RelayList {
+            countries,
+            ..Default::default()
+        };
+
+        assert_eq!(
+            relay_list.get_nearest_country_with_relay(location_sweden),
+            Some("se".to_string())
+        );
+        assert_eq!(
+            relay_list.get_nearest_country_with_relay(location_japan),
+            Some("jp".to_string())
+        );
+        assert_eq!(
+            relay_list.get_nearest_country_with_relay(location_germany),
+            Some("se".to_string())
+        );
+        assert_eq!(
+            relay_list.get_nearest_country_with_relay(location_south_korea),
+            Some("jp".to_string())
+        );
+    }
+}
diff --git a/mullvad-types/src/settings/mod.rs b/mullvad-types/src/settings/mod.rs
index b71280a7e9..ae55fd3d79 100644
--- a/mullvad-types/src/settings/mod.rs
+++ b/mullvad-types/src/settings/mod.rs
@@ -81,6 +81,8 @@ pub struct Settings {
     pub custom_lists: CustomListsSettings,
     /// API access methods
     pub api_access_methods: access_method::Settings,
+    // If the default location in `relay_settings` should be updated based on the user's geolocation.
+    pub update_default_location: bool,
     /// If the daemon should allow communication with private (LAN) networks.
     pub allow_lan: bool,
     /// Extra level of kill switch. When this setting is on, the disconnected state will block
@@ -257,6 +259,7 @@ impl Default for Settings {
                 },
                 ..Default::default()
             }),
+            update_default_location: false,
             bridge_settings: BridgeSettings::default(),
             obfuscation_settings: ObfuscationSettings {
                 selected_obfuscation: SelectedObfuscation::Auto,
diff --git a/mullvad-types/src/states.rs b/mullvad-types/src/states.rs
index e0c097236d..7479708894 100644
--- a/mullvad-types/src/states.rs
+++ b/mullvad-types/src/states.rs
@@ -124,4 +124,14 @@ impl TunnelState {
             None => None,
         }
     }
+
+    /// Returns the geolocation of the tunnel if it exists.
+    pub fn get_location(&self) -> Option<&GeoIpLocation> {
+        match self {
+            TunnelState::Connected { location, .. }
+            | TunnelState::Connecting { location, .. }
+            | TunnelState::Disconnected { location, .. } => location.as_ref(),
+            _ => None,
+        }
+    }
 }