Refactor daemon init and deinit on Android

11 files changed, 333 insertions, 541 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 63e99f8cca..c6a1e328fb 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2252,6 +2252,7 @@ dependencies = [
  "talpid-tunnel",
  "talpid-types",
  "thiserror",
+ "tokio",
 ]
 
 [[package]]
diff --git a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
index 2c441483f6..cdaddec614 100644
--- a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
+++ b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
@@ -1,75 +1,13 @@
 package net.mullvad.mullvadvpn.service
 
-import android.annotation.SuppressLint
-import android.content.Context
-import java.io.File
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.async
-import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withContext
 import net.mullvad.mullvadvpn.lib.endpoint.ApiEndpoint
-import net.mullvad.mullvadvpn.lib.endpoint.ApiEndpointConfiguration
-import net.mullvad.mullvadvpn.service.migration.MigrateSplitTunneling
-
-private const val RELAYS_FILE = "relays.json"
-
-@SuppressLint("SdCardPath")
-class MullvadDaemon(
-    vpnService: MullvadVpnService,
-    rpcSocketFile: File,
-    apiEndpointConfiguration: ApiEndpointConfiguration,
-    migrateSplitTunneling: MigrateSplitTunneling
-) {
-    // Used by JNI
-    @Suppress("ProtectedMemberInFinalClass") protected var daemonInterfaceAddress = 0L
-
-    private val shutdownSignal = Channel<Unit>()
 
+object MullvadDaemon {
     init {
         System.loadLibrary("mullvad_jni")
-
-        prepareFiles(vpnService)
-
-        migrateSplitTunneling.migrate()
-
-        initialize(
-            vpnService = vpnService,
-            rpcSocketPath = rpcSocketFile.absolutePath,
-            filesDirectory = vpnService.filesDir.absolutePath,
-            cacheDirectory = vpnService.cacheDir.absolutePath,
-            apiEndpoint = apiEndpointConfiguration.apiEndpoint()
-        )
-    }
-
-    suspend fun shutdown() =
-        withContext(Dispatchers.IO) {
-            val shutdownSignal = async { shutdownSignal.receive() }
-            shutdown(daemonInterfaceAddress)
-            shutdownSignal.await()
-            deinitialize()
-        }
-
-    private fun prepareFiles(context: Context) {
-        val shouldOverwriteRelayList =
-            lastUpdatedTime(context) > File(context.filesDir, RELAYS_FILE).lastModified()
-
-        FileResourceExtractor(context).apply { extract(RELAYS_FILE, shouldOverwriteRelayList) }
     }
 
-    private fun lastUpdatedTime(context: Context): Long =
-        context.packageManager.getPackageInfo(context.packageName, 0).lastUpdateTime
-
-    // Used by JNI
-    @Suppress("unused")
-    private fun notifyDaemonStopped() {
-        runBlocking {
-            shutdownSignal.send(Unit)
-            shutdownSignal.close()
-        }
-    }
-
-    private external fun initialize(
+    external fun initialize(
         vpnService: MullvadVpnService,
         rpcSocketPath: String,
         filesDirectory: String,
@@ -77,7 +15,5 @@ class MullvadDaemon(
         apiEndpoint: ApiEndpoint?
     )
 
-    private external fun deinitialize()
-
-    private external fun shutdown(daemonInterfaceAddress: Long)
+    external fun shutdown()
 }
diff --git a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt
index 87f6076c3f..8d806500c8 100644
--- a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt
+++ b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt
@@ -1,6 +1,7 @@
 package net.mullvad.mullvadvpn.service
 
 import android.app.KeyguardManager
+import android.content.Context
 import android.content.Intent
 import android.os.Binder
 import android.os.Build
@@ -18,6 +19,7 @@ import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import net.mullvad.mullvadvpn.lib.common.constant.BuildTypes
 import net.mullvad.mullvadvpn.lib.common.constant.GRPC_SOCKET_FILE_NAMED_ARGUMENT
 import net.mullvad.mullvadvpn.lib.common.constant.KEY_CONNECT_ACTION
 import net.mullvad.mullvadvpn.lib.common.constant.KEY_DISCONNECT_ACTION
@@ -40,12 +42,13 @@ import org.koin.android.ext.android.getKoin
 import org.koin.core.context.loadKoinModules
 import org.koin.core.qualifier.named
 
+private const val RELAYS_FILE = "relays.json"
+
 class MullvadVpnService : TalpidVpnService(), ShouldBeOnForegroundProvider {
     private val _shouldBeOnForeground = MutableStateFlow(false)
     override val shouldBeOnForeground: StateFlow<Boolean> = _shouldBeOnForeground
 
     private lateinit var keyguardManager: KeyguardManager
-    private lateinit var daemonInstance: MullvadDaemon
 
     private lateinit var apiEndpointConfiguration: ApiEndpointConfiguration
     private lateinit var managementService: ManagementService
@@ -90,15 +93,11 @@ class MullvadVpnService : TalpidVpnService(), ShouldBeOnForegroundProvider {
         // with intent from API)
         lifecycleScope.launch(context = Dispatchers.IO) {
             managementService.start()
-            daemonInstance =
-                MullvadDaemon(
-                    vpnService = this@MullvadVpnService,
-                    rpcSocketFile = rpcSocketFile,
-                    apiEndpointConfiguration =
-                        intentProvider.getLatestIntent()?.getApiEndpointConfigurationExtras()
-                            ?: apiEndpointConfiguration,
-                    migrateSplitTunneling = migrateSplitTunneling
-                )
+
+            prepareFiles(this@MullvadVpnService)
+            migrateSplitTunneling.migrate()
+
+            startDaemon()
         }
     }
 
@@ -146,6 +145,24 @@ class MullvadVpnService : TalpidVpnService(), ShouldBeOnForegroundProvider {
         return super.onBind(intent) ?: emptyBinder()
     }
 
+    private fun startDaemon() {
+        val apiEndpointConfiguration =
+            if (Build.TYPE == BuildTypes.DEBUG) {
+                intentProvider.getLatestIntent()?.getApiEndpointConfigurationExtras()
+                    ?: apiEndpointConfiguration
+            } else {
+                apiEndpointConfiguration
+            }
+
+        MullvadDaemon.initialize(
+            vpnService = this@MullvadVpnService,
+            rpcSocketPath = rpcSocketFile.absolutePath,
+            filesDirectory = filesDir.absolutePath,
+            cacheDirectory = cacheDir.absolutePath,
+            apiEndpoint = apiEndpointConfiguration.apiEndpoint()
+        )
+    }
+
     private fun emptyBinder() =
         if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
             Binder(this.toString())
@@ -193,7 +210,7 @@ class MullvadVpnService : TalpidVpnService(), ShouldBeOnForegroundProvider {
         managementService.stop()
 
         // Shutting down the daemon gracefully
-        runBlocking { daemonInstance.shutdown() }
+        MullvadDaemon.shutdown()
         super.onDestroy()
     }
 
@@ -202,6 +219,16 @@ class MullvadVpnService : TalpidVpnService(), ShouldBeOnForegroundProvider {
         return this?.action == SERVICE_INTERFACE
     }
 
+    private fun prepareFiles(context: Context) {
+        val shouldOverwriteRelayList =
+            lastUpdatedTime(context) > File(context.filesDir, RELAYS_FILE).lastModified()
+
+        FileResourceExtractor(context).apply { extract(RELAYS_FILE, shouldOverwriteRelayList) }
+    }
+
+    private fun lastUpdatedTime(context: Context): Long =
+        context.packageManager.getPackageInfo(context.packageName, 0).lastUpdateTime
+
     companion object {
         init {
             System.loadLibrary("mullvad_jni")
diff --git a/mullvad-daemon/src/access_method.rs b/mullvad-daemon/src/access_method.rs
index edae229ae0..e683165685 100644
--- a/mullvad-daemon/src/access_method.rs
+++ b/mullvad-daemon/src/access_method.rs
@@ -30,7 +30,7 @@ pub enum Error {
 
 impl<L> Daemon<L>
 where
-    L: EventListener + Clone + Send + 'static,
+    L: EventListener,
 {
     /// Add a [`AccessMethod`] to the daemon's settings.
     ///
diff --git a/mullvad-daemon/src/custom_list.rs b/mullvad-daemon/src/custom_list.rs
index 71c5fc77e9..f3e92f3c97 100644
--- a/mullvad-daemon/src/custom_list.rs
+++ b/mullvad-daemon/src/custom_list.rs
@@ -8,7 +8,7 @@ use talpid_types::net::TunnelType;
 
 impl<L> Daemon<L>
 where
-    L: EventListener + Clone + Send + 'static,
+    L: EventListener,
 {
     /// Create a new custom list.
     ///
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index ae680e7fcf..269bd1dbb1 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -33,7 +33,7 @@ use api::AccessMethodEvent;
 use device::{AccountEvent, PrivateAccountAndDevice, PrivateDeviceEvent};
 use futures::{
     channel::{mpsc, oneshot},
-    future::{abortable, AbortHandle, Future, LocalBoxFuture},
+    future::{abortable, AbortHandle, Future},
     StreamExt,
 };
 use geoip::GeoIpHandler;
@@ -467,7 +467,7 @@ impl DaemonCommandChannel {
     }
 }
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct DaemonCommandSender(Arc<mpsc::UnboundedSender<InternalDaemonEvent>>);
 
 impl DaemonCommandSender {
@@ -540,7 +540,7 @@ where
 }
 
 /// Trait representing something that can broadcast daemon events.
-pub trait EventListener {
+pub trait EventListener: Clone + Send + Sync + 'static {
     /// Notify that the tunnel state changed.
     fn notify_new_state(&self, new_state: TunnelState);
 
@@ -585,7 +585,7 @@ pub struct Daemon<L: EventListener> {
     relay_selector: RelaySelector,
     relay_list_updater: RelayListUpdaterHandle,
     parameters_generator: tunnel::ParametersGenerator,
-    shutdown_tasks: Vec<Pin<Box<dyn Future<Output = ()>>>>,
+    shutdown_tasks: Vec<Pin<Box<dyn Future<Output = ()> + Send + Sync>>>,
     tunnel_state_machine_handle: TunnelStateMachineHandle,
     #[cfg(target_os = "windows")]
     volume_update_tx: mpsc::UnboundedSender<()>,
@@ -594,7 +594,7 @@ pub struct Daemon<L: EventListener> {
 
 impl<L> Daemon<L>
 where
-    L: EventListener + Clone + Send + 'static,
+    L: EventListener,
 {
     pub async fn start(
         log_dir: Option<PathBuf>,
@@ -897,46 +897,27 @@ where
     /// Destroy daemon safely, by dropping all objects in the correct order, waiting for them to
     /// be destroyed, and executing shutdown tasks
     async fn finalize(self) {
-        let (event_listener, shutdown_tasks, api_runtime, tunnel_state_machine_handle) =
-            self.shutdown();
-        for future in shutdown_tasks {
-            future.await;
-        }
-
-        tunnel_state_machine_handle.try_join().await;
-
-        drop(event_listener);
-        drop(api_runtime);
-    }
-
-    /// Shuts down the daemon without shutting down the underlying event listener and the shutdown
-    /// callbacks
-    fn shutdown<'a>(
-        self,
-    ) -> (
-        L,
-        Vec<LocalBoxFuture<'a, ()>>,
-        mullvad_api::Runtime,
-        TunnelStateMachineHandle,
-    ) {
         let Daemon {
             event_listener,
-            mut shutdown_tasks,
+            shutdown_tasks,
             api_runtime,
             tunnel_state_machine_handle,
             target_state,
             account_manager,
             ..
         } = self;
-        shutdown_tasks.push(Box::pin(target_state.finalize()));
-        shutdown_tasks.push(Box::pin(account_manager.shutdown()));
 
-        (
-            event_listener,
-            shutdown_tasks,
-            api_runtime,
-            tunnel_state_machine_handle,
-        )
+        for future in shutdown_tasks {
+            future.await;
+        }
+
+        target_state.finalize().await;
+        account_manager.shutdown().await;
+
+        tunnel_state_machine_handle.try_join().await;
+
+        drop(event_listener);
+        drop(api_runtime);
     }
 
     async fn handle_event(&mut self, event: InternalDaemonEvent) -> bool {
@@ -1686,7 +1667,7 @@ where
 
     #[cfg(not(target_os = "android"))]
     async fn on_factory_reset(&mut self, tx: ResponseTx<(), Error>) {
-        let mut last_error = Ok(());
+        let mut last_error = None;
 
         if let Err(error) = self.account_manager.logout().await {
             log::error!(
@@ -1700,12 +1681,12 @@ where
                 "{}",
                 error.display_chain_with_msg("Failed to clear account history")
             );
-            last_error = Err(Error::FactoryResetError("Failed to clear account history"));
+            last_error = Some("Failed to clear account history");
         }
 
         if let Err(e) = self.settings.reset().await {
             log::error!("Failed to reset settings: {}", e);
-            last_error = Err(Error::FactoryResetError("Failed to reset settings"));
+            last_error = Some("Failed to reset settings");
         }
 
         // Shut the daemon down.
@@ -1717,11 +1698,12 @@ where
                     "{}",
                     e.display_chain_with_msg("Failed to clear cache and log directories")
                 );
-                last_error = Err(Error::FactoryResetError(
-                    "Failed to clear cache and log directories",
-                ));
+                last_error = Some("Failed to clear cache and log directories");
             }
-            Self::oneshot_send(tx, last_error, "factory_reset response");
+            let result = last_error
+                .map(|error| Err(Error::FactoryResetError(error)))
+                .unwrap_or(Ok(()));
+            Self::oneshot_send(tx, result, "factory_reset response");
         }));
     }
 
diff --git a/mullvad-daemon/src/settings/mod.rs b/mullvad-daemon/src/settings/mod.rs
index ef666c34ef..f16a777b54 100644
--- a/mullvad-daemon/src/settings/mod.rs
+++ b/mullvad-daemon/src/settings/mod.rs
@@ -94,7 +94,7 @@ pub struct SettingsPersister {
     settings: Settings,
     path: PathBuf,
     #[allow(clippy::type_complexity)]
-    on_change_listeners: Vec<Box<dyn Fn(&Settings)>>,
+    on_change_listeners: Vec<Box<dyn Fn(&Settings) + Send + Sync>>,
 }
 
 pub type MadeChanges = bool;
@@ -353,7 +353,10 @@ impl SettingsPersister {
         }
     }
 
-    pub fn register_change_listener(&mut self, change_listener: impl Fn(&Settings) + 'static) {
+    pub fn register_change_listener(
+        &mut self,
+        change_listener: impl Fn(&Settings) + Send + Sync + 'static,
+    ) {
         self.on_change_listeners.push(Box::new(change_listener));
     }
 
diff --git a/mullvad-jni/Cargo.toml b/mullvad-jni/Cargo.toml
index c964c3adc9..16db8958a7 100644
--- a/mullvad-jni/Cargo.toml
+++ b/mullvad-jni/Cargo.toml
@@ -18,6 +18,8 @@ api-override = ["mullvad-api/api-override"]
 crate-type = ["cdylib"]
 
 [target.'cfg(target_os = "android")'.dependencies]
+tokio = { workspace = true, features = ["rt"] }
+
 thiserror = { workspace = true }
 futures = "0.3"
 ipnetwork = "0.16"
diff --git a/mullvad-jni/src/api.rs b/mullvad-jni/src/api.rs
new file mode 100644
index 0000000000..3f76e74839
--- /dev/null
+++ b/mullvad-jni/src/api.rs
@@ -0,0 +1,103 @@
+use jnix::{
+    jni::{
+        objects::JObject,
+        signature::{JavaType, Primitive},
+    },
+    FromJava, JnixEnv,
+};
+use std::net::{IpAddr, SocketAddr};
+
+pub fn api_endpoint_from_java(
+    env: &JnixEnv<'_>,
+    object: JObject<'_>,
+) -> Option<mullvad_api::ApiEndpoint> {
+    if object.is_null() {
+        return None;
+    }
+
+    let mut endpoint = mullvad_api::ApiEndpoint::from_env_vars();
+
+    let address = env
+        .call_method(object, "component1", "()Ljava/net/InetSocketAddress;", &[])
+        .expect("missing ApiEndpoint.address")
+        .l()
+        .expect("ApiEndpoint.address is not an InetSocketAddress");
+
+    endpoint.address = Some(
+        try_socketaddr_from_java(env, address).expect("received unresolved InetSocketAddress"),
+    );
+    endpoint.host = try_hostname_from_java(env, address);
+    #[cfg(feature = "api-override")]
+    {
+        endpoint.disable_address_cache = env
+            .call_method(object, "component2", "()Z", &[])
+            .expect("missing ApiEndpoint.disableAddressCache")
+            .z()
+            .expect("ApiEndpoint.disableAddressCache is not a bool");
+        endpoint.disable_tls = env
+            .call_method(object, "component3", "()Z", &[])
+            .expect("missing ApiEndpoint.disableTls")
+            .z()
+            .expect("ApiEndpoint.disableTls is not a bool");
+    }
+
+    Some(endpoint)
+}
+
+/// Converts InetSocketAddress to a SocketAddr. Return `None` if the hostname is unresolved.
+fn try_socketaddr_from_java(env: &JnixEnv<'_>, address: JObject<'_>) -> Option<SocketAddr> {
+    let class = env.get_class("java/net/InetSocketAddress");
+
+    let method_id = env
+        .get_method_id(&class, "getAddress", "()Ljava/net/InetAddress;")
+        .expect("Failed to get method ID for InetSocketAddress.getAddress()");
+    let return_type = JavaType::Object("java/net/InetAddress".to_owned());
+
+    let ip_addr = env
+        .call_method_unchecked(address, method_id, return_type, &[])
+        .expect("Failed to call InetSocketAddress.getAddress()")
+        .l()
+        .expect("Call to InetSocketAddress.getAddress() did not return an object");
+
+    if ip_addr.is_null() {
+        return None;
+    }
+
+    let method_id = env
+        .get_method_id(&class, "getPort", "()I")
+        .expect("Failed to get method ID for InetSocketAddress.getPort()");
+    let return_type = JavaType::Primitive(Primitive::Int);
+
+    let port = env
+        .call_method_unchecked(address, method_id, return_type, &[])
+        .expect("Failed to call InetSocketAddress.getPort()")
+        .i()
+        .expect("Call to InetSocketAddress.getPort() did not return an int");
+
+    Some(SocketAddr::new(
+        IpAddr::from_java(env, ip_addr),
+        u16::try_from(port).expect("invalid port"),
+    ))
+}
+
+/// Returns the hostname for an InetSocketAddress. This may be an IP address converted to a string.
+fn try_hostname_from_java(env: &JnixEnv<'_>, address: JObject<'_>) -> Option<String> {
+    let class = env.get_class("java/net/InetSocketAddress");
+
+    let method_id = env
+        .get_method_id(&class, "getHostString", "()Ljava/lang/String;")
+        .expect("Failed to get method ID for InetSocketAddress.getHostString()");
+    let return_type = JavaType::Object("java/lang/String".to_owned());
+
+    let hostname = env
+        .call_method_unchecked(address, method_id, return_type, &[])
+        .expect("Failed to call InetSocketAddress.getHostString()")
+        .l()
+        .expect("Call to InetSocketAddress.getHostString() did not return an object");
+
+    if hostname.is_null() {
+        return None;
+    }
+
+    Some(String::from_java(env, hostname))
+}
diff --git a/mullvad-jni/src/daemon_interface.rs b/mullvad-jni/src/daemon_interface.rs
deleted file mode 100644
index 19a592e2b7..0000000000
--- a/mullvad-jni/src/daemon_interface.rs
+++ /dev/null
@@ -1,13 +0,0 @@
-use mullvad_daemon::{DaemonCommandSender, Error};
-
-pub struct DaemonInterface(DaemonCommandSender);
-
-impl DaemonInterface {
-    pub fn new(command_sender: DaemonCommandSender) -> Self {
-        DaemonInterface(command_sender)
-    }
-
-    pub fn shutdown(&self) -> Result<(), Error> {
-        self.0.shutdown()
-    }
-}
diff --git a/mullvad-jni/src/lib.rs b/mullvad-jni/src/lib.rs
index 67e40ea534..eb9467920c 100644
--- a/mullvad-jni/src/lib.rs
+++ b/mullvad-jni/src/lib.rs
@@ -1,45 +1,37 @@
 #![cfg(target_os = "android")]
 
+mod api;
 mod classes;
-mod daemon_interface;
 mod is_null;
 mod problem_report;
 mod talpid_vpn_service;
 
-use crate::daemon_interface::DaemonInterface;
 use jnix::{
     jni::{
-        objects::{GlobalRef, JObject, JValue},
-        signature::{JavaType, Primitive},
-        sys::jlong,
-        JNIEnv, JavaVM,
+        objects::{JClass, JObject},
+        JNIEnv,
     },
     FromJava, JnixEnv,
 };
 use mullvad_daemon::{
-    cleanup_old_rpc_socket, exception_logging, logging, runtime::new_multi_thread, version, Daemon,
-    DaemonCommandChannel,
+    cleanup_old_rpc_socket, exception_logging, logging,
+    management_interface::ManagementInterfaceServer, runtime::new_multi_thread, version, Daemon,
+    DaemonCommandChannel, DaemonCommandSender,
 };
 use std::{
     io,
     path::{Path, PathBuf},
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        mpsc, Arc, Once,
-    },
-    thread,
+    sync::{Arc, Mutex, Once, OnceLock},
 };
 use talpid_types::{android::AndroidContext, ErrorExt};
 
-#[cfg(feature = "api-override")]
-use std::net::{IpAddr, SocketAddr};
-
 const LOG_FILENAME: &str = "daemon.log";
 
-static DAEMON_INSTANCE_COUNT: AtomicUsize = AtomicUsize::new(0);
+/// Mullvad daemon instance. It must be initialized and destroyed by `MullvadDaemon.initialize` and
+/// `MullvadDaemon.shutdown`, respectively.
+static DAEMON_CONTEXT: Mutex<Option<DaemonContext>> = Mutex::new(None);
+
 static LOAD_CLASSES: Once = Once::new();
-static LOG_START: Once = Once::new();
-static mut LOG_INIT_RESULT: Option<Result<(), String>> = None;
 
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
@@ -49,455 +41,214 @@ pub enum Error {
     #[error("Failed to get Java VM instance")]
     GetJvmInstance(#[source] jnix::jni::errors::Error),
 
+    #[error("Failed to initialize logging: {0}")]
+    InitializeLogging(String),
+
     #[error("Failed to initialize the mullvad daemon")]
     InitializeDaemon(#[source] mullvad_daemon::Error),
 
-    #[error("Failed to spawn the tokio runtime")]
-    InitializeTokioRuntime(#[source] io::Error),
+    #[error("Failed to init Tokio runtime")]
+    InitTokio(#[source] io::Error),
 
     #[error("Failed to spawn the management interface")]
     SpawnManagementInterface(#[source] mullvad_daemon::management_interface::Error),
 }
 
+/// Throw a Java exception and return if `result` is an error
+macro_rules! ok_or_throw {
+    ($env:expr, $result:expr) => {{
+        match $result {
+            Ok(val) => val,
+            Err(err) => {
+                let env = $env;
+                env.throw(err.to_string())
+                    .expect("Failed to throw exception");
+                return;
+            }
+        }
+    }};
+}
+
+#[derive(Debug)]
+struct DaemonContext {
+    runtime: tokio::runtime::Runtime,
+    daemon_command_tx: DaemonCommandSender,
+    running_daemon: tokio::task::JoinHandle<()>,
+}
+
+/// Spawn Mullvad daemon. There can only be a single instance, which must be shut down using
+/// `MullvadDaemon.shutdown`. On success, nothing is returned. On error, an exception is thrown.
 #[no_mangle]
-#[allow(non_snake_case)]
 pub extern "system" fn Java_net_mullvad_mullvadvpn_service_MullvadDaemon_initialize(
     env: JNIEnv<'_>,
-    this: JObject<'_>,
-    vpnService: JObject<'_>,
-    rpcSocketPath: JObject<'_>,
-    filesDirectory: JObject<'_>,
-    cacheDirectory: JObject<'_>,
-    apiEndpoint: JObject<'_>,
+    _class: JClass<'_>,
+    vpn_service: JObject<'_>,
+    rpc_socket_path: JObject<'_>,
+    files_directory: JObject<'_>,
+    cache_directory: JObject<'_>,
+    api_endpoint: JObject<'_>,
 ) {
-    let env = JnixEnv::from(env);
-    let rpc_socket = PathBuf::from(String::from_java(&env, rpcSocketPath));
-    let files_dir = PathBuf::from(String::from_java(&env, filesDirectory));
-    let cache_dir = PathBuf::from(String::from_java(&env, cacheDirectory));
+    let mut ctx = DAEMON_CONTEXT.lock().unwrap();
+    assert!(ctx.is_none(), "multiple calls to MullvadDaemon.initialize");
 
-    let api_endpoint = if !apiEndpoint.is_null() {
-        #[cfg(feature = "api-override")]
-        {
-            Some(api_endpoint_from_java(&env, apiEndpoint))
-        }
-        #[cfg(not(feature = "api-override"))]
-        {
-            log::warn!("apiEndpoint will be ignored since 'api-override' is not enabled");
-            None
-        }
-    } else {
-        None
-    };
-
-    match start_logging(&files_dir) {
-        Ok(()) => {
-            version::log_version();
+    let env = JnixEnv::from(env);
 
-            LOAD_CLASSES.call_once(|| env.preload_classes(classes::CLASSES.iter().cloned()));
+    LOAD_CLASSES.call_once(|| env.preload_classes(classes::CLASSES.iter().cloned()));
 
-            if let Err(error) = initialize(
-                &env,
-                &this,
-                &vpnService,
-                rpc_socket,
-                files_dir,
-                cache_dir,
-                api_endpoint,
-            ) {
-                log::error!("{}", error.display_chain());
-            }
-        }
-        Err(message) => env
-            .throw(message.as_str())
-            .expect("Failed to throw exception"),
-    }
-}
+    let rpc_socket = pathbuf_from_java(&env, rpc_socket_path);
+    let files_dir = pathbuf_from_java(&env, files_directory);
+    let cache_dir = pathbuf_from_java(&env, cache_directory);
 
-#[cfg(feature = "api-override")]
-fn api_endpoint_from_java(env: &JnixEnv<'_>, object: JObject<'_>) -> mullvad_api::ApiEndpoint {
-    let mut endpoint = mullvad_api::ApiEndpoint::from_env_vars();
+    let android_context = ok_or_throw!(&env, create_android_context(&env, vpn_service));
 
-    let address = env
-        .call_method(object, "component1", "()Ljava/net/InetSocketAddress;", &[])
-        .expect("missing ApiEndpoint.address")
-        .l()
-        .expect("ApiEndpoint.address is not an InetSocketAddress");
+    let api_endpoint = api::api_endpoint_from_java(&env, api_endpoint);
 
-    endpoint.address = Some(
-        try_socketaddr_from_java(env, address).expect("received unresolved InetSocketAddress"),
+    let daemon = ok_or_throw!(
+        &env,
+        start(
+            android_context,
+            rpc_socket,
+            files_dir,
+            cache_dir,
+            api_endpoint,
+        )
     );
-    endpoint.host = try_hostname_from_java(env, address);
-    endpoint.disable_address_cache = env
-        .call_method(object, "component2", "()Z", &[])
-        .expect("missing ApiEndpoint.disableAddressCache")
-        .z()
-        .expect("ApiEndpoint.disableAddressCache is not a bool");
-    endpoint.disable_tls = env
-        .call_method(object, "component3", "()Z", &[])
-        .expect("missing ApiEndpoint.disableTls")
-        .z()
-        .expect("ApiEndpoint.disableTls is not a bool");
 
-    endpoint
+    *ctx = Some(daemon);
 }
 
-/// Converts InetSocketAddress to a SocketAddr. Return `None` if the
-/// hostname is unresolved.
-#[cfg(feature = "api-override")]
-fn try_socketaddr_from_java(env: &JnixEnv<'_>, address: JObject<'_>) -> Option<SocketAddr> {
-    let class = env.get_class("java/net/InetSocketAddress");
-
-    let method_id = env
-        .get_method_id(&class, "getAddress", "()Ljava/net/InetAddress;")
-        .expect("Failed to get method ID for InetSocketAddress.getAddress()");
-    let return_type = JavaType::Object("java/net/InetAddress".to_owned());
-
-    let ip_addr = env
-        .call_method_unchecked(address, method_id, return_type, &[])
-        .expect("Failed to call InetSocketAddress.getAddress()")
-        .l()
-        .expect("Call to InetSocketAddress.getAddress() did not return an object");
-
-    if ip_addr.is_null() {
-        return None;
+/// Shut down Mullvad daemon that was initialized using `MullvadDaemon.initialize`.
+#[no_mangle]
+#[allow(non_snake_case)]
+pub extern "system" fn Java_net_mullvad_mullvadvpn_service_MullvadDaemon_shutdown(
+    _: JNIEnv<'_>,
+    _class: JClass<'_>,
+) {
+    if let Some(context) = DAEMON_CONTEXT.lock().unwrap().take() {
+        _ = context.daemon_command_tx.shutdown();
+        _ = context.runtime.block_on(context.running_daemon);
     }
-
-    let method_id = env
-        .get_method_id(&class, "getPort", "()I")
-        .expect("Failed to get method ID for InetSocketAddress.getPort()");
-    let return_type = JavaType::Primitive(Primitive::Int);
-
-    let port = env
-        .call_method_unchecked(address, method_id, return_type, &[])
-        .expect("Failed to call InetSocketAddress.getPort()")
-        .i()
-        .expect("Call to InetSocketAddress.getPort() did not return an int");
-
-    Some(SocketAddr::new(
-        IpAddr::from_java(env, ip_addr),
-        u16::try_from(port).expect("invalid port"),
-    ))
 }
 
-/// Returns the hostname for an InetSocketAddress. This may be an IP address converted to
-/// a string.
-#[cfg(feature = "api-override")]
-fn try_hostname_from_java(env: &JnixEnv<'_>, address: JObject<'_>) -> Option<String> {
-    let class = env.get_class("java/net/InetSocketAddress");
-
-    let method_id = env
-        .get_method_id(&class, "getHostString", "()Ljava/lang/String;")
-        .expect("Failed to get method ID for InetSocketAddress.getHostString()");
-    let return_type = JavaType::Object("java/lang/String".to_owned());
-
-    let hostname = env
-        .call_method_unchecked(address, method_id, return_type, &[])
-        .expect("Failed to call InetSocketAddress.getHostString()")
-        .l()
-        .expect("Call to InetSocketAddress.getHostString() did not return an object");
+fn start(
+    android_context: AndroidContext,
+    rpc_socket: PathBuf,
+    files_dir: PathBuf,
+    cache_dir: PathBuf,
+    api_endpoint: Option<mullvad_api::ApiEndpoint>,
+) -> Result<DaemonContext, Error> {
+    start_logging(&files_dir).map_err(Error::InitializeLogging)?;
+    version::log_version();
 
-    if hostname.is_null() {
-        return None;
+    #[cfg(feature = "api-override")]
+    if let Some(api_endpoint) = api_endpoint {
+        log::debug!("Overriding API endpoint: {api_endpoint:?}");
+        if mullvad_api::API.override_init(api_endpoint).is_err() {
+            log::warn!("Ignoring API settings (already initialized)");
+        }
     }
-
-    Some(String::from_java(env, hostname))
-}
-
-fn start_logging(log_dir: &Path) -> Result<(), String> {
-    unsafe {
-        LOG_START.call_once(|| LOG_INIT_RESULT = Some(initialize_logging(log_dir)));
-        LOG_INIT_RESULT
-            .clone()
-            .expect("Logging not properly initialized")
+    #[cfg(not(feature = "api-override"))]
+    if api_endpoint.is_some() {
+        log::warn!("api_endpoint will be ignored since 'api-override' is not enabled");
     }
-}
-
-fn initialize_logging(log_dir: &Path) -> Result<(), String> {
-    let log_file = log_dir.join(LOG_FILENAME);
 
-    logging::init_logger(log::LevelFilter::Debug, Some(&log_file), true)
-        .map_err(|error| error.display_chain_with_msg("Failed to start logger"))?;
-    exception_logging::enable();
-    log_panics::init();
-
-    Ok(())
+    spawn_daemon(android_context, rpc_socket, files_dir, cache_dir)
 }
 
-fn initialize(
-    env: &JnixEnv<'_>,
-    this: &JObject<'_>,
-    vpn_service: &JObject<'_>,
+fn spawn_daemon(
+    android_context: AndroidContext,
     rpc_socket: PathBuf,
     files_dir: PathBuf,
     cache_dir: PathBuf,
-    api_endpoint: Option<mullvad_api::ApiEndpoint>,
-) -> Result<(), Error> {
-    let android_context = create_android_context(env, *vpn_service)?;
+) -> Result<DaemonContext, Error> {
     let daemon_command_channel = DaemonCommandChannel::new();
-    let daemon_interface = Box::new(DaemonInterface::new(daemon_command_channel.sender()));
+    let daemon_command_tx = daemon_command_channel.sender();
+
+    let runtime = new_multi_thread().build().map_err(Error::InitTokio)?;
 
-    spawn_daemon(
-        env,
-        this,
+    let running_daemon = runtime.block_on(spawn_daemon_inner(
         rpc_socket,
         files_dir,
         cache_dir,
-        api_endpoint,
         daemon_command_channel,
         android_context,
-    )?;
+    ))?;
 
-    set_daemon_interface_address(env, this, Box::into_raw(daemon_interface) as jlong);
-
-    Ok(())
-}
-
-fn create_android_context(
-    env: &JnixEnv<'_>,
-    vpn_service: JObject<'_>,
-) -> Result<AndroidContext, Error> {
-    Ok(AndroidContext {
-        jvm: Arc::new(env.get_java_vm().map_err(Error::GetJvmInstance)?),
-        vpn_service: env
-            .new_global_ref(vpn_service)
-            .map_err(Error::CreateGlobalReference)?,
+    Ok(DaemonContext {
+        runtime,
+        daemon_command_tx,
+        running_daemon,
     })
 }
 
-#[allow(clippy::too_many_arguments)]
-fn spawn_daemon(
-    env: &JnixEnv<'_>,
-    this: &JObject<'_>,
+async fn spawn_daemon_inner(
     rpc_socket: PathBuf,
     files_dir: PathBuf,
     cache_dir: PathBuf,
-    #[cfg_attr(not(feature = "api-override"), allow(unused_variables))] api_endpoint: Option<
-        mullvad_api::ApiEndpoint,
-    >,
     command_channel: DaemonCommandChannel,
     android_context: AndroidContext,
-) -> Result<(), Error> {
-    let daemon_object = env
-        .new_global_ref(*this)
-        .map_err(Error::CreateGlobalReference)?;
-    let (tx, rx) = mpsc::channel();
+) -> Result<tokio::task::JoinHandle<()>, Error> {
+    cleanup_old_rpc_socket(&rpc_socket).await;
 
-    let runtime = new_multi_thread()
-        .build()
-        .map_err(Error::InitializeTokioRuntime)?;
+    let event_listener = ManagementInterfaceServer::start(command_channel.sender(), &rpc_socket)
+        .map_err(Error::SpawnManagementInterface)?;
 
-    thread::spawn(move || {
-        let jvm = android_context.jvm.clone();
-        let running_instances = DAEMON_INSTANCE_COUNT.fetch_add(1, Ordering::AcqRel);
+    log::info!("Management interface listening on {}", rpc_socket.display());
 
-        if running_instances != 0 {
-            log::error!(
-                "It seems that there are already {} instances of the Mullvad daemon running",
-                running_instances
-            );
-        }
-
-        #[cfg(feature = "api-override")]
-        if let Some(api_endpoint) = api_endpoint {
-            log::debug!("Overriding API endpoint: {api_endpoint:?}");
-            if mullvad_api::API.override_init(api_endpoint).is_err() {
-                log::warn!("Ignoring API settings (already initialized)");
-            }
-        }
-
-        runtime.block_on(cleanup_old_rpc_socket(&rpc_socket));
-
-        let event_listener = match runtime
-            .block_on(async { spawn_management_interface(command_channel.sender(), &rpc_socket) })
-        {
-            Ok(event_listener) => event_listener,
-            Err(error) => {
-                let _ = tx.send(Err(error));
-                return;
-            }
-        };
-
-        let daemon = runtime.block_on(Daemon::start(
-            Some(files_dir.clone()),
-            files_dir.clone(),
-            files_dir,
-            cache_dir,
-            event_listener,
-            command_channel,
-            android_context,
-        ));
-
-        DAEMON_INSTANCE_COUNT.fetch_sub(1, Ordering::AcqRel);
-
-        match daemon {
-            Ok(daemon) => {
-                let _ = tx.send(Ok(()));
-                match runtime.block_on(daemon.run()) {
-                    Ok(()) => log::info!("Mullvad daemon has stopped"),
-                    Err(error) => log::error!("{}", error.display_chain()),
-                }
-            }
-            Err(error) => {
-                let _ = tx.send(Err(Error::InitializeDaemon(error)));
-            }
-        }
-
-        notify_daemon_stopped(jvm, daemon_object);
-    });
-
-    rx.recv().unwrap()
-}
-
-use mullvad_daemon::{
-    management_interface::{ManagementInterfaceEventBroadcaster, ManagementInterfaceServer},
-    DaemonCommandSender,
-};
+    let daemon = Daemon::start(
+        Some(files_dir.clone()),
+        files_dir.clone(),
+        files_dir,
+        cache_dir,
+        event_listener,
+        command_channel,
+        android_context,
+    )
+    .await
+    .map_err(Error::InitializeDaemon)?;
 
-fn spawn_management_interface(
-    command_sender: DaemonCommandSender,
-    rpc_socket_path: impl AsRef<Path>,
-) -> Result<ManagementInterfaceEventBroadcaster, Error> {
-    let event_broadcaster = ManagementInterfaceServer::start(command_sender, &rpc_socket_path)
-        .map_err(|error| {
-            log::error!(
+    let running_daemon = tokio::spawn(async move {
+        match daemon.run().await {
+            Ok(()) => log::info!("Mullvad daemon has stopped"),
+            Err(error) => log::error!(
                 "{}",
-                error.display_chain_with_msg("Unable to start management interface server")
-            );
-            Error::SpawnManagementInterface(error)
-        })?;
-
-    log::info!(
-        "Management interface listening on {}",
-        rpc_socket_path.as_ref().display()
-    );
-
-    Ok(event_broadcaster)
-}
-
-fn notify_daemon_stopped(jvm: Arc<JavaVM>, daemon_object: GlobalRef) {
-    match jvm.attach_current_thread_as_daemon() {
-        Ok(env) => {
-            let env = JnixEnv::from(env);
-            let class = env.get_class("net/mullvad/mullvadvpn/service/MullvadDaemon");
-            let object = daemon_object.as_obj();
-            let method_id = env
-                .get_method_id(&class, "notifyDaemonStopped", "()V")
-                .expect("Failed to get method ID for MullvadDaemon.notifyDaemonStopped");
-            let return_type = JavaType::Primitive(Primitive::Void);
-
-            let result = env.call_method_unchecked(object, method_id, return_type, &[]);
-
-            match result {
-                Ok(JValue::Void) => {}
-                Ok(value) => panic!(
-                    "Unexpected return value from MullvadDaemon.notifyDaemonStopped: {:?}",
-                    value
-                ),
-                Err(error) => panic!(
-                    "{}",
-                    error
-                        .display_chain_with_msg("Failed to call MullvadDaemon.notifyDaemonStopped")
-                ),
-            }
+                error.display_chain_with_msg("Mullvad daemon exited with an error")
+            ),
         }
-        Err(error) => log::error!(
-            "{}",
-            error.display_chain_with_msg("Failed to notify that the daemon stopped")
-        ),
-    }
-}
-
-fn set_daemon_interface_address(env: &JnixEnv<'_>, this: &JObject<'_>, address: jlong) {
-    let class = env.get_class("net/mullvad/mullvadvpn/service/MullvadDaemon");
-    let method_id = env
-        .get_method_id(&class, "setDaemonInterfaceAddress", "(J)V")
-        .expect("Failed to get method ID for MullvadDaemon.setDaemonInterfaceAddress");
-    let return_type = JavaType::Primitive(Primitive::Void);
-
-    let result = env.call_method_unchecked(*this, method_id, return_type, &[JValue::Long(address)]);
+    });
 
-    match result {
-        Ok(JValue::Void) => {}
-        Ok(value) => panic!(
-            "Unexpected return value from MullvadDaemon.setDaemonInterfaceAddress: {:?}",
-            value
-        ),
-        Err(error) => panic!(
-            "{}",
-            error.display_chain_with_msg("Failed to call MullvadDaemon.setDaemonInterfaceAddress")
-        ),
-    }
+    Ok(running_daemon)
 }
 
-fn get_daemon_interface_address(env: &JnixEnv<'_>, this: &JObject<'_>) -> *mut DaemonInterface {
-    let class = env.get_class("net/mullvad/mullvadvpn/service/MullvadDaemon");
-    let method_id = env
-        .get_method_id(&class, "getDaemonInterfaceAddress", "()J")
-        .expect("Failed to get method ID for MullvadDaemon.getDaemonInterfaceAddress");
-    let return_type = JavaType::Primitive(Primitive::Long);
-
-    let result = env.call_method_unchecked(*this, method_id, return_type, &[]);
-
-    match result {
-        Ok(JValue::Long(address)) => address as *mut DaemonInterface,
-        Ok(value) => panic!(
-            "Invalid return value from MullvadDaemon.getDaemonInterfaceAddress: {:?}",
-            value
-        ),
-        Err(error) => panic!(
-            "{}",
-            error.display_chain_with_msg("Failed to call MullvadDaemon.getDaemonInterfaceAddress")
-        ),
-    }
+fn start_logging(log_dir: &Path) -> Result<(), String> {
+    static LOGGER_RESULT: OnceLock<Result<(), String>> = OnceLock::new();
+    LOGGER_RESULT
+        .get_or_init(|| start_logging_inner(log_dir).map_err(|e| e.display_chain()))
+        .to_owned()
 }
 
-#[no_mangle]
-#[allow(non_snake_case)]
-pub extern "system" fn Java_net_mullvad_mullvadvpn_service_MullvadDaemon_deinitialize(
-    env: JNIEnv<'_>,
-    this: JObject<'_>,
-) {
-    let env = JnixEnv::from(env);
-    let daemon_interface_address = get_daemon_interface_address(&env, &this);
+fn start_logging_inner(log_dir: &Path) -> Result<(), logging::Error> {
+    let log_file = log_dir.join(LOG_FILENAME);
 
-    set_daemon_interface_address(&env, &this, 0);
+    logging::init_logger(log::LevelFilter::Debug, Some(&log_file), true)?;
+    exception_logging::enable();
+    log_panics::init();
 
-    if !daemon_interface_address.is_null() {
-        let _ = unsafe { Box::from_raw(daemon_interface_address) };
-    }
+    Ok(())
 }
 
-/// # Safety
-///
-/// `address` must either be zero or a valid pointer to a `DaemonInterface` instance.
-/// This function has no concept of lifetimes, so the caller must ensure that the
-/// pointed to `DaemonInterface` is valid for at least as long as the return value
-/// of this function is still alive.
-unsafe fn get_daemon_interface(address: jlong) -> Option<&'static mut DaemonInterface> {
-    let pointer = address as *mut DaemonInterface;
-
-    if !pointer.is_null() {
-        Some(&mut *pointer)
-    } else {
-        log::error!("Attempt to get daemon interface while it is null");
-        None
-    }
+fn create_android_context(
+    env: &JnixEnv<'_>,
+    vpn_service: JObject<'_>,
+) -> Result<AndroidContext, Error> {
+    Ok(AndroidContext {
+        jvm: Arc::new(env.get_java_vm().map_err(Error::GetJvmInstance)?),
+        vpn_service: env
+            .new_global_ref(vpn_service)
+            .map_err(Error::CreateGlobalReference)?,
+    })
 }
 
-#[no_mangle]
-#[allow(non_snake_case)]
-pub extern "system" fn Java_net_mullvad_mullvadvpn_service_MullvadDaemon_shutdown(
-    _: JNIEnv<'_>,
-    _: JObject<'_>,
-    daemon_interface_address: jlong,
-) {
-    // SAFETY: The address points to an instance valid for the duration of this function call
-    if let Some(daemon_interface) = unsafe { get_daemon_interface(daemon_interface_address) } {
-        if let Err(error) = daemon_interface.shutdown() {
-            log::error!(
-                "{}",
-                error.display_chain_with_msg("Failed to shutdown daemon thread")
-            );
-        }
-    }
+fn pathbuf_from_java(env: &JnixEnv<'_>, path: JObject<'_>) -> PathBuf {
+    PathBuf::from(String::from_java(env, path))
 }