Propagate device list

4 files changed, 115 insertions, 6 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
index 380ae0dedf..8d983ad883 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
@@ -5,6 +5,7 @@ import kotlinx.coroutines.flow.asSharedFlow
 import net.mullvad.mullvadvpn.model.AppVersionInfo
 import net.mullvad.mullvadvpn.model.Device
 import net.mullvad.mullvadvpn.model.DeviceEvent
+import net.mullvad.mullvadvpn.model.DeviceListEvent
 import net.mullvad.mullvadvpn.model.DeviceState
 import net.mullvad.mullvadvpn.model.DnsOptions
 import net.mullvad.mullvadvpn.model.GeoIpLocation
@@ -22,7 +23,6 @@ import net.mullvad.talpid.util.EventNotifier
 class MullvadDaemon(vpnService: MullvadVpnService) {
     protected var daemonInterfaceAddress = 0L
 
-    var onDeviceRemoved = EventNotifier<RemoveDeviceEvent?>(null)
     val onSettingsChange = EventNotifier<Settings?>(null)
     var onTunnelStateChange = EventNotifier<TunnelState>(TunnelState.Disconnected)
 
@@ -33,6 +33,9 @@ class MullvadDaemon(vpnService: MullvadVpnService) {
     private val _deviceStateUpdates = MutableSharedFlow<DeviceState>(extraBufferCapacity = 1)
     val deviceStateUpdates = _deviceStateUpdates.asSharedFlow()
 
+    private val _deviceListUpdates = MutableSharedFlow<DeviceListEvent>(extraBufferCapacity = 1)
+    val deviceListUpdates = _deviceListUpdates.asSharedFlow()
+
     init {
         System.loadLibrary("mullvad_jni")
         initialize(vpnService, vpnService.cacheDir.absolutePath, vpnService.filesDir.absolutePath)
@@ -104,8 +107,16 @@ class MullvadDaemon(vpnService: MullvadVpnService) {
 
     fun logoutAccount() = logoutAccount(daemonInterfaceAddress)
 
-    fun listDevices(accountToken: String?): List<Device>? {
-        return listDevices(daemonInterfaceAddress, accountToken)
+    fun getAndEmitDeviceList(accountToken: String): List<Device>? {
+        return listDevices(daemonInterfaceAddress, accountToken).also { deviceList ->
+            _deviceListUpdates.tryEmit(
+                if (deviceList == null) {
+                    DeviceListEvent.Error
+                } else {
+                    DeviceListEvent.Available(accountToken, deviceList)
+                }
+            )
+        }
     }
 
     fun getAndEmitDeviceState(): DeviceState {
@@ -247,6 +258,6 @@ class MullvadDaemon(vpnService: MullvadVpnService) {
     }
 
     private fun notifyRemoveDeviceEvent(event: RemoveDeviceEvent) {
-        onDeviceRemoved.notify(event)
+        _deviceListUpdates.tryEmit(DeviceListEvent.Available(event.accountToken, event.newDevices))
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt
index cc23b3fe01..cb290a6d27 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt
@@ -22,17 +22,41 @@ class DaemonDeviceDataSource(
     }
 
     private fun launchDeviceEndpointJobs(daemon: MullvadDaemon) {
-        tracker.newBackgroundJob("propagateDeviceUpdates") {
+        tracker.newBackgroundJob("propagateDeviceUpdatesJob") {
             daemon.deviceStateUpdates.collect { newState ->
                 endpoint.sendEvent(Event.DeviceStateEvent(newState))
             }
         }
 
+        tracker.newBackgroundJob("propagateDeviceListUpdatesJob") {
+            daemon.deviceListUpdates.collect { newState ->
+                endpoint.sendEvent(Event.DeviceListUpdate(newState))
+            }
+        }
+
+        endpoint.dispatcher.registerHandler(Request.GetDevice::class) {
+            tracker.newBackgroundJob("getDeviceJob") {
+                daemon.getAndEmitDeviceState()
+            }
+        }
+
         endpoint.dispatcher.registerHandler(Request.RefreshDeviceState::class) {
             tracker.newBackgroundJob("refreshDeviceJob") {
                 daemon.refreshDevice()
             }
         }
+
+        endpoint.dispatcher.registerHandler(Request.RemoveDevice::class) { request ->
+            tracker.newBackgroundJob("removeDeviceJob") {
+                daemon.removeDevice(request.accountToken, request.deviceId)
+            }
+        }
+
+        endpoint.dispatcher.registerHandler(Request.GetDeviceList::class) { request ->
+            tracker.newBackgroundJob("getDeviceListJob") {
+                daemon.getAndEmitDeviceList(request.accountToken)
+            }
+        }
     }
 
     fun onDestroy() {
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 6e50b401ec..08290ef7d2 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
@@ -3,23 +3,34 @@ package net.mullvad.mullvadvpn.ui.serviceconnection
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.withTimeoutOrNull
+import net.mullvad.mullvadvpn.model.Device
+import net.mullvad.mullvadvpn.model.DeviceListEvent
 import net.mullvad.mullvadvpn.model.DeviceState
 
 class DeviceRepository(
     private val serviceConnectionManager: ServiceConnectionManager,
+    private val deviceListTimeoutMillis: Long = 5000L,
     dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) {
+    private val cachedDeviceList = MutableStateFlow<List<Device>>(emptyList())
+
     val deviceState = serviceConnectionManager.connectionState
         .flatMapLatest { state ->
             if (state is ServiceConnectionState.ConnectedReady) {
                 state.container.deviceDataSource.deviceStateUpdates
                     .onStart {
-                        state.container.deviceDataSource.refreshDevice()
+                        state.container.deviceDataSource.getDevice()
                     }
             } else {
                 flowOf(DeviceState.Unknown)
@@ -31,6 +42,24 @@ class DeviceRepository(
             DeviceState.Initial
         )
 
+    private val deviceListEvents = serviceConnectionManager.connectionState
+        .flatMapLatest { state ->
+            if (state is ServiceConnectionState.ConnectedReady) {
+                state.container.deviceDataSource.deviceListUpdates
+            } else {
+                emptyFlow()
+            }
+        }
+
+    val deviceList = deviceListEvents
+        .map { (it as? DeviceListEvent.Available)?.devices ?: emptyList() }
+        .onStart {
+            if (cachedDeviceList.value.isNotEmpty()) {
+                emit(cachedDeviceList.value)
+            }
+        }
+        .stateIn(CoroutineScope(Dispatchers.IO), SharingStarted.WhileSubscribed(), emptyList())
+
     fun refreshDeviceState() {
         container()?.deviceDataSource?.refreshDevice()
     }
@@ -38,4 +67,27 @@ class DeviceRepository(
     private fun container(): ServiceConnectionContainer? {
         return serviceConnectionManager.connectionState.value.readyContainer()
     }
+
+    fun removeDevice(accountToken: String, deviceId: String) {
+        cachedDeviceList.value = emptyList()
+        container()?.deviceDataSource?.removeDevice(accountToken, deviceId)
+    }
+
+    fun refreshDeviceList(accountToken: String) {
+        container()?.deviceDataSource?.refreshDeviceList(accountToken)
+    }
+
+    suspend fun getDeviceList(accountToken: String): DeviceListEvent {
+        return withTimeoutOrNull(deviceListTimeoutMillis) {
+            deviceListEvents
+                .onStart {
+                    refreshDeviceList(accountToken)
+                }
+                .onEach {
+                    cachedDeviceList.value =
+                        (it as? DeviceListEvent.Available)?.devices ?: emptyList()
+                }
+                .firstOrNull() ?: DeviceListEvent.Error
+        } ?: DeviceListEvent.Error
+    }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt
index 23892257d6..6f018a27b1 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt
@@ -21,8 +21,30 @@ class ServiceConnectionDeviceDataSource(
         }
     }
 
+    val deviceListUpdates = callbackFlow {
+        val handler: (Event.DeviceListUpdate) -> Unit = { event ->
+            trySend(event.event)
+        }
+        dispatcher.registerHandler(Event.DeviceListUpdate::class, handler)
+        awaitClose {
+            // The current dispatcher doesn't support unregistration of handlers.
+        }
+    }
+
     // Async result: Event.DeviceChanged
     fun refreshDevice() {
         connection.send(Request.RefreshDeviceState.message)
     }
+
+    fun getDevice() {
+        connection.send(Request.GetDevice.message)
+    }
+
+    fun removeDevice(accountToken: String, deviceId: String) {
+        connection.send(Request.RemoveDevice(accountToken, deviceId).message)
+    }
+
+    fun refreshDeviceList(accountToken: String) {
+        connection.send(Request.GetDeviceList(accountToken).message)
+    }
 }