Merge branch 'fix-custom-last-location'

3 files changed, 42 insertions, 18 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc997580dd..a51712bdca 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -39,6 +39,8 @@ Line wrap the file at 100 chars.                                              Th
 - Experimental: Upgrade the support for quantum-resistant WireGuard tunnels to a newer protocol.
 
 ### Fixed
+- Reset known location when connecting to a custom relay.
+
 #### Linux
 - Fix app crashing immediately when using some icon themes.
 
diff --git a/mullvad-cli/src/format.rs b/mullvad-cli/src/format.rs
index 0e0f24c8dd..5a2fd8b654 100644
--- a/mullvad-cli/src/format.rs
+++ b/mullvad-cli/src/format.rs
@@ -34,7 +34,7 @@ pub fn print_state(state: &TunnelState, verbose: bool) {
 
 fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) -> String {
     let endpoint = relay_info.tunnel_endpoint.as_ref().unwrap();
-    let location = &relay_info.location.as_ref().unwrap();
+    let location = &relay_info.location.as_ref();
 
     let prefix_separator = if verbose { "\n\t" } else { " " };
     let mut obfuscator_overlaps = false;
@@ -43,7 +43,10 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
         let mut address = Cow::Borrowed(endpoint.address.as_str());
         let mut protocol = endpoint.protocol;
         if let Some(obfuscator) = endpoint.obfuscation.as_ref() {
-            if location.hostname == location.obfuscator_hostname {
+            if location
+                .map(|l| l.hostname == l.obfuscator_hostname)
+                .unwrap_or(false)
+            {
                 obfuscator_overlaps = true;
                 address = Cow::Owned(format!("{}:{}", obfuscator.address, obfuscator.port));
                 protocol = obfuscator.protocol;
@@ -51,11 +54,16 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
         };
 
         let exit = format_endpoint(
-            &location.hostname,
+            location.map(|l| l.hostname.as_str()),
             protocol,
-            Some(address).filter(|_| verbose).as_deref(),
+            &*address,
+            verbose,
         );
-        format!("{exit} in {}, {}", &location.city, &location.country)
+        if let Some(location) = location {
+            format!("{exit} in {}, {}", &location.city, &location.country)
+        } else {
+            exit
+        }
     };
 
     let first_hop = endpoint.entry_endpoint.as_ref().map(|entry| {
@@ -63,16 +71,20 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
         let mut protocol = entry.protocol;
         if let Some(obfuscator) = endpoint.obfuscation.as_ref() {
             obfuscator_overlaps = true;
-            if location.entry_hostname == location.obfuscator_hostname {
+            if location
+                .map(|l| l.hostname == l.obfuscator_hostname)
+                .unwrap_or(false)
+            {
                 address = &obfuscator.address;
                 protocol = obfuscator.protocol;
             }
         };
 
         let endpoint = format_endpoint(
-            &location.entry_hostname,
+            location.map(|l| l.entry_hostname.as_str()),
             protocol,
-            Some(address).filter(|_| verbose),
+            address,
+            verbose,
         );
         format!("{prefix_separator}via {endpoint}")
     });
@@ -80,9 +92,10 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
     let obfuscator = endpoint.obfuscation.as_ref().map(|obfuscator| {
         if !obfuscator_overlaps {
             let endpoint_str = format_endpoint(
-                &location.obfuscator_hostname,
+                location.map(|l| l.obfuscator_hostname.as_str()),
                 obfuscator.protocol,
-                Some(obfuscator.address.as_str()).filter(|_| verbose),
+                obfuscator.address.as_str(),
+                verbose,
             );
             format!("{prefix_separator}obfuscated via {endpoint_str}")
         } else {
@@ -92,9 +105,10 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
 
     let bridge = endpoint.proxy.as_ref().map(|proxy| {
         let proxy_endpoint = format_endpoint(
-            &location.bridge_hostname,
+            location.map(|l| l.bridge_hostname.as_str()),
             proxy.protocol,
-            Some(proxy.address.as_str()).filter(|_| verbose),
+            proxy.address.as_str(),
+            verbose,
         );
 
         format!("{prefix_separator}via {proxy_endpoint}")
@@ -148,14 +162,22 @@ fn convert_obfuscator_type(obfuscator: i32) -> &'static str {
     }
 }
 
-fn format_endpoint(hostname: &String, protocol_enum: i32, addr: Option<&str>) -> String {
+fn format_endpoint(
+    hostname: Option<&str>,
+    protocol_enum: i32,
+    addr: &str,
+    verbose: bool,
+) -> String {
     let protocol = format_protocol(
         TransportProtocol::from_i32(protocol_enum).expect("invalid transport protocol"),
     );
-    let sockaddr_suffix = addr
-        .map(|addr| format!(" ({addr}/{protocol})"))
-        .unwrap_or_else(String::new);
-    format!("{hostname}{sockaddr_suffix}")
+
+    match (hostname, verbose) {
+        (Some(hostname), true) => format!("{hostname} ({addr}/{protocol})"),
+        (None, true) => format!("{addr}/{protocol}"),
+        (Some(hostname), false) => hostname.to_string(),
+        (None, false) => addr.to_string(),
+    }
 }
 
 fn print_error_state(error_state: &ErrorState) {
diff --git a/mullvad-daemon/src/tunnel.rs b/mullvad-daemon/src/tunnel.rs
index 18b8bf2c20..7f7fb0c8b3 100644
--- a/mullvad-daemon/src/tunnel.rs
+++ b/mullvad-daemon/src/tunnel.rs
@@ -41,7 +41,6 @@ struct InnerParametersGenerator {
     tunnel_options: TunnelOptions,
     account_manager: AccountManagerHandle,
 
-    // TODO: Move this to `RelaySelector`?
     last_generated_relays: Option<LastSelectedRelays>,
 }
 
@@ -124,6 +123,7 @@ impl InnerParametersGenerator {
         let _data = self.device().await?;
         match self.relay_selector.get_relay(retry_attempt) {
             Ok((SelectedRelay::Custom(custom_relay), _bridge, _obfsucator)) => {
+                self.last_generated_relays = None;
                 custom_relay
                     // TODO: generate proxy settings for custom tunnels
                     .to_tunnel_parameters(self.tunnel_options.clone(), None)