Merge branch 'connect-when-selecting-location'

5 files changed, 188 insertions, 141 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
index c31b1691dc..53e00a4c31 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
@@ -1,28 +1,21 @@
 package net.mullvad.mullvadvpn
 
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.Job
 
-import android.app.Activity
 import android.content.Context
-import android.content.Intent
-import android.net.VpnService
 import android.os.Bundle
-import android.os.Handler
 import android.support.v4.app.Fragment
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
-import android.widget.Button
 import android.widget.ImageButton
 
+import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.dataproxy.LocationInfoCache
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
-import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.TunnelStateTransition
 
 class ConnectFragment : Fragment() {
@@ -34,27 +27,19 @@ class ConnectFragment : Fragment() {
     private lateinit var locationInfo: LocationInfo
 
     private lateinit var parentActivity: MainActivity
+    private lateinit var connectionProxy: ConnectionProxy
     private lateinit var locationInfoCache: LocationInfoCache
     private lateinit var relayListListener: RelayListListener
 
-    private var daemon = CompletableDeferred<MullvadDaemon>()
-    private var vpnPermission = CompletableDeferred<Unit>()
-
-    private var fetchInitialStateJob = fetchInitialState()
-    private var generateWireguardKeyJob = generateWireguardKey()
-
-    private var activeAction: Job? = null
-    private var attachListenerJob: Job? = null
     private var updateViewJob: Job? = null
-    private var waitForDaemonJob: Job? = null
 
     override fun onAttach(context: Context) {
         super.onAttach(context)
 
         parentActivity = context as MainActivity
+        connectionProxy = parentActivity.connectionProxy
         locationInfoCache = parentActivity.locationInfoCache
         relayListListener = parentActivity.relayListListener
-        waitForDaemonJob = waitForDaemon(parentActivity.daemon)
     }
 
     override fun onCreateView(
@@ -75,15 +60,22 @@ class ConnectFragment : Fragment() {
 
         actionButton = ConnectActionButton(view)
         actionButton.apply {
-            onConnect = { connect() }
-            onCancel = { disconnect() }
-            onDisconnect = { disconnect() }
+            onConnect = { connectionProxy.connect() }
+            onCancel = { connectionProxy.disconnect() }
+            onDisconnect = { connectionProxy.disconnect() }
         }
 
         switchLocationButton = SwitchLocationButton(view)
         switchLocationButton.onClick = { openSwitchLocationScreen() }
 
-        attachListenerJob = attachListener()
+        updateView(connectionProxy.uiState)
+
+        connectionProxy.onUiStateChange = { uiState ->
+            updateViewJob?.cancel()
+            updateViewJob = GlobalScope.launch(Dispatchers.Main) {
+                updateView(uiState)
+            }
+        }
 
         return view
     }
@@ -106,110 +98,23 @@ class ConnectFragment : Fragment() {
         locationInfo.onDestroy()
         switchLocationButton.onDestroy()
 
-        waitForDaemonJob?.cancel()
-        attachListenerJob?.cancel()
-
-        detachListener()
-
-        generateWireguardKeyJob.cancel()
+        connectionProxy.onUiStateChange = null
         updateViewJob?.cancel()
 
         super.onDestroyView()
     }
 
-    override fun onActivityResult(requestCode: Int, resultCode: Int, resultData: Intent?) {
-        if (resultCode == Activity.RESULT_OK) {
-            vpnPermission.complete(Unit)
-        }
-    }
-
-    private fun waitForDaemon(originalDaemon: Deferred<MullvadDaemon>) =
-            GlobalScope.launch(Dispatchers.Default) {
-        daemon.complete(originalDaemon.await())
-    }
-
-    private fun attachListener() = GlobalScope.launch(Dispatchers.Default) {
-        daemon.await().onTunnelStateChange = { state ->
-            synchronized(this@ConnectFragment) {
-                updateViewJob = updateView(state)
-            }
-        }
-    }
-
-    private fun detachListener() = GlobalScope.launch(Dispatchers.Default) {
-        daemon.await().onTunnelStateChange = null
-    }
-
-    private fun fetchInitialState() = GlobalScope.launch(Dispatchers.Default) {
-        val state = daemon.await().getState()
-
-        synchronized(this@ConnectFragment) {
-            updateViewJob = updateViewJob ?: updateView(state)
-        }
-    }
-
-    private fun generateWireguardKey() = GlobalScope.launch(Dispatchers.Default) {
-        val daemon = this@ConnectFragment.daemon.await()
-        val key = daemon.getWireguardKey()
-
-        if (key == null) {
-            daemon.generateWireguardKey()
-        }
-    }
-
-    private fun requestVpnPermission() {
-        val intent = VpnService.prepare(parentActivity)
-
-        vpnPermission = CompletableDeferred<Unit>()
-
-        if (intent != null) {
-            startActivityForResult(intent, 0)
-        } else {
-            onActivityResult(0, Activity.RESULT_OK, null)
-        }
-    }
-
-    private fun connect() {
-        updateViewToPreConnecting()
-        activeAction?.cancel()
-
-        requestVpnPermission()
+    private fun updateView(uiState: TunnelStateTransition) {
+        val realState = connectionProxy.state
 
-        activeAction = GlobalScope.launch(Dispatchers.Default) {
-            vpnPermission.await()
-            generateWireguardKeyJob.join()
-            daemon.await().connect()
-        }
-    }
-
-    private fun disconnect() {
-        updateView(TunnelStateTransition.Disconnecting())
-        activeAction?.cancel()
-
-        activeAction = GlobalScope.launch(Dispatchers.Default) {
-            daemon.await().disconnect()
-        }
-    }
-
-    private fun updateViewToPreConnecting() {
-        val connecting = TunnelStateTransition.Connecting()
-        val disconnected = TunnelStateTransition.Disconnected()
-
-        headerBar.setState(disconnected)
-
-        actionButton.state = connecting
-        notificationBanner.setState(connecting)
-        status.setState(connecting)
-    }
+        locationInfoCache.state = realState
+        headerBar.setState(realState)
 
-    private fun updateView(state: TunnelStateTransition) = GlobalScope.launch(Dispatchers.Main) {
-        actionButton.state = state
-        switchLocationButton.state = state
+        actionButton.state = uiState
+        switchLocationButton.state = uiState
 
-        headerBar.setState(state)
-        notificationBanner.setState(state)
-        status.setState(state)
-        locationInfoCache.setState(state)
+        notificationBanner.setState(uiState)
+        status.setState(uiState)
     }
 
     private fun openSwitchLocationScreen() {
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt
index 457cefa3dc..b75a408ece 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt
@@ -9,14 +9,17 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.Job
 
+import android.app.Activity
 import android.content.ComponentName
 import android.content.Intent
 import android.content.ServiceConnection
+import android.net.VpnService
 import android.os.Bundle
 import android.os.IBinder
 import android.support.v4.app.FragmentActivity
 
 import net.mullvad.mullvadvpn.dataproxy.AccountCache
+import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.dataproxy.LocationInfoCache
 import net.mullvad.mullvadvpn.dataproxy.MullvadProblemReport
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
@@ -27,9 +30,12 @@ import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.relaylist.RelayList
 
 class MainActivity : FragmentActivity() {
+    private var vpnPermission: CompletableDeferred<Boolean>? = null
+
     var daemon = CompletableDeferred<MullvadDaemon>()
         private set
 
+    val connectionProxy = ConnectionProxy(this)
     val locationInfoCache = LocationInfoCache(daemon)
     val problemReport = MullvadProblemReport()
     var settingsListener = SettingsListener(this)
@@ -71,6 +77,14 @@ class MainActivity : FragmentActivity() {
         bindService(intent, serviceConnection, 0)
     }
 
+    override fun onActivityResult(requestCode: Int, resultCode: Int, resultData: Intent?) {
+        if (resultCode == Activity.RESULT_OK) {
+            vpnPermission?.complete(true)
+        } else {
+            vpnPermission?.complete(false)
+        }
+    }
+
     override fun onStop() {
         unbindService(serviceConnection)
 
@@ -102,6 +116,21 @@ class MainActivity : FragmentActivity() {
         }
     }
 
+    fun requestVpnPermission(): Deferred<Boolean> {
+        val intent = VpnService.prepare(this)
+        val request = CompletableDeferred<Boolean>()
+
+        vpnPermission = request
+
+        if (intent != null) {
+            startActivityForResult(intent, 0)
+        } else {
+            request.complete(true)
+        }
+
+        return request
+    }
+
     private fun addInitialFragment() {
         supportFragmentManager?.beginTransaction()?.apply {
             add(R.id.main_fragment, LaunchFragment())
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt
index f0d9af7913..eae0ed75f2 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt
@@ -16,6 +16,7 @@ import android.view.ViewGroup
 import android.widget.ImageButton
 import android.widget.ViewSwitcher
 
+import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
 import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.LocationConstraint
@@ -27,6 +28,7 @@ import net.mullvad.mullvadvpn.relaylist.RelayListAdapter
 
 class SelectLocationFragment : Fragment() {
     private lateinit var parentActivity: MainActivity
+    private lateinit var connectionProxy: ConnectionProxy
     private lateinit var relayListListener: RelayListListener
 
     private lateinit var relayListContainer: ViewSwitcher
@@ -38,6 +40,7 @@ class SelectLocationFragment : Fragment() {
     init {
         relayListAdapter.onSelect = { relayItem ->
             updateLocationConstraint(relayItem)
+            connectionProxy.connect()
             close()
         }
     }
@@ -46,6 +49,7 @@ class SelectLocationFragment : Fragment() {
         super.onAttach(context)
 
         parentActivity = context as MainActivity
+        connectionProxy = parentActivity.connectionProxy
         relayListListener = parentActivity.relayListListener
     }
 
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/ConnectionProxy.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/ConnectionProxy.kt
new file mode 100644
index 0000000000..fb74e70160
--- /dev/null
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/ConnectionProxy.kt
@@ -0,0 +1,93 @@
+package net.mullvad.mullvadvpn.dataproxy
+
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.Job
+
+import net.mullvad.mullvadvpn.MainActivity
+import net.mullvad.mullvadvpn.model.TunnelStateTransition
+
+class ConnectionProxy(val parentActivity: MainActivity) {
+    val daemon = parentActivity.daemon
+
+    private var activeAction: Job? = null
+
+    private val attachListenerJob = attachListener()
+    private val fetchInitialStateJob = fetchInitialState()
+
+    private var realState: TunnelStateTransition? = null
+        set(value) {
+            field = value
+            uiState = value ?: TunnelStateTransition.Disconnected()
+        }
+
+    val state: TunnelStateTransition
+        get() {
+            return realState ?: TunnelStateTransition.Disconnected()
+        }
+
+    var uiState: TunnelStateTransition = TunnelStateTransition.Disconnected()
+        private set(value) {
+            field = value
+            onUiStateChange?.invoke(value)
+        }
+
+    var onUiStateChange: ((TunnelStateTransition) -> Unit)? = null
+
+    fun connect() {
+        uiState = TunnelStateTransition.Connecting()
+
+        cancelActiveAction()
+
+        val vpnPermission = parentActivity.requestVpnPermission()
+
+        activeAction = GlobalScope.launch(Dispatchers.Default) {
+            if (vpnPermission.await()) {
+                daemon.await().connect()
+            }
+        }
+    }
+
+    fun disconnect() {
+        uiState = TunnelStateTransition.Disconnecting()
+
+        cancelActiveAction()
+        activeAction = GlobalScope.launch(Dispatchers.Default) {
+            daemon.await().disconnect()
+        }
+    }
+
+    fun cancelActiveAction() {
+        activeAction?.cancel()
+    }
+
+    fun onDestroy() {
+        attachListenerJob.cancel()
+        detachListener()
+        fetchInitialStateJob.cancel()
+        cancelActiveAction()
+    }
+
+    private fun fetchInitialState() = GlobalScope.launch(Dispatchers.Default) {
+        val initialState = daemon.await().getState()
+
+        synchronized(this) {
+            if (realState == null) {
+                realState = initialState
+            }
+        }
+    }
+
+    private fun attachListener() = GlobalScope.launch(Dispatchers.Default) {
+        daemon.await().onTunnelStateChange = { newState ->
+            synchronized(this) {
+                realState = newState
+            }
+        }
+    }
+
+    private fun detachListener() = GlobalScope.launch(Dispatchers.Default) {
+        daemon.await().onTunnelStateChange = null
+    }
+}
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/LocationInfoCache.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/LocationInfoCache.kt
index 98e529fbb8..4f5684b157 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/LocationInfoCache.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/LocationInfoCache.kt
@@ -27,18 +27,18 @@ class LocationInfoCache(val daemon: Deferred<MullvadDaemon>) {
             notifyNewLocation()
         }
 
-    fun setState(state: TunnelStateTransition) {
-        activeFetch?.cancel()
-        activeFetch = null
+    var state: TunnelStateTransition = TunnelStateTransition.Disconnected()
+        set(value) {
+            field = value
 
-        when (state) {
-            is TunnelStateTransition.Disconnected -> activeFetch = fetchRealLocation()
-            is TunnelStateTransition.Connecting -> activeFetch = fetchRelayLocation()
-            is TunnelStateTransition.Connected -> activeFetch = fetchRelayLocation()
-            is TunnelStateTransition.Disconnecting -> location = lastKnownRealLocation
-            is TunnelStateTransition.Blocked -> location = null
+            when (value) {
+                is TunnelStateTransition.Disconnected -> fetchLocation()
+                is TunnelStateTransition.Connecting -> fetchLocation()
+                is TunnelStateTransition.Connected -> fetchLocation()
+                is TunnelStateTransition.Disconnecting -> location = lastKnownRealLocation
+                is TunnelStateTransition.Blocked -> location = null
+            }
         }
-    }
 
     fun notifyNewLocation() {
         val location = this.location
@@ -49,24 +49,40 @@ class LocationInfoCache(val daemon: Deferred<MullvadDaemon>) {
         onNewLocation?.invoke(country, city, hostname)
     }
 
-    private fun fetchRealLocation() = GlobalScope.launch(Dispatchers.Main) {
-        var realLocation: GeoIpLocation? = null
-        var remainingAttempts = 10
+    private fun fetchLocation() {
+        val previousFetch = activeFetch
+        val initialState = state
 
-        while (realLocation == null && remainingAttempts > 0) {
-            realLocation = fetchLocation().await()
-            remainingAttempts -= 1
-        }
+        activeFetch = GlobalScope.launch(Dispatchers.Main) {
+            var newLocation: GeoIpLocation? = null
 
-        lastKnownRealLocation = realLocation
-        location = realLocation
-    }
+            previousFetch?.join()
 
-    private fun fetchRelayLocation() = GlobalScope.launch(Dispatchers.Main) {
-        location = fetchLocation().await()
+            while (newLocation == null && shouldRetryFetch() && state == initialState) {
+                newLocation = executeFetch().await()
+            }
+
+            if (newLocation != null && state == initialState) {
+                when (state) {
+                    is TunnelStateTransition.Disconnected -> {
+                        lastKnownRealLocation = newLocation
+                        location = newLocation
+                    }
+                    is TunnelStateTransition.Connecting -> location = newLocation
+                    is TunnelStateTransition.Connected -> location = newLocation
+                }
+            }
+        }
     }
 
-    private fun fetchLocation() = GlobalScope.async(Dispatchers.Default) {
+    private fun executeFetch() = GlobalScope.async(Dispatchers.Default) {
         daemon.await().getCurrentLocation()
     }
+
+    private fun shouldRetryFetch(): Boolean {
+        val state = this.state
+
+        return state is TunnelStateTransition.Disconnected ||
+            state is TunnelStateTransition.Connected
+    }
 }