Merge branch 'handle-key-generation'

13 files changed, 254 insertions, 24 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectActionButton.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectActionButton.kt
index a430ab84c7..00ebdd8f64 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectActionButton.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectActionButton.kt
@@ -3,12 +3,47 @@ package net.mullvad.mullvadvpn
 import android.view.View
 import android.widget.Button
 
+import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.TunnelState
 
 class ConnectActionButton(val parentView: View) {
     private val button: Button = parentView.findViewById(R.id.action_button)
 
-    var state: TunnelState = TunnelState.Disconnected()
+    private var enabled = true
+        set(value) {
+            if (field != value) {
+                field = value
+
+                button.setEnabled(value)
+                button.setAlpha(if (value) 1.0F else 0.5F)
+            }
+        }
+
+    private var canConnect = true
+        set(value) {
+            field = value
+            updateEnabled()
+        }
+
+    private var showingConnect = true
+        set(value) {
+            field = value
+            updateEnabled()
+        }
+
+    var keyState: KeygenEvent? = null
+        set(value) {
+            when (value) {
+                null -> canConnect = true
+                is KeygenEvent.NewKey -> canConnect = true
+                is KeygenEvent.TooManyKeys -> canConnect = false
+                is KeygenEvent.GenerationFailure -> canConnect = false
+            }
+
+            field = value
+        }
+
+    var tunnelState: TunnelState = TunnelState.Disconnected()
         set(value) {
             when (value) {
                 is TunnelState.Disconnected -> disconnected()
@@ -30,7 +65,7 @@ class ConnectActionButton(val parentView: View) {
     }
 
     private fun action() {
-        when (state) {
+        when (tunnelState) {
             is TunnelState.Disconnected -> onConnect?.invoke()
             is TunnelState.Disconnecting -> onConnect?.invoke()
             is TunnelState.Connecting -> onCancel?.invoke()
@@ -42,15 +77,22 @@ class ConnectActionButton(val parentView: View) {
     private fun disconnected() {
         button.setBackgroundResource(R.drawable.green_button_background)
         button.setText(R.string.connect)
+        showingConnect = true
     }
 
     private fun connecting() {
         button.setBackgroundResource(R.drawable.transparent_red_button_background)
         button.setText(R.string.cancel)
+        showingConnect = false
     }
 
     private fun connected() {
         button.setBackgroundResource(R.drawable.transparent_red_button_background)
         button.setText(R.string.disconnect)
+        showingConnect = false
+    }
+
+    private fun updateEnabled() {
+        enabled = !showingConnect || canConnect
     }
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
index 9a7c50a9af..a5f14cfa92 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
@@ -14,8 +14,10 @@ import android.view.ViewGroup
 import android.widget.ImageButton
 
 import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
+import net.mullvad.mullvadvpn.dataproxy.KeyStatusListener
 import net.mullvad.mullvadvpn.dataproxy.LocationInfoCache
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
+import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.TunnelState
 
 class ConnectFragment : Fragment() {
@@ -28,16 +30,19 @@ class ConnectFragment : Fragment() {
 
     private lateinit var parentActivity: MainActivity
     private lateinit var connectionProxy: ConnectionProxy
+    private lateinit var keyStatusListener: KeyStatusListener
     private lateinit var locationInfoCache: LocationInfoCache
     private lateinit var relayListListener: RelayListListener
 
-    private var updateViewJob: Job? = null
+    private lateinit var updateKeyStatusJob: Job
+    private lateinit var updateTunnelStateJob: Job
 
     override fun onAttach(context: Context) {
         super.onAttach(context)
 
         parentActivity = context as MainActivity
         connectionProxy = parentActivity.connectionProxy
+        keyStatusListener = parentActivity.keyStatusListener
         locationInfoCache = parentActivity.locationInfoCache
         relayListListener = parentActivity.relayListListener
     }
@@ -68,13 +73,12 @@ class ConnectFragment : Fragment() {
         switchLocationButton = SwitchLocationButton(view)
         switchLocationButton.onClick = { openSwitchLocationScreen() }
 
-        updateView(connectionProxy.uiState)
+        updateKeyStatusJob = updateKeyStatus(keyStatusListener.keyStatus)
+        updateTunnelStateJob = updateTunnelState(connectionProxy.uiState)
 
         connectionProxy.onUiStateChange = { uiState ->
-            updateViewJob?.cancel()
-            updateViewJob = GlobalScope.launch(Dispatchers.Main) {
-                updateView(uiState)
-            }
+            updateTunnelStateJob.cancel()
+            updateTunnelStateJob = updateTunnelState(uiState)
         }
 
         return view
@@ -83,12 +87,18 @@ class ConnectFragment : Fragment() {
     override fun onResume() {
         super.onResume()
 
+        keyStatusListener.onKeyStatusChange = { keyStatus ->
+            updateKeyStatusJob.cancel()
+            updateKeyStatusJob = updateKeyStatus(keyStatus)
+        }
+
         relayListListener.onRelayListChange = { relayList, selectedRelayItem ->
             switchLocationButton.location = selectedRelayItem
         }
     }
 
     override fun onPause() {
+        keyStatusListener.onKeyStatusChange = null
         relayListListener.onRelayListChange = null
 
         super.onPause()
@@ -99,24 +109,28 @@ class ConnectFragment : Fragment() {
         switchLocationButton.onDestroy()
 
         connectionProxy.onUiStateChange = null
-        updateViewJob?.cancel()
+        updateTunnelStateJob.cancel()
 
         super.onDestroyView()
     }
 
-    private fun updateView(uiState: TunnelState) {
+    private fun updateTunnelState(uiState: TunnelState) = GlobalScope.launch(Dispatchers.Main) {
         val realState = connectionProxy.state
 
         locationInfoCache.state = realState
         headerBar.setState(realState)
 
-        actionButton.state = uiState
+        actionButton.tunnelState = uiState
         switchLocationButton.state = uiState
-
-        notificationBanner.setState(uiState)
+        notificationBanner.tunnelState = uiState
         status.setState(uiState)
     }
 
+    private fun updateKeyStatus(keyStatus: KeygenEvent?) = GlobalScope.launch(Dispatchers.Main) {
+        notificationBanner.keyState = keyStatus
+        actionButton.keyState = keyStatus
+    }
+
     private fun openSwitchLocationScreen() {
         fragmentManager?.beginTransaction()?.apply {
             setCustomAnimations(
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt
index b75a408ece..f5fbab1526 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/MainActivity.kt
@@ -20,6 +20,7 @@ import android.support.v4.app.FragmentActivity
 
 import net.mullvad.mullvadvpn.dataproxy.AccountCache
 import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
+import net.mullvad.mullvadvpn.dataproxy.KeyStatusListener
 import net.mullvad.mullvadvpn.dataproxy.LocationInfoCache
 import net.mullvad.mullvadvpn.dataproxy.MullvadProblemReport
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
@@ -36,6 +37,7 @@ class MainActivity : FragmentActivity() {
         private set
 
     val connectionProxy = ConnectionProxy(this)
+    val keyStatusListener = KeyStatusListener(daemon)
     val locationInfoCache = LocationInfoCache(daemon)
     val problemReport = MullvadProblemReport()
     var settingsListener = SettingsListener(this)
@@ -93,6 +95,7 @@ class MainActivity : FragmentActivity() {
 
     override fun onDestroy() {
         accountCache.onDestroy()
+        keyStatusListener.onDestroy()
         relayListListener.onDestroy()
         settingsListener.onDestroy()
 
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/MullvadDaemon.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/MullvadDaemon.kt
index eacedcdb69..769a2f1ad9 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/MullvadDaemon.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/MullvadDaemon.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn
 
 import net.mullvad.mullvadvpn.model.AccountData
 import net.mullvad.mullvadvpn.model.GeoIpLocation
+import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.PublicKey
 import net.mullvad.mullvadvpn.model.RelayList
 import net.mullvad.mullvadvpn.model.RelaySettingsUpdate
@@ -14,6 +15,7 @@ class MullvadDaemon(val vpnService: MullvadVpnService) {
         initialize(vpnService)
     }
 
+    var onKeygenEvent: ((KeygenEvent) -> Unit)? = null
     var onRelayListChange: ((RelayList) -> Unit)? = null
     var onSettingsChange: ((Settings) -> Unit)? = null
     var onTunnelStateChange: ((TunnelState) -> Unit)? = null
@@ -32,6 +34,10 @@ class MullvadDaemon(val vpnService: MullvadVpnService) {
 
     private external fun initialize(vpnService: MullvadVpnService)
 
+    private fun notifyKeygenEvent(event: KeygenEvent) {
+        onKeygenEvent?.invoke(event)
+    }
+
     private fun notifyRelayListEvent(relayList: RelayList) {
         onRelayListChange?.invoke(relayList)
     }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
index 1af739dcc6..b90b3aa251 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/NotificationBanner.kt
@@ -1,30 +1,65 @@
 package net.mullvad.mullvadvpn
 
+import android.widget.TextView
 import android.view.View
 
+import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.TunnelState
 
 class NotificationBanner(val parentView: View) {
     private val banner: View = parentView.findViewById(R.id.notification_banner)
+    private val title: TextView = parentView.findViewById(R.id.notification_title)
+
     private var visible = false
 
-    fun setState(state: TunnelState) {
-        when (state) {
+    var keyState: KeygenEvent? = null
+        set(value) {
+            field = value
+            update()
+        }
+
+    var tunnelState: TunnelState = TunnelState.Disconnected()
+        set(value) {
+            field = value
+            update()
+        }
+
+    private fun update() {
+        updateBasedOnKeyState() || updateBasedOnTunnelState()
+    }
+
+    private fun updateBasedOnKeyState(): Boolean {
+        when (keyState) {
+            null -> return false
+            is KeygenEvent.NewKey -> return false
+            is KeygenEvent.TooManyKeys -> show(R.string.too_many_keys)
+            is KeygenEvent.GenerationFailure -> show(R.string.failed_to_generate_key)
+        }
+
+        return true
+    }
+
+    private fun updateBasedOnTunnelState(): Boolean {
+        when (tunnelState) {
             is TunnelState.Disconnecting -> hide()
             is TunnelState.Disconnected -> hide()
-            is TunnelState.Connecting -> show()
+            is TunnelState.Connecting -> show(R.string.blocking_internet)
             is TunnelState.Connected -> hide()
-            is TunnelState.Blocked -> show()
+            is TunnelState.Blocked -> show(R.string.blocking_internet)
         }
+
+        return true
     }
 
-    private fun show() {
+    private fun show(message: Int) {
         if (!visible) {
             visible = true
             banner.visibility = View.VISIBLE
             banner.translationY = -banner.height.toFloat()
             banner.animate().translationY(0.0F).setDuration(350).start()
         }
+
+        title.setText(message)
     }
 
     private fun hide() {
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt
index eae0ed75f2..4a1849318c 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/SelectLocationFragment.kt
@@ -19,6 +19,7 @@ import android.widget.ViewSwitcher
 import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
 import net.mullvad.mullvadvpn.model.Constraint
+import net.mullvad.mullvadvpn.model.KeygenEvent
 import net.mullvad.mullvadvpn.model.LocationConstraint
 import net.mullvad.mullvadvpn.model.RelaySettingsUpdate
 import net.mullvad.mullvadvpn.relaylist.RelayItem
@@ -40,7 +41,7 @@ class SelectLocationFragment : Fragment() {
     init {
         relayListAdapter.onSelect = { relayItem ->
             updateLocationConstraint(relayItem)
-            connectionProxy.connect()
+            maybeConnect()
             close()
         }
     }
@@ -123,4 +124,12 @@ class SelectLocationFragment : Fragment() {
             relayListContainer.showNext()
         }
     }
+
+    private fun maybeConnect() {
+        val keyStatus = parentActivity.keyStatusListener.keyStatus
+
+        if (keyStatus == null || keyStatus is KeygenEvent.NewKey) {
+            connectionProxy.connect()
+        }
+    }
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
new file mode 100644
index 0000000000..1279f59567
--- /dev/null
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/KeyStatusListener.kt
@@ -0,0 +1,45 @@
+package net.mullvad.mullvadvpn.dataproxy
+
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+
+import net.mullvad.mullvadvpn.MullvadDaemon
+import net.mullvad.mullvadvpn.model.KeygenEvent
+
+class KeyStatusListener(val asyncDaemon: Deferred<MullvadDaemon>) {
+    private var daemon: MullvadDaemon? = null
+
+    private val setUpJob = setUp()
+
+    var keyStatus: KeygenEvent? = null
+        private set(value) {
+            synchronized(this) {
+                field = value
+
+                if (value != null) {
+                    onKeyStatusChange?.invoke(value)
+                }
+            }
+        }
+
+    var onKeyStatusChange: ((KeygenEvent) -> Unit)? = null
+        set(value) {
+            field = value
+
+            synchronized(this) {
+                keyStatus?.let { status -> value?.invoke(status) }
+            }
+        }
+
+    private fun setUp() = GlobalScope.launch(Dispatchers.Default) {
+        daemon = asyncDaemon.await()
+        daemon?.onKeygenEvent = { event -> keyStatus = event }
+    }
+
+    fun onDestroy() {
+        setUpJob.cancel()
+        daemon?.onKeygenEvent = null
+    }
+}
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
new file mode 100644
index 0000000000..b5938d65f6
--- /dev/null
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/model/KeygenEvent.kt
@@ -0,0 +1,7 @@
+package net.mullvad.mullvadvpn.model
+
+sealed class KeygenEvent {
+    class NewKey(var publicKey: PublicKey) : KeygenEvent()
+    class TooManyKeys : KeygenEvent()
+    class GenerationFailure : KeygenEvent()
+}
diff --git a/android/src/main/res/layout/connect.xml b/android/src/main/res/layout/connect.xml
index 5c2beb1e2e..38298d2882 100644
--- a/android/src/main/res/layout/connect.xml
+++ b/android/src/main/res/layout/connect.xml
@@ -62,7 +62,7 @@
                 android:layout_marginLeft="19dp"
                 android:src="@drawable/icon_notification_error"
                 />
-        <TextView
+        <TextView android:id="@+id/notification_title"
                 android:layout_width="0dp"
                 android:layout_height="wrap_content"
                 android:layout_weight="1"
diff --git a/android/src/main/res/values/strings.xml b/android/src/main/res/values/strings.xml
index a3da5d4e96..711be12793 100644
--- a/android/src/main/res/values/strings.xml
+++ b/android/src/main/res/values/strings.xml
@@ -54,12 +54,15 @@
     <string name="unsecured_connection">Unsecured connection</string>
     <string name="creating_secure_connection">Creating secure connection</string>
     <string name="secure_connection">Secure connection</string>
-    <string name="blocking_internet">Blocking internet</string>
     <string name="connect">Secure my connection</string>
     <string name="cancel">Cancel</string>
     <string name="disconnect">Disconnect</string>
     <string name="switch_location">Switch location</string>
 
+    <string name="blocking_internet">Blocking internet</string>
+    <string name="too_many_keys">Too many WireGuard keys registered to account</string>
+    <string name="failed_to_generate_key">Failed to generate WireGuard key</string>
+
     <string name="select_location">Select location</string>
     <string name="select_location_description">
         While connected, your real location is masked with a private and secure location in the
diff --git a/mullvad-jni/src/into_java.rs b/mullvad-jni/src/into_java.rs
index bb29c20bc9..fc899bf12e 100644
--- a/mullvad-jni/src/into_java.rs
+++ b/mullvad-jni/src/into_java.rs
@@ -13,6 +13,7 @@ use mullvad_types::{
     relay_list::{Relay, RelayList, RelayListCity, RelayListCountry},
     settings::Settings,
     states::TunnelState,
+    wireguard::KeygenEvent,
     CustomTunnelEndpoint,
 };
 use std::{fmt::Debug, net::IpAddr};
@@ -423,6 +424,39 @@ impl<'env> IntoJava<'env> for CustomTunnelEndpoint {
     }
 }
 
+impl<'env> IntoJava<'env> for KeygenEvent {
+    type JavaType = JObject<'env>;
+
+    fn into_java(self, env: &JNIEnv<'env>) -> Self::JavaType {
+        match self {
+            KeygenEvent::NewKey(public_key) => {
+                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$NewKey");
+                let java_public_key = env.auto_local(public_key.into_java(env));
+                let parameters = [JValue::Object(java_public_key.as_obj())];
+
+                env.new_object(
+                    &class,
+                    "(Lnet/mullvad/mullvadvpn/model/PublicKey;)V",
+                    &parameters,
+                )
+                .expect("Failed to create KeygenEvent.NewKey Java object")
+            }
+            KeygenEvent::TooManyKeys => {
+                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$TooManyKeys");
+
+                env.new_object(&class, "()V", &[])
+                    .expect("Failed to create KeygenEvent.TooManyKeys Java object")
+            }
+            KeygenEvent::GenerationFailure => {
+                let class = get_class("net/mullvad/mullvadvpn/model/KeygenEvent$GenerationFailure");
+
+                env.new_object(&class, "()V", &[])
+                    .expect("Failed to create KeygenEvent.GenerationFailure Java object")
+            }
+        }
+    }
+}
+
 impl<'env> IntoJava<'env> for RelayConstraints {
     type JavaType = JObject<'env>;
 
diff --git a/mullvad-jni/src/jni_event_listener.rs b/mullvad-jni/src/jni_event_listener.rs
index 8fa2456d5b..cb0e14fb40 100644
--- a/mullvad-jni/src/jni_event_listener.rs
+++ b/mullvad-jni/src/jni_event_listener.rs
@@ -24,6 +24,7 @@ pub enum Error {
 }
 
 enum Event {
+    KeygenEvent(KeygenEvent),
     RelayList(RelayList),
     Settings(Settings),
     Tunnel(TunnelState),
@@ -39,6 +40,10 @@ impl JniEventListener {
 }
 
 impl EventListener for JniEventListener {
+    fn notify_key_event(&self, key_event: KeygenEvent) {
+        let _ = self.0.send(Event::KeygenEvent(key_event));
+    }
+
     fn notify_new_state(&self, state: TunnelState) {
         let _ = self.0.send(Event::Tunnel(state));
     }
@@ -50,14 +55,12 @@ impl EventListener for JniEventListener {
     fn notify_relay_list(&self, relay_list: RelayList) {
         let _ = self.0.send(Event::RelayList(relay_list));
     }
-
-    // TODO: manage key events properly
-    fn notify_key_event(&self, _key_event: KeygenEvent) {}
 }
 
 struct JniEventHandler<'env> {
     env: AttachGuard<'env>,
     mullvad_ipc_client: JObject<'env>,
+    notify_keygen_event: JMethodID<'env>,
     notify_relay_list_event: JMethodID<'env>,
     notify_settings_event: JMethodID<'env>,
     notify_tunnel_event: JMethodID<'env>,
@@ -101,6 +104,12 @@ impl<'env> JniEventHandler<'env> {
         events: mpsc::Receiver<Event>,
     ) -> Result<Self, Error> {
         let class = get_class("net/mullvad/mullvadvpn/MullvadDaemon");
+        let notify_keygen_event = Self::get_method_id(
+            &env,
+            &class,
+            "notifyKeygenEvent",
+            "(Lnet/mullvad/mullvadvpn/model/KeygenEvent;)V",
+        )?;
         let notify_relay_list_event = Self::get_method_id(
             &env,
             &class,
@@ -123,6 +132,7 @@ impl<'env> JniEventHandler<'env> {
         Ok(JniEventHandler {
             env,
             mullvad_ipc_client,
+            notify_keygen_event,
             notify_relay_list_event,
             notify_settings_event,
             notify_tunnel_event,
@@ -143,6 +153,7 @@ impl<'env> JniEventHandler<'env> {
     fn run(&mut self) {
         while let Ok(event) = self.events.recv() {
             match event {
+                Event::KeygenEvent(keygen_event) => self.handle_keygen_event(keygen_event),
                 Event::RelayList(relay_list) => self.handle_relay_list_event(relay_list),
                 Event::Settings(settings) => self.handle_settings(settings),
                 Event::Tunnel(tunnel_event) => self.handle_tunnel_event(tunnel_event),
@@ -150,6 +161,24 @@ impl<'env> JniEventHandler<'env> {
         }
     }
 
+    fn handle_keygen_event(&self, event: KeygenEvent) {
+        let java_keygen_event = self.env.auto_local(event.into_java(&self.env));
+
+        let result = self.env.call_method_unchecked(
+            self.mullvad_ipc_client,
+            self.notify_keygen_event,
+            JavaType::Primitive(Primitive::Void),
+            &[JValue::Object(java_keygen_event.as_obj())],
+        );
+
+        if let Err(error) = result {
+            log::error!(
+                "{}",
+                error.display_chain_with_msg("Failed to call MullvadDaemon.notifyKeygenEvent")
+            );
+        }
+    }
+
     fn handle_relay_list_event(&self, relay_list: RelayList) {
         let java_relay_list = self.env.auto_local(relay_list.into_java(&self.env));
 
diff --git a/mullvad-jni/src/lib.rs b/mullvad-jni/src/lib.rs
index 2ec4cac08a..dfc55010ac 100644
--- a/mullvad-jni/src/lib.rs
+++ b/mullvad-jni/src/lib.rs
@@ -39,6 +39,9 @@ const CLASSES_TO_LOAD: &[&str] = &[
     "net/mullvad/mullvadvpn/model/Constraint$Only",
     "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/LocationConstraint$City",
     "net/mullvad/mullvadvpn/model/LocationConstraint$Country",
     "net/mullvad/mullvadvpn/model/LocationConstraint$Hostname",