fix: disconnect VPN when sleep/shutdown

apps/gpclient/Cargo.toml

@@ -16,7 +16,7 @@ clap.workspace = true
 env_logger.workspace = true
 inquire = "0.7"
 log.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["rt-multi-thread"] }
 sysinfo.workspace = true
 serde_json.workspace = true
 whoami.workspace = true

apps/gpclient/src/disconnect.rs

@@ -1,6 +1,7 @@
 use crate::GP_CLIENT_LOCK_FILE;
+use gpapi::GP_SERVICE_LOCK_FILE;
 use log::{info, warn};
-use std::fs;
+use std::{fs, str::FromStr, thread, time::Duration};
 use sysinfo::{Pid, Signal, System};
 
 pub(crate) struct DisconnectHandler;
@@ -11,21 +12,40 @@ impl DisconnectHandler {
   }
 
   pub(crate) fn handle(&self) -> anyhow::Result<()> {
-    if fs::metadata(GP_CLIENT_LOCK_FILE).is_err() {
-      warn!("PID file not found, maybe the client is not running");
-      return Ok(());
-    }
+    if let Ok(c) = fs::read_to_string(GP_CLIENT_LOCK_FILE) {
+      send_signal(c.trim(), Signal::Interrupt).unwrap_or_else(|err| {
+        warn!("Failed to send signal to client: {}", err);
+      });
+    };
 
-    let pid = fs::read_to_string(GP_CLIENT_LOCK_FILE)?;
-    let pid = pid.trim().parse::<usize>()?;
-    let s = System::new_all();
+    if let Ok(c) = fs::read_to_string(GP_SERVICE_LOCK_FILE) {
+      c.split(':').next().map_or_else(
+        || info!("Failed to extract PID from: {}", c),
+        |pid| {
+          send_signal(pid, Signal::User1).unwrap_or_else(|err| {
+            warn!("Failed to send signal to service: {}", err);
+          });
+        },
+      );
+    };
+
+    // sleep for 3 seconds to give the client and service time to shut down
+    thread::sleep(Duration::from_secs(3));
 
-    if let Some(process) = s.process(Pid::from(pid)) {
-      info!("Found process {}, killing...", pid);
-      if process.kill_with(Signal::Interrupt).is_none() {
-        warn!("Failed to kill process {}", pid);
-      }
-    }
     Ok(())
   }
 }
+
+fn send_signal(pid: &str, signal: Signal) -> anyhow::Result<()> {
+  let s = System::new_all();
+  let pid = Pid::from_str(pid)?;
+
+  if let Some(process) = s.process(pid) {
+    info!("Found process {}, sending signal...", pid);
+
+    if process.kill_with(signal).is_none() {
+      warn!("Failed to kill process {}", pid);
+    }
+  }
+  Ok(())
+}

apps/gpservice/Cargo.toml

@@ -9,7 +9,7 @@ gpapi = { path = "../../crates/gpapi", features = ["clap", "logger"] }
 openconnect = { path = "../../crates/openconnect" }
 clap.workspace = true
 anyhow.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["rt-multi-thread"] }
 tokio-util.workspace = true
 axum = { workspace = true, features = ["ws"] }
 futures.workspace = true

apps/gpservice/src/cli.rs

@@ -38,7 +38,8 @@ impl Cli {
     let redaction = self.init_logger();
     info!("gpservice started: {}", VERSION);
 
-    let lock_file = Arc::new(LockFile::new(GP_SERVICE_LOCK_FILE));
+    let pid = std::process::id();
+    let lock_file = Arc::new(LockFile::new(GP_SERVICE_LOCK_FILE, pid));
 
     if lock_file.check_health().await {
       bail!("Another instance of the service is already running");
@@ -56,9 +57,10 @@ impl Cli {
 
     let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(4);
     let shutdown_tx_clone = shutdown_tx.clone();
-    let vpn_task_token = vpn_task.cancel_token();
+    let vpn_task_cancel_token = vpn_task.cancel_token();
     let server_token = ws_server.cancel_token();
 
+    let vpn_cxt = vpn_task.context();
     let vpn_task_handle = tokio::spawn(async move { vpn_task.start(server_token).await });
     let ws_server_handle = tokio::spawn(async move { ws_server.start(shutdown_tx_clone).await });
 
@@ -81,6 +83,26 @@ impl Cli {
       });
     }
 
+    // Handle SIGUSR1 signal to disconnect the VPN
+    #[cfg(unix)]
+    tokio::spawn(async move {
+      use tokio::signal::unix::{signal, SignalKind};
+
+      let mut sig = match signal(SignalKind::user_defined1()) {
+        Ok(sig) => sig,
+        Err(err) => {
+          warn!("Failed to create signal: {}", err);
+          return;
+        }
+      };
+
+      loop {
+        sig.recv().await;
+        info!("Received SIGUSR1 signal");
+        vpn_cxt.disconnect().await;
+      }
+    });
+
     tokio::select! {
         _ = shutdown_signal() => {
             info!("Shutdown signal received");
@@ -90,7 +112,7 @@ impl Cli {
         }
     }
 
-    vpn_task_token.cancel();
+    vpn_task_cancel_token.cancel();
     let _ = tokio::join!(vpn_task_handle, ws_server_handle);
 
     lock_file.unlock()?;

apps/gpservice/src/handlers.rs

@@ -1,5 +1,4 @@
 use std::{
-  borrow::Cow,
   fs::{File, Permissions},
   io::BufReader,
   ops::ControlFlow,
@@ -12,7 +11,7 @@ use anyhow::bail;
 use axum::{
   body::Bytes,
   extract::{
-    ws::{self, CloseFrame, Message, WebSocket},
+    ws::{self, CloseFrame, Message, Utf8Bytes, WebSocket},
     State, WebSocketUpgrade,
   },
   http::StatusCode,
@@ -133,7 +132,7 @@ async fn handle_socket(mut socket: WebSocket, ctx: Arc<WsServerContext>) {
 
     let close_msg = Message::Close(Some(CloseFrame {
       code: ws::close_code::NORMAL,
-      reason: Cow::from("Goodbye"),
+      reason: Utf8Bytes::from("Goodbye"),
     }));
 
     if let Err(err) = sender.send(close_msg).await {

apps/gpservice/src/vpn_task.rs

@@ -146,6 +146,10 @@ impl VpnTask {
     server_cancel_token.cancel();
   }
 
+  pub fn context(&self) -> Arc<VpnTaskContext> {
+    return Arc::clone(&self.ctx);
+  }
+
   async fn recv(&mut self) {
     while let Some(req) = self.ws_req_rx.recv().await {
       tokio::spawn(process_ws_req(req, self.ctx.clone()));

apps/gpservice/src/ws_connection.rs

@@ -20,7 +20,7 @@ impl WsConnection {
 
   pub async fn send_event(&self, event: &WsEvent) -> anyhow::Result<()> {
     let encrypted = self.crypto.encrypt(event)?;
-    let msg = Message::Binary(encrypted);
+    let msg = Message::Binary(encrypted.into());
 
     self.tx.send(msg).await?;
 
@@ -29,7 +29,7 @@ impl WsConnection {
 
   pub fn recv_msg(&self, msg: Message) -> ControlFlow<(), WsRequest> {
     match msg {
-      Message::Binary(data) => match self.crypto.decrypt(data) {
+      Message::Binary(data) => match self.crypto.decrypt(data.into()) {
         Ok(ws_req) => ControlFlow::Continue(ws_req),
         Err(err) => {
           info!("Failed to decrypt message: {}", err);

apps/gpservice/src/ws_server.rs

@@ -124,7 +124,7 @@ impl WsServer {
         warn!("Failed to start WS server: {}", err);
         let _ = shutdown_tx.send(()).await;
         return;
-      },
+      }
     };
 
     tokio::select! {
@@ -149,7 +149,7 @@ impl WsServer {
 
     info!("WS server listening on port: {}", port);
 
-    self.lock_file.lock(port.to_string())?;
+    self.lock_file.lock(&port.to_string())?;
 
     Ok(listener)
   }