2 files changed, 39 insertions, 13 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
index 1b676c96ec..ddcb9079fd 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
@@ -38,6 +38,8 @@ class ConnectFragment : Fragment() {
     private var fetchInitialStateJob = fetchInitialState()
     private var generateWireguardKeyJob = generateWireguardKey()
 
+    private var lastKnownRealLocation: GeoIpLocation? = null
+
     private var activeAction: Job? = null
     private var attachListenerJob: Job? = null
     private var updateViewJob: Job? = null
@@ -64,7 +66,7 @@ class ConnectFragment : Fragment() {
         headerBar = HeaderBar(view, context!!)
         notificationBanner = NotificationBanner(view)
         status = ConnectionStatus(view, context!!)
-        locationInfo = LocationInfo(view)
+        locationInfo = LocationInfo(view, daemon)
 
         actionButton = ConnectActionButton(view)
         actionButton.apply {
@@ -154,7 +156,6 @@ class ConnectFragment : Fragment() {
 
     private fun disconnect() {
         activeAction?.cancel()
-        clearLocation()
 
         activeAction = GlobalScope.launch(Dispatchers.Default) {
             daemon.await().disconnect()
@@ -177,16 +178,7 @@ class ConnectFragment : Fragment() {
         headerBar.setState(state)
         notificationBanner.setState(state)
         status.setState(state)
-
-        locationInfo.location = fetchLocation().await()
-    }
-
-    private fun fetchLocation() = GlobalScope.async(Dispatchers.Default) {
-        daemon.await().getCurrentLocation()
-    }
-
-    private fun clearLocation() = GlobalScope.launch(Dispatchers.Main) {
-        locationInfo.location = null
+        locationInfo.setState(state)
     }
 
     private fun openSwitchLocationScreen() {
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/LocationInfo.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/LocationInfo.kt
index 5a3290a523..d49db5db7a 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/LocationInfo.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/LocationInfo.kt
@@ -1,24 +1,58 @@
 package net.mullvad.mullvadvpn
 
+import kotlinx.coroutines.async
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+
 import android.view.View
 import android.widget.TextView
 
 import net.mullvad.mullvadvpn.model.GeoIpLocation
+import net.mullvad.mullvadvpn.model.TunnelStateTransition
 
-class LocationInfo(val parentView: View) {
+class LocationInfo(val parentView: View, val daemon: Deferred<MullvadDaemon>) {
     private val country: TextView = parentView.findViewById(R.id.country)
     private val city: TextView = parentView.findViewById(R.id.city)
     private val hostname: TextView = parentView.findViewById(R.id.hostname)
 
+    private var lastKnownRealLocation: GeoIpLocation? = null
+
     var location: GeoIpLocation? = null
         set(value) {
             field = value
             updateViews(value)
         }
 
+    fun setState(state: TunnelStateTransition) {
+        when (state) {
+            is TunnelStateTransition.Disconnected -> fetchRealLocation()
+            is TunnelStateTransition.Connecting -> fetchRelayLocation()
+            is TunnelStateTransition.Connected -> fetchRelayLocation()
+            is TunnelStateTransition.Disconnecting -> location = lastKnownRealLocation
+            is TunnelStateTransition.Blocked -> location = null
+        }
+    }
+
     fun updateViews(location: GeoIpLocation?) {
         country.text = location?.country ?: ""
         city.text = location?.city ?: ""
         hostname.text = location?.hostname ?: ""
     }
+
+    private fun fetchRealLocation() = GlobalScope.launch(Dispatchers.Main) {
+        val realLocation = fetchLocation().await()
+
+        lastKnownRealLocation = realLocation
+        location = realLocation
+    }
+
+    private fun fetchRelayLocation() = GlobalScope.launch(Dispatchers.Main) {
+        location = fetchLocation().await()
+    }
+
+    private fun fetchLocation() = GlobalScope.async(Dispatchers.Default) {
+        daemon.await().getCurrentLocation()
+    }
 }