Merge branch 'android-allow-key-regeneration'

9 files changed, 174 insertions, 69 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
index acc11784fb..a9053b02c3 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
@@ -12,6 +12,7 @@ import net.mullvad.mullvadvpn.model.ActionAfterDisconnect
 import net.mullvad.mullvadvpn.model.BlockReason
 import net.mullvad.mullvadvpn.model.ParameterGenerationError
 import net.mullvad.mullvadvpn.model.KeygenEvent
+import net.mullvad.mullvadvpn.model.KeygenFailure
 import net.mullvad.mullvadvpn.model.TunnelState
 
 class NotificationBanner(
@@ -61,15 +62,20 @@ class NotificationBanner(
     }
 
     private fun updateBasedOnKeyState(): Boolean {
+        val keyState = keyState
         when (keyState) {
             null -> return false
             is KeygenEvent.NewKey -> return false
-            is KeygenEvent.TooManyKeys -> {
-                externalLink = accountUrl
-                showError(R.string.wireguard_error, R.string.too_many_keys)
-            }
-            is KeygenEvent.GenerationFailure -> {
-                showError(R.string.wireguard_error, R.string.failed_to_generate_key)
+            is KeygenEvent.Failure -> {
+                when (keyState.failure) {
+                    is KeygenFailure.TooManyKeys -> {
+                        externalLink = accountUrl
+                        showError(R.string.wireguard_error, R.string.too_many_keys)
+                    }
+                    is KeygenFailure.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 e9564d38b4..b8ce500d3c 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/WireguardKeyFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/WireguardKeyFragment.kt
@@ -11,6 +11,7 @@ import android.net.Uri
 import android.os.Bundle
 import android.support.v4.app.Fragment
 import android.util.Base64
+import android.util.Log
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
@@ -21,12 +22,12 @@ import android.widget.Toast
 
 import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.dataproxy.KeyStatusListener
+import net.mullvad.mullvadvpn.model.KeygenFailure
 import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.mullvadvpn.util.SmartDeferred
 
 class WireguardKeyFragment : Fragment() {
-    private var keyState: KeygenEvent? = null
     private var currentJob: Job? = null
     private var updateViewsJob: Job? = null
     private var tunnelStateListener: Int? = null
@@ -41,9 +42,10 @@ class WireguardKeyFragment : Fragment() {
     private lateinit var publicKey: TextView
     private lateinit var statusMessage: TextView
     private lateinit var visitWebsiteView: View
-    private lateinit var actionButton: Button
-    private lateinit var actionSpinner: ProgressBar
-
+    private lateinit var generateButton: Button
+    private lateinit var generateSpinner: ProgressBar
+    private lateinit var verifyButton: Button
+    private lateinit var verifySpinner: ProgressBar
 
     override fun onAttach(context: Context) {
         super.onAttach(context)
@@ -53,9 +55,9 @@ class WireguardKeyFragment : Fragment() {
     }
 
     override fun onCreateView(
-            inflater: LayoutInflater,
-            container: ViewGroup?,
-            savedInstanceState: Bundle?
+        inflater: LayoutInflater,
+        container: ViewGroup?,
+        savedInstanceState: Bundle?
     ): View {
         val view = inflater.inflate(R.layout.wireguard_key, container, false)
 
@@ -63,12 +65,13 @@ class WireguardKeyFragment : Fragment() {
             parentActivity.onBackPressed()
         }
 
-
         statusMessage = view.findViewById<TextView>(R.id.wireguard_key_status)
         visitWebsiteView = view.findViewById<View>(R.id.wireguard_manage_keys)
         publicKey = view.findViewById<TextView>(R.id.wireguard_public_key)
-        actionButton = view.findViewById<Button>(R.id.wg_key_button)
-        actionSpinner = view.findViewById<ProgressBar>(R.id.wg_action_spinner)
+        generateButton = view.findViewById<Button>(R.id.wg_generate_key_button)
+        generateSpinner = view.findViewById<ProgressBar>(R.id.wg_generate_key_spinner)
+        verifyButton = view.findViewById<Button>(R.id.wg_verify_key_button)
+        verifySpinner = view.findViewById<ProgressBar>(R.id.wg_verify_key_spinner)
 
         visitWebsiteView.visibility = View.VISIBLE
         visitWebsiteView.setOnClickListener {
@@ -85,39 +88,35 @@ class WireguardKeyFragment : Fragment() {
         updateViews()
     }
 
-
     private fun updateViews() {
         clearErrorMessage()
 
+        setGenerateButton()
+        setVerifyButton()
+
         when (val keyState = keyStatusListener.keyStatus) {
             null -> {
                 publicKey.visibility = View.INVISIBLE
-                setGenerateButton()
             }
-            is KeygenEvent.TooManyKeys -> {
 
-                setStatusMessage(R.string.too_many_keys, R.color.red)
-                setGenerateButton()
-            }
-            is KeygenEvent.GenerationFailure -> {
-                setStatusMessage(R.string.failed_to_generate_key, R.color.red)
-                setGenerateButton()
-            }
             is KeygenEvent.NewKey -> {
                 val publicKeyString = Base64.encodeToString(keyState.publicKey.key, Base64.DEFAULT)
                 publicKey.visibility = View.VISIBLE
                 publicKey.setText(publicKeyString)
 
-                setVerifyButton()
-
                 if (keyState.verified != null) {
                     if (keyState.verified) {
                         setStatusMessage(R.string.wireguard_key_valid, R.color.green)
                     } else {
                         setStatusMessage(R.string.wireguard_key_invalid, R.color.red)
-                        setGenerateButton()
                     }
                 }
+                if (keyState.replacementFailure != null) {
+                    showKeygenFailure(keyState.replacementFailure)
+                }
+            }
+            is KeygenEvent.Failure -> {
+                showKeygenFailure(keyState.failure)
             }
         }
         drawNoConnectionState()
@@ -133,55 +132,85 @@ class WireguardKeyFragment : Fragment() {
         statusMessage.visibility = View.GONE
     }
 
+    private fun showKeygenFailure(failure: KeygenFailure) {
+        when (failure) {
+            is KeygenFailure.TooManyKeys -> {
+                setStatusMessage(R.string.too_many_keys, R.color.red)
+            }
+            is KeygenFailure.GenerationFailure -> {
+                setStatusMessage(R.string.failed_to_generate_key, R.color.red)
+            }
+        }
+    }
+
     private fun setGenerateButton() {
+        generateButton.setClickable(true)
+        generateButton.setAlpha(1f)
+        if (validatingKey) {
+            generateButton.setClickable(false)
+            generateButton.setAlpha(0.5f)
+            return
+        }
         if (generatingKey) {
-            showActionSpinner()
+            generateButton.visibility = View.GONE
+            generateSpinner.visibility = View.VISIBLE
             return
         }
-        actionSpinner.visibility = View.GONE
-        actionButton.visibility = View.VISIBLE
-        actionButton.setText(R.string.wireguard_generate_key)
-        actionButton.setOnClickListener {
+        generateSpinner.visibility = View.GONE
+        generateButton.visibility = View.VISIBLE
+        if (keyStatusListener.keyStatus is KeygenEvent.NewKey) {
+            generateButton.setText(R.string.wireguard_replace_key)
+        } else {
+            generateButton.setText(R.string.wireguard_generate_key)
+        }
+
+        generateButton.setOnClickListener {
             onGenerateKeyPress()
         }
     }
 
     private fun setVerifyButton() {
+        verifyButton.setClickable(true)
+        verifyButton.setAlpha(1f)
+        val keyState = keyStatusListener.keyStatus
+        if (generatingKey || keyState is KeygenEvent.Failure) {
+            verifyButton.setClickable(false)
+            verifyButton.setAlpha(0.5f)
+            return
+        }
         if (validatingKey) {
-            showActionSpinner()
+            verifyButton.visibility = View.GONE
+            verifySpinner.visibility = View.VISIBLE
             return
         }
-        actionSpinner.visibility = View.GONE
-        actionButton.visibility = View.VISIBLE
-        actionButton.setText(R.string.wireguard_verify_key)
-        actionButton.setOnClickListener {
+        verifySpinner.visibility = View.GONE
+        verifyButton.visibility = View.VISIBLE
+        verifyButton.setText(R.string.wireguard_verify_key)
+        verifyButton.setOnClickListener {
             onValidateKeyPress()
         }
     }
 
-    private fun showActionSpinner() {
-        actionButton.visibility = View.GONE
-        actionSpinner.visibility = View.VISIBLE
-    }
-
     private fun drawNoConnectionState() {
-        actionButton.setClickable(true)
         visitWebsiteView.setClickable(true)
-        actionButton.setAlpha(1f)
         visitWebsiteView.setAlpha(1f)
 
         when (tunnelState) {
             is TunnelState.Connecting, is TunnelState.Disconnecting -> {
                 statusMessage.setText(R.string.wireguard_key_connectivity)
                 statusMessage.visibility = View.VISIBLE
-                actionButton.visibility = View.GONE
-                actionSpinner.visibility = View.VISIBLE
+                generateButton.visibility = View.GONE
+                generateSpinner.visibility = View.VISIBLE
+                verifyButton.visibility = View.GONE
+                verifySpinner.visibility = View.VISIBLE
             }
             is TunnelState.Blocked -> {
                 statusMessage.setText(R.string.wireguard_key_blocked_state_message)
                 statusMessage.visibility = View.VISIBLE
-                actionButton.setClickable(false)
-                actionButton.setAlpha(0.5f)
+                generateButton.setClickable(false)
+                generateButton.setAlpha(0.5f)
+                verifyButton.setClickable(false)
+                verifyButton.setAlpha(0.5f)
                 visitWebsiteView.setClickable(false)
                 visitWebsiteView.setAlpha(0.5f)
             }
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 74ad587eea..b39b23cf04 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
@@ -40,13 +40,19 @@ class KeyStatusListener(val asyncDaemon: Deferred<MullvadDaemon>) {
         daemon?.onKeygenEvent = { event -> keyStatus = event }
         val wireguardKey = daemon?.getWireguardKey()
         if (wireguardKey != null) {
-            keyStatus = KeygenEvent.NewKey(wireguardKey, null)
+            keyStatus = KeygenEvent.NewKey(wireguardKey, null, null)
         }
     }
 
     fun generateKey() = GlobalScope.launch(Dispatchers.Default) {
             setUpJob.join()
-            keyStatus = daemon?.generateWireguardKey()
+            val oldStatus = keyStatus
+            val newStatus = daemon?.generateWireguardKey()
+            if (oldStatus is KeygenEvent.NewKey && newStatus is KeygenEvent.Failure) {
+                keyStatus = KeygenEvent.NewKey(oldStatus.publicKey, oldStatus.verified, newStatus.failure)
+            } else {
+                keyStatus = newStatus
+            }
     }
 
     fun verifyKey() = GlobalScope.launch(Dispatchers.Default) {
@@ -55,7 +61,7 @@ class KeyStatusListener(val asyncDaemon: Deferred<MullvadDaemon>) {
             // Only update verification status if the key is actually there
             when (val state = keyStatus) {
                 is KeygenEvent.NewKey -> {
-                    keyStatus = KeygenEvent.NewKey(state.publicKey, verified)
+                    keyStatus = KeygenEvent.NewKey(state.publicKey, verified, state.replacementFailure)
                 }
             }
     }
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 b93794d44c..46b2789c9f 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
@@ -1,7 +1,11 @@
 package net.mullvad.mullvadvpn.model
 
 sealed class KeygenEvent {
-    class NewKey(val publicKey: PublicKey, val verified: Boolean?) : KeygenEvent()
-    class TooManyKeys : KeygenEvent()
-    class GenerationFailure : KeygenEvent()
+    class NewKey(val publicKey: PublicKey, val verified: Boolean?, val replacementFailure: KeygenFailure?) : KeygenEvent()
+    class Failure(val failure: KeygenFailure) : KeygenEvent()
+}
+
+sealed class KeygenFailure {
+    class TooManyKeys : KeygenFailure()
+    class GenerationFailure : KeygenFailure()
 }
diff --git a/android/src/main/res/layout/wireguard_key.xml b/android/src/main/res/layout/wireguard_key.xml
index 90d41fbebd..da2bb39207 100644
--- a/android/src/main/res/layout/wireguard_key.xml
+++ b/android/src/main/res/layout/wireguard_key.xml
@@ -87,7 +87,7 @@
     </LinearLayout>
 
     <LinearLayout
-        android:id="@+id/wireguard_button_container"
+        android:id="@+id/wireguard_generate_button_container"
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:layout_marginTop="15dp"
@@ -104,14 +104,50 @@
             android:orientation="vertical">
 
             <Button
-                android:id="@+id/wg_key_button"
+                android:id="@+id/wg_generate_key_button"
+                style="@style/Button"
+                android:layout_width="fill_parent"
+                android:layout_height="fill_parent"
+                android:text="@string/wireguard_generate_key"/>
+
+            <ProgressBar
+                android:id="@+id/wg_generate_key_spinner"
+                android:layout_width="30dp"
+                android:layout_height="30dp"
+                android:indeterminate="true"
+                android:indeterminateDrawable="@drawable/icon_spinner"
+                android:indeterminateDuration="600"
+                android:indeterminateOnly="true"
+                android:visibility="gone" />
+        </RelativeLayout>
+    </LinearLayout>
+
+    <LinearLayout
+        android:id="@+id/wireguard_verify_button_container"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:layout_marginTop="15dp"
+        android:layout_marginBottom="15dp"
+        android:background="@drawable/cell_button_background"
+        android:clickable="true"
+        android:gravity="center"
+        android:orientation="vertical">
+
+        <RelativeLayout
+            android:layout_width="wrap_content"
+            android:layout_height="50dp"
+            android:gravity="center"
+            android:orientation="vertical">
+
+            <Button
+                android:id="@+id/wg_verify_key_button"
                 style="@style/Button"
                 android:layout_width="fill_parent"
                 android:layout_height="fill_parent"
                 android:text="@string/wireguard_verify_key"/>
 
             <ProgressBar
-                android:id="@+id/wg_action_spinner"
+                android:id="@+id/wg_verify_key_spinner"
                 android:layout_width="30dp"
                 android:layout_height="30dp"
                 android:indeterminate="true"
diff --git a/android/src/main/res/values/strings.xml b/android/src/main/res/values/strings.xml
index 229cb6c933..0ddba635d5 100644
--- a/android/src/main/res/values/strings.xml
+++ b/android/src/main/res/values/strings.xml
@@ -116,6 +116,7 @@
     <string name="wireguard_public_key">Public key</string>
     <string name="wireguard_verify_key">Verify key</string>
     <string name="wireguard_generate_key">Generate key</string>
+    <string name="wireguard_replace_key">Replace key</string>
     <string name="wireguard_manage_keys">Manage keys</string>
     <string name="wireguard_key_connectivity">
         Connectivity required to manage your key.
diff --git a/mullvad-daemon/src/logging.rs b/mullvad-daemon/src/logging.rs
index bd2b00864c..2ffac0f1bc 100644
--- a/mullvad-daemon/src/logging.rs
+++ b/mullvad-daemon/src/logging.rs
@@ -101,7 +101,9 @@ pub fn init_logger(
     #[cfg(all(target_os = "android", debug_assertions))]
     {
         use android_logger::{AndroidLogger, Config};
-        let logger: Box<dyn log::Log> = Box::new(AndroidLogger::new(Config::default().with_tag("mullvad-daemon")));
+        let logger: Box<dyn log::Log> = Box::new(AndroidLogger::new(
+            Config::default().with_tag("mullvad-daemon"),
+        ));
         top_dispatcher = top_dispatcher.chain(logger);
     }
     top_dispatcher.apply().map_err(Error::SetLoggerError)?;
diff --git a/mullvad-jni/src/into_java.rs b/mullvad-jni/src/into_java.rs
index c99ed45113..e85e5b2f3f 100644
--- a/mullvad-jni/src/into_java.rs
+++ b/mullvad-jni/src/into_java.rs
@@ -551,26 +551,46 @@ impl<'env> IntoJava<'env> for KeygenEvent {
                 let parameters = [
                     JValue::Object(java_public_key.as_obj()),
                     JValue::Object(JObject::null()),
+                    JValue::Object(JObject::null()),
                 ];
 
                 env.new_object(
                     &class,
-                    "(Lnet/mullvad/mullvadvpn/model/PublicKey;Ljava/lang/Boolean;)V",
+                    "(Lnet/mullvad/mullvadvpn/model/PublicKey;Ljava/lang/Boolean;Lnet/mullvad/mullvadvpn/model/KeygenFailure;)V",
                     &parameters,
                 )
                 .expect("Failed to create KeygenEvent.NewKey Java object")
             }
             KeygenEvent::TooManyKeys => {
-                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$TooManyKeys");
+                let failure_class =
+                    get_class("net/mullvad/mullvadvpn/model/KeygenFailure$TooManyKeys");
 
-                env.new_object(&class, "()V", &[])
-                    .expect("Failed to create KeygenEvent.TooManyKeys Java object")
+                let failure = env
+                    .new_object(&failure_class, "()V", &[])
+                    .expect("Failed to create KeygenFailure.TooManyKeys Java object");
+
+                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$Failure");
+                env.new_object(
+                    &class,
+                    "(Lnet/mullvad/mullvadvpn/model/KeygenFailure;)V",
+                    &[JValue::Object(failure)],
+                )
+                .expect("Failed to create KeygenEvent.Failure Java object")
             }
             KeygenEvent::GenerationFailure => {
-                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$GenerationFailure");
+                let failure_class =
+                    get_class("net/mullvad/mullvadvpn/model/KeygenFailure$GenerationFailure");
+                let failure = env
+                    .new_object(&failure_class, "()V", &[])
+                    .expect("Failed to create KeygenFailure.GenerationFailure Java object");
 
-                env.new_object(&class, "()V", &[])
-                    .expect("Failed to create KeygenEvent.GenerationFailure Java object")
+                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$Failure");
+                env.new_object(
+                    &class,
+                    "(Lnet/mullvad/mullvadvpn/model/KeygenFailure;)V",
+                    &[JValue::Object(failure)],
+                )
+                .expect("Failed to create KeygenEvent.GenerationFailure Java object")
             }
         }
     }
diff --git a/mullvad-jni/src/lib.rs b/mullvad-jni/src/lib.rs
index 0eaf92b795..afe7ce7229 100644
--- a/mullvad-jni/src/lib.rs
+++ b/mullvad-jni/src/lib.rs
@@ -54,8 +54,9 @@ const CLASSES_TO_LOAD: &[&str] = &[
     "net/mullvad/mullvadvpn/model/GeoIpLocation",
     "net/mullvad/mullvadvpn/model/InetNetwork",
     "net/mullvad/mullvadvpn/model/KeygenEvent$NewKey",
-    "net/mullvad/mullvadvpn/model/KeygenEvent$TooManyKeys",
-    "net/mullvad/mullvadvpn/model/KeygenEvent$GenerationFailure",
+    "net/mullvad/mullvadvpn/model/KeygenEvent$Failure",
+    "net/mullvad/mullvadvpn/model/KeygenFailure$TooManyKeys",
+    "net/mullvad/mullvadvpn/model/KeygenFailure$GenerationFailure",
     "net/mullvad/mullvadvpn/model/LocationConstraint$City",
     "net/mullvad/mullvadvpn/model/LocationConstraint$Country",
     "net/mullvad/mullvadvpn/model/LocationConstraint$Hostname",