Merge branch 'split-key-status-listener'

16 files changed, 202 insertions, 113 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt
index ba94206a73..23bc01d5a4 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt
@@ -1,24 +1,27 @@
 package net.mullvad.mullvadvpn.ipc
 
 import android.os.Message as RawMessage
-import android.os.Parcelable
 import kotlinx.parcelize.Parcelize
 import net.mullvad.mullvadvpn.model.GeoIpLocation
+import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.Settings
 
 // Events that can be sent from the service
-sealed class Event : Message(), Parcelable {
+sealed class Event : Message() {
     protected override val messageId = 1
     protected override val messageKey = MESSAGE_KEY
 
     @Parcelize
-    object ListenerReady : Event(), Parcelable
+    object ListenerReady : Event()
 
     @Parcelize
-    data class NewLocation(val location: GeoIpLocation?) : Event(), Parcelable
+    data class NewLocation(val location: GeoIpLocation?) : Event()
 
     @Parcelize
-    data class SettingsUpdate(val settings: Settings?) : Event(), Parcelable
+    data class SettingsUpdate(val settings: Settings?) : Event()
+
+    @Parcelize
+    data class WireGuardKeyStatus(val keyStatus: KeygenEvent?) : Event()
 
     companion object {
         private const val MESSAGE_KEY = "event"
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Request.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Request.kt
index 388f60b07e..d85090ac05 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Request.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Request.kt
@@ -2,16 +2,21 @@ package net.mullvad.mullvadvpn.ipc
 
 import android.os.Message as RawMessage
 import android.os.Messenger
-import android.os.Parcelable
 import kotlinx.parcelize.Parcelize
 
 // Requests that the service can handle
-sealed class Request : Message(), Parcelable {
+sealed class Request : Message() {
     protected override val messageId = 2
     protected override val messageKey = MESSAGE_KEY
 
     @Parcelize
-    data class RegisterListener(val listener: Messenger) : Request(), Parcelable
+    data class RegisterListener(val listener: Messenger) : Request()
+
+    @Parcelize
+    object WireGuardGenerateKey : Request()
+
+    @Parcelize
+    object WireGuardVerifyKey : Request()
 
     companion object {
         private const val MESSAGE_KEY = "request"
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/Constraint.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/Constraint.kt
index 55779fd720..37b98298ab 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/Constraint.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/Constraint.kt
@@ -6,8 +6,8 @@ import kotlinx.parcelize.Parcelize
 sealed class Constraint<T>() : Parcelable {
     @Parcelize
     @Suppress("PARCELABLE_PRIMARY_CONSTRUCTOR_IS_EMPTY")
-    class Any<T>() : Constraint<T>(), Parcelable
+    class Any<T>() : Constraint<T>()
 
     @Parcelize
-    data class Only<T : Parcelable>(val value: T) : Constraint<T>(), Parcelable
+    data class Only<T : Parcelable>(val value: T) : Constraint<T>()
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
index 8c476ef5ff..ced83db74a 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
@@ -1,23 +1,22 @@
 package net.mullvad.mullvadvpn.model
 
-sealed class KeygenEvent {
-    class NewKey(val publicKey: PublicKey) : KeygenEvent() {
-        var verified: Boolean? = null
-            private set
-        var replacementFailure: KeygenFailure? = null
-            private set
+import android.os.Parcelable
+import kotlinx.parcelize.Parcelize
 
-        constructor(
-            publicKey: PublicKey,
-            verified: Boolean?,
-            replacementFailure: KeygenFailure?
-        ) : this(publicKey) {
-            this.verified = verified
-            this.replacementFailure = replacementFailure
-        }
+sealed class KeygenEvent : Parcelable {
+    @Parcelize
+    class NewKey(
+        val publicKey: PublicKey,
+        val verified: Boolean?,
+        val replacementFailure: KeygenFailure?
+    ) : KeygenEvent() {
+        constructor(publicKey: PublicKey) : this (publicKey, null, null)
     }
 
+    @Parcelize
     object TooManyKeys : KeygenEvent()
+
+    @Parcelize
     object GenerationFailure : KeygenEvent()
 
     fun failure(): KeygenFailure? {
@@ -29,7 +28,8 @@ sealed class KeygenEvent {
     }
 }
 
-enum class KeygenFailure {
+@Parcelize
+enum class KeygenFailure : Parcelable {
     TooManyKeys,
     GenerationFailure,
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/LocationConstraint.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/LocationConstraint.kt
index 55757e5b8f..6734ff418e 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/LocationConstraint.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/LocationConstraint.kt
@@ -7,16 +7,13 @@ sealed class LocationConstraint : Parcelable {
     abstract val location: GeoIpLocation
 
     @Parcelize
-    data class Country(val countryCode: String) : LocationConstraint(), Parcelable {
+    data class Country(val countryCode: String) : LocationConstraint() {
         override val location: GeoIpLocation
             get() = GeoIpLocation(null, null, countryCode, null, null)
     }
 
     @Parcelize
-    data class City(
-        val countryCode: String,
-        val cityCode: String
-    ) : LocationConstraint(), Parcelable {
+    data class City(val countryCode: String, val cityCode: String) : LocationConstraint() {
         override val location: GeoIpLocation
             get() = GeoIpLocation(null, null, countryCode, cityCode, null)
     }
@@ -26,7 +23,7 @@ sealed class LocationConstraint : Parcelable {
         val countryCode: String,
         val cityCode: String,
         val hostname: String
-    ) : LocationConstraint(), Parcelable {
+    ) : LocationConstraint() {
         override val location: GeoIpLocation
             get() = GeoIpLocation(null, null, countryCode, cityCode, hostname)
     }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/PublicKey.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/PublicKey.kt
index 1931f40ec7..4ee6ad51df 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/PublicKey.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/PublicKey.kt
@@ -1,3 +1,7 @@
 package net.mullvad.mullvadvpn.model
 
-data class PublicKey(val key: ByteArray, val dateCreated: String)
+import android.os.Parcelable
+import kotlinx.parcelize.Parcelize
+
+@Parcelize
+data class PublicKey(val key: ByteArray, val dateCreated: String) : Parcelable
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/RelaySettings.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/RelaySettings.kt
index e1e56c3c51..6a247997db 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/RelaySettings.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/RelaySettings.kt
@@ -5,8 +5,8 @@ import kotlinx.parcelize.Parcelize
 
 sealed class RelaySettings : Parcelable {
     @Parcelize
-    object CustomTunnelEndpoint : RelaySettings(), Parcelable
+    object CustomTunnelEndpoint : RelaySettings()
 
     @Parcelize
-    class Normal(val relayConstraints: RelayConstraints) : RelaySettings(), Parcelable
+    class Normal(val relayConstraints: RelayConstraints) : RelaySettings()
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/KeyStatusListener.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/KeyStatusListener.kt
deleted file mode 100644
index 2b8c2245dd..0000000000
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/KeyStatusListener.kt
+++ /dev/null
@@ -1,57 +0,0 @@
-package net.mullvad.mullvadvpn.service
-
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.GlobalScope
-import kotlinx.coroutines.launch
-import net.mullvad.mullvadvpn.model.KeygenEvent
-import net.mullvad.talpid.util.EventNotifier
-
-class KeyStatusListener(val daemon: MullvadDaemon) {
-    val onKeyStatusChange = EventNotifier(getInitialKeyStatus())
-
-    var keyStatus by onKeyStatusChange.notifiable()
-
-    init {
-        daemon.onKeygenEvent = { event -> keyStatus = event }
-    }
-
-    private fun getInitialKeyStatus(): KeygenEvent? {
-        return daemon.getWireguardKey()?.let { wireguardKey ->
-            KeygenEvent.NewKey(wireguardKey, null, null)
-        }
-    }
-
-    fun generateKey() = GlobalScope.launch(Dispatchers.Default) {
-        val oldStatus = keyStatus
-        val newStatus = daemon.generateWireguardKey()
-        val newFailure = newStatus?.failure()
-        if (oldStatus is KeygenEvent.NewKey && newFailure != null) {
-            keyStatus = KeygenEvent.NewKey(
-                oldStatus.publicKey,
-                oldStatus.verified,
-                newFailure
-            )
-        } else {
-            keyStatus = newStatus ?: KeygenEvent.GenerationFailure
-        }
-    }
-
-    fun verifyKey() = GlobalScope.launch(Dispatchers.Default) {
-        val verified = daemon.verifyWireguardKey()
-        // Only update verification status if the key is actually there
-        when (val state = keyStatus) {
-            is KeygenEvent.NewKey -> {
-                keyStatus = KeygenEvent.NewKey(
-                    state.publicKey,
-                    verified,
-                    state.replacementFailure
-                )
-            }
-        }
-    }
-
-    fun onDestroy() {
-        daemon.onKeygenEvent = null
-        onKeyStatusChange.unsubscribeAll()
-    }
-}
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt
index b09a80696e..251802dce9 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt
@@ -14,12 +14,10 @@ class ServiceInstance(
     val splitTunneling: SplitTunneling
 ) {
     val accountCache = AccountCache(daemon, settingsListener)
-    val keyStatusListener = KeyStatusListener(daemon)
 
     fun onDestroy() {
         accountCache.onDestroy()
         connectionProxy.onDestroy()
         customDns.onDestroy()
-        keyStatusListener.onDestroy()
     }
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/KeyStatusListener.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/KeyStatusListener.kt
new file mode 100644
index 0000000000..70ac7ef827
--- /dev/null
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/KeyStatusListener.kt
@@ -0,0 +1,95 @@
+package net.mullvad.mullvadvpn.service.endpoint
+
+import kotlin.properties.Delegates.observable
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.ClosedReceiveChannelException
+import kotlinx.coroutines.channels.actor
+import kotlinx.coroutines.channels.sendBlocking
+import net.mullvad.mullvadvpn.ipc.Event
+import net.mullvad.mullvadvpn.ipc.Request
+import net.mullvad.mullvadvpn.model.KeygenEvent
+
+class KeyStatusListener(endpoint: ServiceEndpoint) {
+    companion object {
+        private enum class Command {
+            GenerateKey,
+            VerifyKey,
+        }
+    }
+
+    private val daemon = endpoint.intermittentDaemon
+
+    private val commandChannel = spawnActor()
+
+    var keyStatus by observable<KeygenEvent?>(null) { _, _, status ->
+        endpoint.sendEvent(Event.WireGuardKeyStatus(status))
+    }
+        private set
+
+    init {
+        daemon.registerListener(this) { newDaemon ->
+            newDaemon?.apply {
+                keyStatus = getWireguardKey()?.let { wireguardKey ->
+                    KeygenEvent.NewKey(wireguardKey, null, null)
+                }
+
+                onKeygenEvent = { event -> keyStatus = event }
+            }
+        }
+
+        endpoint.dispatcher.apply {
+            registerHandler(Request.WireGuardGenerateKey::class) { _ ->
+                commandChannel.sendBlocking(Command.GenerateKey)
+            }
+
+            registerHandler(Request.WireGuardVerifyKey::class) { _ ->
+                commandChannel.sendBlocking(Command.VerifyKey)
+            }
+        }
+    }
+
+    fun onDestroy() {
+        commandChannel.close()
+        daemon.unregisterListener(this)
+    }
+
+    private fun spawnActor() = GlobalScope.actor<Command>(Dispatchers.Default, Channel.UNLIMITED) {
+        try {
+            for (command in channel) {
+                when (command) {
+                    Command.GenerateKey -> generateKey()
+                    Command.VerifyKey -> verifyKey()
+                }
+            }
+        } catch (exception: ClosedReceiveChannelException) {
+        }
+    }
+
+    private suspend fun generateKey() {
+        val oldStatus = keyStatus
+        val newStatus = daemon.await().generateWireguardKey()
+        val newFailure = newStatus?.failure()
+        if (oldStatus is KeygenEvent.NewKey && newFailure != null) {
+            keyStatus = KeygenEvent.NewKey(
+                oldStatus.publicKey,
+                oldStatus.verified,
+                newFailure
+            )
+        } else {
+            keyStatus = newStatus ?: KeygenEvent.GenerationFailure
+        }
+    }
+
+    private suspend fun verifyKey() {
+        // Only update verification status if the key is actually there
+        (keyStatus as? KeygenEvent.NewKey)?.let { currentStatus ->
+            keyStatus = KeygenEvent.NewKey(
+                currentStatus.publicKey,
+                daemon.await().verifyWireguardKey(),
+                currentStatus.replacementFailure
+            )
+        }
+    }
+}
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/ServiceEndpoint.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/ServiceEndpoint.kt
index 5f0ecc5ecd..05f903e66a 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/ServiceEndpoint.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/ServiceEndpoint.kt
@@ -33,6 +33,7 @@ class ServiceEndpoint(
 
     val settingsListener = SettingsListener(this)
 
+    val keyStatusListener = KeyStatusListener(this)
     val locationInfoCache = LocationInfoCache(this)
 
     init {
@@ -45,6 +46,7 @@ class ServiceEndpoint(
         dispatcher.onDestroy()
         registrationQueue.close()
 
+        keyStatusListener.onDestroy()
         locationInfoCache.onDestroy()
         settingsListener.onDestroy()
     }
@@ -89,6 +91,7 @@ class ServiceEndpoint(
             val initialEvents = listOf(
                 Event.SettingsUpdate(settingsListener.settings),
                 Event.NewLocation(locationInfoCache.location),
+                Event.WireGuardKeyStatus(keyStatusListener.keyStatus),
                 Event.ListenerReady
             )
 
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
index 5e5ecba072..1e892f9160 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
@@ -10,9 +10,9 @@ import net.mullvad.mullvadvpn.dataproxy.RelayListListener
 import net.mullvad.mullvadvpn.service.AccountCache
 import net.mullvad.mullvadvpn.service.ConnectionProxy
 import net.mullvad.mullvadvpn.service.CustomDns
-import net.mullvad.mullvadvpn.service.KeyStatusListener
 import net.mullvad.mullvadvpn.service.MullvadDaemon
 import net.mullvad.mullvadvpn.service.SplitTunneling
+import net.mullvad.mullvadvpn.ui.serviceconnection.KeyStatusListener
 import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnection
 import net.mullvad.mullvadvpn.ui.serviceconnection.SettingsListener
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/WireguardKeyFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/WireguardKeyFragment.kt
index 973d57a298..b5494cd82a 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/WireguardKeyFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/WireguardKeyFragment.kt
@@ -7,6 +7,7 @@ import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
 import android.widget.TextView
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.delay
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.KeygenEvent
@@ -40,6 +41,7 @@ class WireguardKeyFragment : ServiceDependentFragment(OnNoService.GoToLaunchScre
     private var greenColor: Int = 0
     private var redColor: Int = 0
 
+    private var actionCompletion: CompletableDeferred<Unit>? = null
     private var tunnelState: TunnelState = TunnelState.Disconnected
 
     private var actionState: ActionState = ActionState.Idle(false)
@@ -63,6 +65,8 @@ class WireguardKeyFragment : ServiceDependentFragment(OnNoService.GoToLaunchScre
                 updateStatusMessage()
                 updateGenerateKeyButtonText()
                 updateVerifyKeyButtonState()
+
+                actionCompletion?.complete(Unit)
             }
         }
 
@@ -145,18 +149,16 @@ class WireguardKeyFragment : ServiceDependentFragment(OnNoService.GoToLaunchScre
     override fun onSafelyStart() {
         connectionProxy.onUiStateChange.subscribe(this) { uiState ->
             jobTracker.newUiJob("tunnelStateUpdate") {
-                synchronized(this@WireguardKeyFragment) {
-                    tunnelState = uiState
-
-                    if (actionState is ActionState.Generating) {
-                        reconnectionExpected = !(tunnelState is TunnelState.Disconnected)
-                    } else if (tunnelState is TunnelState.Connected) {
-                        reconnectionExpected = false
-                    }
+                tunnelState = uiState
 
-                    isOffline = uiState is TunnelState.Error &&
-                        uiState.errorState.cause is ErrorStateCause.IsOffline
+                if (actionState is ActionState.Generating) {
+                    reconnectionExpected = !(tunnelState is TunnelState.Disconnected)
+                } else if (tunnelState is TunnelState.Connected) {
+                    reconnectionExpected = false
                 }
+
+                isOffline = uiState is TunnelState.Error &&
+                    uiState.errorState.cause is ErrorStateCause.IsOffline
             }
         }
 
@@ -312,25 +314,30 @@ class WireguardKeyFragment : ServiceDependentFragment(OnNoService.GoToLaunchScre
     }
 
     private suspend fun onGenerateKeyPress() {
-        synchronized(this) {
-            actionState = ActionState.Generating(keyStatus is KeygenEvent.NewKey)
-            reconnectionExpected = !(tunnelState is TunnelState.Disconnected)
-        }
+        actionState = ActionState.Generating(keyStatus is KeygenEvent.NewKey)
+        reconnectionExpected = !(tunnelState is TunnelState.Disconnected)
 
         keyStatus = null
-        keyStatusListener.generateKey().join()
+
+        actionCompletion = CompletableDeferred()
+        keyStatusListener.generateKey()
+        actionCompletion?.await()
 
         actionState = ActionState.Idle(false)
     }
 
     private suspend fun onValidateKeyPress() {
         actionState = ActionState.Verifying()
-        keyStatusListener.verifyKey().join()
+
+        actionCompletion = CompletableDeferred()
+        keyStatusListener.verifyKey()
+        actionCompletion?.await()
+
         actionState = ActionState.Idle(true)
     }
 
     private fun resetReconnectionExpected() {
-        jobTracker.newBackgroundJob("resetReconnectionExpected") {
+        jobTracker.newUiJob("resetReconnectionExpected") {
             delay(20_000)
 
             if (reconnectionExpected) {
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/notification/KeyStatusNotification.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/notification/KeyStatusNotification.kt
index 880cea9f8c..fba57eca03 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/notification/KeyStatusNotification.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/notification/KeyStatusNotification.kt
@@ -3,8 +3,8 @@ package net.mullvad.mullvadvpn.ui.notification
 import android.content.Context
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.KeygenEvent
-import net.mullvad.mullvadvpn.service.KeyStatusListener
 import net.mullvad.mullvadvpn.service.MullvadDaemon
+import net.mullvad.mullvadvpn.ui.serviceconnection.KeyStatusListener
 
 class KeyStatusNotification(
     context: Context,
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/KeyStatusListener.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/KeyStatusListener.kt
new file mode 100644
index 0000000000..7f7425e9c7
--- /dev/null
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/KeyStatusListener.kt
@@ -0,0 +1,33 @@
+package net.mullvad.mullvadvpn.ui.serviceconnection
+
+import android.os.Messenger
+import net.mullvad.mullvadvpn.ipc.DispatchingHandler
+import net.mullvad.mullvadvpn.ipc.Event
+import net.mullvad.mullvadvpn.ipc.Request
+import net.mullvad.mullvadvpn.model.KeygenEvent
+import net.mullvad.talpid.util.EventNotifier
+
+class KeyStatusListener(val connection: Messenger, val eventDispatcher: DispatchingHandler<Event>) {
+    val onKeyStatusChange = EventNotifier<KeygenEvent?>(null)
+
+    var keyStatus by onKeyStatusChange.notifiable()
+        private set
+
+    init {
+        eventDispatcher.registerHandler(Event.WireGuardKeyStatus::class) { event ->
+            keyStatus = event.keyStatus
+        }
+    }
+
+    fun generateKey() {
+        connection.send(Request.WireGuardGenerateKey.message)
+    }
+
+    fun verifyKey() {
+        connection.send(Request.WireGuardVerifyKey.message)
+    }
+
+    fun onDestroy() {
+        onKeyStatusChange.unsubscribeAll()
+    }
+}
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnection.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnection.kt
index da0e09b703..a9ac1da761 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnection.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnection.kt
@@ -25,7 +25,7 @@ class ServiceConnection(private val service: ServiceInstance, val mainActivity:
     val accountCache = service.accountCache
     val connectionProxy = service.connectionProxy
     val customDns = service.customDns
-    val keyStatusListener = service.keyStatusListener
+    val keyStatusListener = KeyStatusListener(service.messenger, dispatcher)
     val locationInfoCache = LocationInfoCache(dispatcher)
     val settingsListener = SettingsListener(dispatcher)
     val splitTunneling = service.splitTunneling
@@ -42,6 +42,7 @@ class ServiceConnection(private val service: ServiceInstance, val mainActivity:
     fun onDestroy() {
         dispatcher.onDestroy()
 
+        keyStatusListener.onDestroy()
         locationInfoCache.onDestroy()
         settingsListener.onDestroy()