Adapt device repo to new service management

9 files changed, 82 insertions, 60 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index ea1ae2b79b..26326fd002 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -6,6 +6,7 @@ import kotlinx.coroutines.Dispatchers
 import net.mullvad.mullvadvpn.applist.ApplicationsIconManager
 import net.mullvad.mullvadvpn.applist.ApplicationsProvider
 import net.mullvad.mullvadvpn.ipc.EventDispatcher
+import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.SplitTunneling
 import net.mullvad.mullvadvpn.viewmodel.SplitTunnelingViewModel
@@ -33,6 +34,7 @@ val uiModule = module {
     }
 
     single { ServiceConnectionManager(androidContext()) }
+    single { DeviceRepository(get()) }
 }
 const val APPS_SCOPE = "APPS_SCOPE"
 const val SERVICE_CONNECTION_SCOPE = "SERVICE_CONNECTION_SCOPE"
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt
index 526b5356f4..4b5fda7bbe 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt
@@ -5,12 +5,16 @@ import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
 import androidx.fragment.app.FragmentManager
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.flowWithLifecycle
+import androidx.lifecycle.lifecycleScope
 import java.text.DateFormat
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.widget.Button
 import net.mullvad.mullvadvpn.ui.widget.CopyableInformationView
 import net.mullvad.mullvadvpn.ui.widget.InformationView
@@ -19,8 +23,11 @@ import net.mullvad.mullvadvpn.ui.widget.SitePaymentButton
 import net.mullvad.mullvadvpn.util.capitalizeFirstCharOfEachWord
 import net.mullvad.talpid.tunnel.ErrorStateCause
 import org.joda.time.DateTime
+import org.koin.android.ext.android.inject
 
 class AccountFragment : ServiceDependentFragment(OnNoService.GoBack) {
+    private val deviceRepository: DeviceRepository by inject()
+
     override val isSecureScreen = true
 
     private val dateStyle = DateFormat.MEDIUM
@@ -97,27 +104,20 @@ class AccountFragment : ServiceDependentFragment(OnNoService.GoBack) {
         return view
     }
 
-    override fun onSafelyStart() {
-        jobTracker.newUiJob("updateAccountNumber") {
-            deviceRepository.deviceState
-                .onEach { state ->
-                    if (state.isInitialState()) deviceRepository.refreshDeviceState()
-                }
-                .collect { state ->
-                    accountNumberView.information = state.token()
-                }
-        }
+    override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
+        super.onViewCreated(view, savedInstanceState)
 
-        jobTracker.newUiJob("updateDeviceName") {
+        lifecycleScope.launch {
             deviceRepository.deviceState
-                .onEach { state ->
-                    if (state.isInitialState()) deviceRepository.refreshDeviceState()
-                }
+                .flowWithLifecycle(lifecycle, Lifecycle.State.RESUMED)
                 .collect { state ->
+                    accountNumberView.information = state.token()
                     deviceNameView.information = state.deviceName()?.capitalizeFirstCharOfEachWord()
                 }
         }
+    }
 
+    override fun onSafelyStart() {
         jobTracker.newUiJob("updateAccountExpiry") {
             accountCache.accountExpiryState
                 .map { state -> state.date() }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LaunchFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LaunchFragment.kt
index 95d72fea9c..420581785b 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LaunchFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LaunchFragment.kt
@@ -4,14 +4,19 @@ import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
-import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.flow.onEach
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.flowWithLifecycle
+import androidx.lifecycle.lifecycleScope
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.DeviceState
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
+import org.koin.android.ext.android.inject
 
 class LaunchFragment : ServiceAwareFragment() {
+    private val deviceRepository: DeviceRepository by inject()
 
     override fun onCreateView(
         inflater: LayoutInflater,
@@ -27,23 +32,16 @@ class LaunchFragment : ServiceAwareFragment() {
         return view
     }
 
-    override fun onStop() {
-        jobTracker.cancelJob("advanceToNextScreen")
-        super.onStop()
-    }
-
     override fun onNewServiceConnection(serviceConnectionContainer: ServiceConnectionContainer) {
-        advanceToNextScreen(serviceConnectionContainer.deviceRepository)
     }
 
-    private fun advanceToNextScreen(deviceRepository: DeviceRepository) {
-        jobTracker.newUiJob("advanceToNextScreen") {
+    override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
+        super.onViewCreated(view, savedInstanceState)
+
+        lifecycleScope.launch {
             deviceRepository.deviceState
-                .onEach { state ->
-                    if (state.isInitialState()) deviceRepository.refreshDeviceState()
-                }
-                .first { state -> state.isInitialState().not() }
-                .let { deviceState ->
+                .flowWithLifecycle(lifecycle, Lifecycle.State.RESUMED)
+                .collect { deviceState ->
                     when (deviceState) {
                         is DeviceState.LoggedIn -> advanceToConnectScreen()
                         is DeviceState.LoggedOut -> advanceToLoginScreen()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
index 5202960e81..2911dcfb9f 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
@@ -10,10 +10,10 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.AppVersionInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.AuthTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
 import net.mullvad.mullvadvpn.ui.serviceconnection.CustomDns
-import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.RelayListListener
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionDeviceDataSource
 import net.mullvad.mullvadvpn.ui.serviceconnection.SettingsListener
 import net.mullvad.mullvadvpn.ui.serviceconnection.SplitTunneling
 
@@ -48,7 +48,7 @@ abstract class ServiceDependentFragment(private val onNoService: OnNoService) :
     lateinit var customDns: CustomDns
         private set
 
-    lateinit var deviceRepository: DeviceRepository
+    lateinit var deviceDataSource: ServiceConnectionDeviceDataSource
         private set
 
     lateinit var locationInfoCache: LocationInfoCache
@@ -70,11 +70,12 @@ abstract class ServiceDependentFragment(private val onNoService: OnNoService) :
         appVersionInfoCache = serviceConnectionContainer.appVersionInfoCache
         authTokenCache = serviceConnectionContainer.authTokenCache
         connectionProxy = serviceConnectionContainer.connectionProxy
-        deviceRepository = serviceConnectionContainer.deviceRepository
+        deviceDataSource = serviceConnectionContainer.deviceDataSource
         customDns = serviceConnectionContainer.customDns
         locationInfoCache = serviceConnectionContainer.locationInfoCache
         relayListListener = serviceConnectionContainer.relayListListener
         settingsListener = serviceConnectionContainer.settingsListener
+
         splitTunneling = serviceConnectionContainer.splitTunneling
 
         synchronized(this) {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt
index cc0de9d2d7..2cf83fc4c7 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt
@@ -6,10 +6,12 @@ import android.view.View
 import android.view.ViewGroup
 import android.widget.ImageButton
 import androidx.core.content.ContextCompat
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.flowWithLifecycle
 import androidx.lifecycle.lifecycleScope
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.DeviceState
 import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
@@ -19,8 +21,11 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.widget.AccountCell
 import net.mullvad.mullvadvpn.ui.widget.AppVersionCell
 import net.mullvad.mullvadvpn.ui.widget.NavigateCell
+import org.koin.android.ext.android.inject
 
 class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBarPainter {
+    private val deviceRepository: DeviceRepository by inject()
+
     private lateinit var accountMenu: AccountCell
     private lateinit var appVersionMenu: AppVersionCell
     private lateinit var preferencesMenu: View
@@ -30,12 +35,10 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
     private var active = false
 
     private var accountCache: AccountCache? = null
-    private var deviceRepository: DeviceRepository? = null
     private var versionInfoCache: AppVersionInfoCache? = null
 
     override fun onNewServiceConnection(serviceConnectionContainer: ServiceConnectionContainer) {
         accountCache = serviceConnectionContainer.accountCache
-        deviceRepository = serviceConnectionContainer.deviceRepository
         versionInfoCache = serviceConnectionContainer.appVersionInfoCache
 
         if (active) {
@@ -45,7 +48,6 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
 
     override fun onNoServiceConnection() {
         accountCache = null
-        deviceRepository = null
         versionInfoCache = null
     }
 
@@ -90,6 +92,14 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
                 paintStatusBar(ContextCompat.getColor(requireContext(), R.color.darkBlue))
             }
         }
+
+        lifecycleScope.launch {
+            deviceRepository.deviceState
+                .flowWithLifecycle(lifecycle, Lifecycle.State.RESUMED)
+                .collect { device ->
+                    updateLoggedInStatus(device is DeviceState.LoggedIn)
+                }
+        }
     }
 
     override fun onResume() {
@@ -131,16 +141,6 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
             fetchAccountExpiry()
         }
 
-        jobTracker.newUiJob("updateLoggedInStatus") {
-            deviceRepository?.let { repository ->
-                repository.deviceState
-                    .onEach { state -> if (state.isInitialState()) repository.refreshDeviceState() }
-                    .collect { device ->
-                        updateLoggedInStatus(device is DeviceState.LoggedIn)
-                    }
-            }
-        }
-
         versionInfoCache?.onUpdate = {
             jobTracker.newUiJob("updateVersionInfo") {
                 updateVersionInfo()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt
index 92364fdf0f..f83660bd06 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt
@@ -13,14 +13,18 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.collect
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.widget.HeaderBar
 import net.mullvad.mullvadvpn.ui.widget.RedeemVoucherButton
 import net.mullvad.mullvadvpn.ui.widget.SitePaymentButton
 import org.joda.time.DateTime
+import org.koin.android.ext.android.inject
 
 val POLL_INTERVAL: Long = 15 /* s */ * 1000 /* ms */
 
 class WelcomeFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
+    private val deviceRepository: DeviceRepository by inject()
+
     private lateinit var accountLabel: TextView
     private lateinit var sitePaymentButton: SitePaymentButton
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
index 7d1eed2961..f2a6fe7ff1 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
@@ -1,20 +1,35 @@
 package net.mullvad.mullvadvpn.ui.serviceconnection
 
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.SharingStarted.Companion.Eagerly
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.stateIn
 import net.mullvad.mullvadvpn.model.DeviceState
 
 class DeviceRepository(
-    private val dataSource: ServiceConnectionDeviceDataSource,
-    externalScope: CoroutineScope
+    private val serviceConnectionManager: ServiceConnectionManager
 ) {
-    val deviceState = dataSource.deviceStateUpdates
-        .stateIn(
-            externalScope,
-            Eagerly,
-            DeviceState.InitialState
-        )
+    val deviceState = serviceConnectionManager.connectionState
+        .flatMapLatest { state ->
+            if (state is ServiceConnectionState.ConnectedReady) {
+                state.container.deviceDataSource.deviceStateUpdates
+                    .onStart {
+                        state.container.deviceDataSource.refreshDevice()
+                    }
+            } else {
+                emptyFlow()
+            }
+        }
+        .stateIn(CoroutineScope(Dispatchers.IO), SharingStarted.Lazily, DeviceState.InitialState)
 
-    fun refreshDeviceState() = dataSource.refreshDevice()
+    fun refreshDeviceState() {
+        container()?.deviceDataSource?.refreshDevice()
+    }
+
+    private fun container(): ServiceConnectionContainer? {
+        return serviceConnectionManager.connectionState.value.readyContainer()
+    }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
index d0d13ac330..e430c2d81f 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
@@ -4,7 +4,6 @@ import android.os.Looper
 import android.os.Messenger
 import android.os.RemoteException
 import android.util.Log
-import kotlinx.coroutines.MainScope
 import net.mullvad.mullvadvpn.di.SERVICE_CONNECTION_SCOPE
 import net.mullvad.mullvadvpn.ipc.DispatchingHandler
 import net.mullvad.mullvadvpn.ipc.Event
@@ -37,8 +36,7 @@ class ServiceConnectionContainer(
     val accountCache = AccountCache(connection, dispatcher)
     val authTokenCache = AuthTokenCache(connection, dispatcher)
     val connectionProxy = ConnectionProxy(connection, dispatcher)
-    val deviceRepository =
-        DeviceRepository(ServiceConnectionDeviceDataSource(connection, dispatcher), MainScope())
+    val deviceDataSource = ServiceConnectionDeviceDataSource(connection, dispatcher)
     val locationInfoCache = LocationInfoCache(dispatcher)
     val settingsListener = SettingsListener(connection, dispatcher)
     // NOTE: `org.koin.core.scope.get` must be used here rather than `org.koin.core.component.get`.
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionState.kt
index 8e48f3a85c..ca868e5cfa 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionState.kt
@@ -7,4 +7,8 @@ sealed class ServiceConnectionState {
         ServiceConnectionState()
 
     object Disconnected : ServiceConnectionState()
+
+    fun readyContainer(): ServiceConnectionContainer? {
+        return (this as? ConnectedReady)?.container
+    }
 }