Fix key generation type mismatch

4 files changed, 23 insertions, 18 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
index f6fd4c9e1e..495a62c387 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
@@ -14,7 +14,6 @@ import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.dataproxy.AppVersionInfoCache
 import net.mullvad.mullvadvpn.dataproxy.WwwAuthTokenRetriever
 import net.mullvad.mullvadvpn.model.KeygenEvent
-import net.mullvad.mullvadvpn.model.KeygenFailure
 import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.talpid.tunnel.ActionAfterDisconnect
 import net.mullvad.talpid.tunnel.BlockReason
@@ -102,16 +101,12 @@ class NotificationBanner(
         when (keyState) {
             null -> return false
             is KeygenEvent.NewKey -> return false
-            is KeygenEvent.Failure -> {
-                when (keyState.failure) {
-                    is KeygenFailure.TooManyKeys -> {
+            is KeygenEvent.TooManyKeys -> {
                         externalLink = ExternalLink.KeyManagement
                         showError(R.string.wireguard_error, R.string.too_many_keys)
-                    }
-                    is KeygenFailure.GenerationFailure -> {
+            }
+            is KeygenEvent.GenerationFailure -> {
                         showError(R.string.wireguard_error, R.string.failed_to_generate_key)
-                    }
-                }
             }
         }
 
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/WireguardKeyFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/WireguardKeyFragment.kt
index 141e373ccf..11f27d2b53 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/WireguardKeyFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/WireguardKeyFragment.kt
@@ -162,8 +162,8 @@ class WireguardKeyFragment : Fragment() {
 
                 keyState.replacementFailure?.let { error -> showKeygenFailure(error) }
             }
-            is KeygenEvent.Failure -> {
-                showKeygenFailure(keyState.failure)
+            else -> {
+                showKeygenFailure(keyState.failure())
             }
         }
         drawNoConnectionState()
@@ -179,7 +179,7 @@ class WireguardKeyFragment : Fragment() {
         statusMessage.visibility = View.GONE
     }
 
-    private fun showKeygenFailure(failure: KeygenFailure) {
+    private fun showKeygenFailure(failure: KeygenFailure?) {
         when (failure) {
             is KeygenFailure.TooManyKeys -> {
                 setStatusMessage(R.string.too_many_keys, R.color.red)
@@ -220,7 +220,7 @@ class WireguardKeyFragment : Fragment() {
         verifyButton.setClickable(true)
         verifyButton.setAlpha(1f)
         val keyState = keyStatusListener.keyStatus
-        if (generatingKey || keyState is KeygenEvent.Failure) {
+        if (generatingKey || keyState?.failure() != null) {
             verifyButton.setClickable(false)
             verifyButton.setAlpha(0.5f)
             return
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
index f36495734e..ce7c5ea543 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
@@ -45,10 +45,11 @@ class KeyStatusListener(val asyncDaemon: Deferred<MullvadDaemon>) {
             setUpJob.join()
             val oldStatus = keyStatus
             val newStatus = daemon?.generateWireguardKey()
-            if (oldStatus is KeygenEvent.NewKey && newStatus is KeygenEvent.Failure) {
+            val newFailure = newStatus?.failure()
+            if (oldStatus is KeygenEvent.NewKey && newStatus != null) {
                 keyStatus = KeygenEvent.NewKey(oldStatus.publicKey,
                                 oldStatus.verified,
-                                newStatus.failure)
+                                newFailure)
             } else {
                 keyStatus = newStatus
             }
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 0e1e801895..00764ce9dc 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
@@ -16,10 +16,19 @@ sealed class KeygenEvent {
             this.replacementFailure = replacementFailure
         }
     }
-    class Failure(val failure: KeygenFailure) : KeygenEvent()
+    class TooManyKeys : KeygenEvent()
+    class GenerationFailure : KeygenEvent()
+
+    fun failure(): KeygenFailure? {
+        return when (this) {
+            is KeygenEvent.TooManyKeys -> KeygenFailure.TooManyKeys()
+            is KeygenEvent.GenerationFailure -> KeygenFailure.GenerationFailure()
+            else -> { null }
+        }
+    }
 }
 
-sealed class KeygenFailure {
-    class TooManyKeys : KeygenFailure()
-    class GenerationFailure : KeygenFailure()
+sealed class KeygenFailure() {
+    class TooManyKeys() : KeygenFailure()
+    class GenerationFailure() : KeygenFailure()
 }