fix: disconnect VPN when sleep

apps/gpservice/Cargo.toml

@@ -9,7 +9,7 @@ gpapi = { path = "../../crates/gpapi" }
 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

@@ -30,7 +30,8 @@ struct Cli {
 
 impl Cli {
   async fn run(&mut self, redaction: Arc<Redaction>) -> anyhow::Result<()> {
-    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");
@@ -48,9 +49,17 @@ 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();
 
+    #[cfg(unix)]
+    {
+      let vpn_ctx = vpn_task.context();
+      let ws_ctx = ws_server.context();
+
+      tokio::spawn(async move { signals::handle_signals(vpn_ctx, ws_ctx).await });
+    }
+
     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 });
 
@@ -74,15 +83,15 @@ impl Cli {
     }
 
     tokio::select! {
-        _ = shutdown_signal() => {
-            info!("Shutdown signal received");
-        }
-        _ = shutdown_rx.recv() => {
-            info!("Shutdown request received, shutting down");
-        }
+      _ = shutdown_signal() => {
+        info!("Shutdown signal received");
+      }
+      _ = shutdown_rx.recv() => {
+        info!("Shutdown request received, shutting down");
+      }
     }
 
-    vpn_task_token.cancel();
+    vpn_task_cancel_token.cancel();
     let _ = tokio::join!(vpn_task_handle, ws_server_handle);
 
     lock_file.unlock()?;
@@ -125,6 +134,54 @@ fn init_logger() -> Arc<Redaction> {
   redaction
 }
 
+#[cfg(unix)]
+mod signals {
+  use std::sync::Arc;
+
+  use log::{info, warn};
+
+  use crate::vpn_task::VpnTaskContext;
+  use crate::ws_server::WsServerContext;
+
+  const DISCONNECTED_PID_FILE: &str = "/tmp/gpservice_disconnected.pid";
+
+  pub async fn handle_signals(vpn_ctx: Arc<VpnTaskContext>, ws_ctx: Arc<WsServerContext>) {
+    use gpapi::service::event::WsEvent;
+    use tokio::signal::unix::{signal, Signal, SignalKind};
+
+    let (mut user_sig1, mut user_sig2) = match || -> anyhow::Result<(Signal, Signal)> {
+      let user_sig1 = signal(SignalKind::user_defined1())?;
+      let user_sig2 = signal(SignalKind::user_defined2())?;
+      Ok((user_sig1, user_sig2))
+    }() {
+      Ok(signals) => signals,
+      Err(err) => {
+        warn!("Failed to create signal: {}", err);
+        return;
+      }
+    };
+
+    loop {
+      tokio::select! {
+        _ = user_sig1.recv() => {
+          info!("Received SIGUSR1 signal");
+          if vpn_ctx.disconnect().await {
+            // Write the PID to a dedicated file to indicate that the VPN task is disconnected via SIGUSR1
+            let pid = std::process::id();
+            if let Err(err) = tokio::fs::write(DISCONNECTED_PID_FILE, pid.to_string()).await {
+              warn!("Failed to write PID to file: {}", err);
+            }
+          }
+        }
+        _ = user_sig2.recv() => {
+          info!("Received SIGUSR2 signal");
+          ws_ctx.send_event(WsEvent::ResumeConnection).await;
+        }
+      }
+    }
+  }
+}
+
 async fn launch_gui(envs: Option<HashMap<String, String>>, api_key: Vec<u8>, mut minimized: bool) {
   loop {
     let gui_launcher = GuiLauncher::new(env!("CARGO_PKG_VERSION"), &api_key)

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,
@@ -137,7 +136,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

@@ -87,7 +87,7 @@ impl VpnTaskContext {
     });
   }
 
-  pub async fn disconnect(&self) {
+  pub async fn disconnect(&self) -> bool {
     if let Some(disconnect_rx) = self.disconnect_rx.write().await.take() {
       info!("Disconnecting VPN...");
       if let Some(vpn) = self.vpn_handle.read().await.as_ref() {
@@ -98,9 +98,13 @@ impl VpnTaskContext {
       // Wait for the VPN to be disconnected
       disconnect_rx.await.ok();
       info!("VPN disconnected");
+
+      true
     } else {
       info!("VPN is not connected, skip disconnect");
       self.vpn_state_tx.send(VpnState::Disconnected).ok();
+
+      false
     }
   }
 }
@@ -143,6 +147,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

@@ -113,6 +113,10 @@ impl WsServer {
     }
   }
 
+  pub fn context(&self) -> Arc<WsServerContext> {
+    Arc::clone(&self.ctx)
+  }
+
   pub fn cancel_token(&self) -> CancellationToken {
     self.cancel_token.clone()
   }
@@ -124,7 +128,7 @@ impl WsServer {
         warn!("Failed to start WS server: {}", err);
         let _ = shutdown_tx.send(()).await;
         return;
-      },
+      }
     };
 
     tokio::select! {
@@ -149,7 +153,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)
   }