Implement safe cross-platform daemon shutdown

2 files changed, 84 insertions, 8 deletions

diff --git a/mullvad-daemon/tests/common/mod.rs b/mullvad-daemon/tests/common/mod.rs
index 45b48c3db7..8259c63795 100644
--- a/mullvad-daemon/tests/common/mod.rs
+++ b/mullvad-daemon/tests/common/mod.rs
@@ -1,5 +1,8 @@
 #![allow(dead_code)]
 
+#[cfg(unix)]
+extern crate libc;
+
 use std::fs::File;
 use std::io::{BufRead, BufReader, Write};
 use std::path::{Path, PathBuf};
@@ -9,6 +12,8 @@ use std::time::Duration;
 
 use duct;
 use os_pipe::{pipe, PipeReader};
+use serde::{Deserialize, Serialize};
+use talpid_ipc::WsIpcClient;
 
 pub use self::platform_specific::*;
 
@@ -45,8 +50,44 @@ fn prepare_relay_list<T: AsRef<Path>>(path: T) {
     }
 }
 
+pub struct DaemonRpcClient {
+    address: String,
+}
+
+impl DaemonRpcClient {
+    fn new() -> Result<Self, String> {
+        let rpc_file = File::open(rpc_file_path())
+            .map_err(|error| format!("failed to open RPC address file: {}", error))?;
+        let reader = BufReader::new(rpc_file);
+        let mut lines = reader.lines();
+        let address = lines
+            .next()
+            .ok_or("RPC address file is empty".to_string())?
+            .map_err(|error| format!("failed to read address from RPC address file: {}", error))?;
+
+        Ok(DaemonRpcClient { address })
+    }
+
+    pub fn shutdown(&self) -> Result<(), String> {
+        self.call("shutdown", &[] as &[u8; 0])
+    }
+
+    pub fn call<A, O>(&self, method: &str, args: &A) -> Result<O, String>
+    where
+        A: Serialize,
+        O: for<'de> Deserialize<'de>,
+    {
+        let mut rpc_client = WsIpcClient::new(self.address.clone())
+            .map_err(|error| format!("unable to create RPC client: {}", error))?;
+
+        rpc_client
+            .call(method, args)
+            .map_err(|error| format!("RPC request failed: {}", error))
+    }
+}
+
 pub struct DaemonRunner {
-    process: duct::Handle,
+    process: Option<duct::Handle>,
     output: Arc<Mutex<BufReader<PipeReader>>>,
 }
 
@@ -58,6 +99,7 @@ impl DaemonRunner {
         let process = cmd!(
             DAEMON_EXECUTABLE_PATH,
             "-v",
+            "--disable-rpc-auth",
             "--resource-dir",
             "dist-assets"
         ).dir("..")
@@ -67,7 +109,7 @@ impl DaemonRunner {
             .expect("failed to start daemon");
 
         DaemonRunner {
-            process,
+            process: Some(process),
             output: Arc::new(Mutex::new(BufReader::new(reader))),
         }
     }
@@ -99,14 +141,46 @@ impl DaemonRunner {
                 .expect("failed to read line from daemon stdout");
         }
     }
+
+    #[cfg(unix)]
+    fn request_clean_shutdown(&mut self, process: &mut duct::Handle) -> bool {
+        use duct::unix::HandleExt;
+
+        process.send_signal(libc::SIGTERM).is_ok()
+    }
+
+    #[cfg(not(unix))]
+    fn request_clean_shutdown(&mut self, _: &mut duct::Handle) -> bool {
+        if let Ok(rpc_client) = DaemonRpcClient::new() {
+            rpc_client.shutdown().is_ok()
+        } else {
+            false
+        }
+    }
 }
 
-#[cfg(unix)]
 impl Drop for DaemonRunner {
     fn drop(&mut self) {
-        use duct::unix::HandleExt;
-        use libc;
+        if let Some(mut process) = self.process.take() {
+            if self.request_clean_shutdown(&mut process) {
+                let process = Arc::new(process);
+                let wait_handle = process.clone();
+                let (finished_tx, finished_rx) = mpsc::channel();
 
-        let _ = self.process.send_signal(libc::SIGTERM);
+                thread::spawn(move || finished_tx.send(wait_handle.wait().map(|_| ())).unwrap());
+
+                let has_finished = finished_rx
+                    .recv_timeout(Duration::from_secs(5))
+                    .map_err(|_| ())
+                    .and_then(|result| result.map_err(|_| ()))
+                    .is_ok();
+
+                if !has_finished {
+                    process.kill().unwrap();
+                }
+            } else {
+                process.kill().unwrap();
+            }
+        }
     }
 }
diff --git a/mullvad-daemon/tests/startup.rs b/mullvad-daemon/tests/startup.rs
index 2f93bd0f99..71ba7a752b 100644
--- a/mullvad-daemon/tests/startup.rs
+++ b/mullvad-daemon/tests/startup.rs
@@ -1,8 +1,8 @@
 #[macro_use]
 extern crate duct;
-#[cfg(unix)]
-extern crate libc;
 extern crate os_pipe;
+extern crate serde;
+extern crate talpid_ipc;
 
 mod common;
 
@@ -39,6 +39,8 @@ fn rpc_info_file_permissions() {
 
 #[cfg(unix)]
 mod platform_specific {
+    extern crate libc;
+
     use super::*;
     use std::os::unix::fs::MetadataExt;