fix: disconnect VPN when sleep/shutdown

yuezk · Jan 12, 2025 · 9740231 · 9740231
1 parent ec85e85
commit 9740231
Show file tree

Hide file tree

Showing 18 changed files with 285 additions and 172 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,13 @@
 [workspace]
 resolver = "2"
 
-members = ["crates/*", "apps/gpclient", "apps/gpservice", "apps/gpauth", "apps/gpgui-helper/src-tauri"]
+members = [
+  "crates/*",
+  "apps/gpclient",
+  "apps/gpservice",
+  "apps/gpauth",
+  "apps/gpgui-helper/src-tauri",
+]
 
 [workspace.package]
 rust-version = "1.80"
@@ -31,11 +37,11 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sysinfo = "0.33"
 tempfile = "3.8"
-tokio = { version = "1", features = ["full"] }
+tokio = { version = "1" }
 tokio-util = "0.7"
 url = "2.4"
 urlencoding = "2.1.3"
-axum = "0.7"
+axum = "0.8"
 futures = "0.3"
 futures-util = "0.3"
 uzers = "0.12"
@@ -44,9 +50,9 @@ thiserror = "2"
 redact-engine = "0.1"
 compile-time = "0.2"
 serde_urlencoded = "0.7"
-md5="0.7"
-sha256="1"
-which="7"
+md5 = "0.7"
+sha256 = "1"
+which = "7"
 
 # Tauri dependencies
 tauri = { version = "2" }

diff --git a/Makefile b/Makefile
@@ -130,6 +130,10 @@ install:
 		install -Dm755 .build/gpgui/gpgui_*/gpgui $(DESTDIR)/usr/bin/gpgui; \
 	fi
 
+	# Install the sleep/shutdown scripts
+	install -Dm755 packaging/files/usr/lib/systemd/system-sleep/gpclient $(DESTDIR)/usr/lib/systemd/system-sleep/gpclient
+	install -Dm755 packaging/files/usr/lib/systemd/system-shutdown/gpclient $(DESTDIR)/usr/lib/systemd/system-shutdown/gpclient
+
 	install -Dm644 packaging/files/usr/share/applications/gpgui.desktop $(DESTDIR)/usr/share/applications/gpgui.desktop
 	install -Dm644 packaging/files/usr/share/icons/hicolor/scalable/apps/gpgui.svg $(DESTDIR)/usr/share/icons/hicolor/scalable/apps/gpgui.svg
 	install -Dm644 packaging/files/usr/share/icons/hicolor/32x32/apps/gpgui.png $(DESTDIR)/usr/share/icons/hicolor/32x32/apps/gpgui.png
@@ -146,6 +150,9 @@ uninstall:
 	rm -f $(DESTDIR)/usr/bin/gpgui-helper
 	rm -f $(DESTDIR)/usr/bin/gpgui
 
+	rm -f $(DESTDIR)/usr/lib/systemd/system-sleep/gpclient
+	rm -f $(DESTDIR)/usr/lib/systemd/system-shutdown/gpclient
+
 	rm -f $(DESTDIR)/usr/share/applications/gpgui.desktop
 	rm -f $(DESTDIR)/usr/share/icons/hicolor/scalable/apps/gpgui.svg
 	rm -f $(DESTDIR)/usr/share/icons/hicolor/32x32/apps/gpgui.png

diff --git a/apps/gpclient/Cargo.toml b/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

diff --git a/apps/gpclient/src/disconnect.rs b/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(())
+}
diff --git a/apps/gpservice/Cargo.toml b/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

diff --git a/apps/gpservice/src/cli.rs b/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()?;

diff --git a/apps/gpservice/src/handlers.rs b/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 {

diff --git a/apps/gpservice/src/vpn_task.rs b/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()));

diff --git a/apps/gpservice/src/ws_connection.rs b/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);

diff --git a/apps/gpservice/src/ws_server.rs b/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)
   }

diff --git a/crates/gpapi/Cargo.toml b/crates/gpapi/Cargo.toml
@@ -18,7 +18,7 @@ roxmltree.workspace = true
 serde.workspace = true
 specta = { workspace = true, features = ["derive"] }
 urlencoding.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["process", "signal", "macros"] }
 serde_json.workspace = true
 whoami.workspace = true
 tempfile.workspace = true

diff --git a/crates/gpapi/src/utils/endpoint.rs b/crates/gpapi/src/utils/endpoint.rs
@@ -3,8 +3,13 @@ use tokio::fs;
 use crate::GP_SERVICE_LOCK_FILE;
 
 async fn read_port() -> anyhow::Result<String> {
-  let port = fs::read_to_string(GP_SERVICE_LOCK_FILE).await?;
-  Ok(port.trim().to_string())
+  // PID:PORT
+  let lock_content = fs::read_to_string(GP_SERVICE_LOCK_FILE).await?;
+  let port = lock_content
+    .split(':')
+    .last()
+    .ok_or_else(|| anyhow::anyhow!("Invalid lock file content"))?;
+  Ok(port.to_string())
 }
 
 pub async fn http_endpoint() -> anyhow::Result<String> {

diff --git a/crates/gpapi/src/utils/lock_file.rs b/crates/gpapi/src/utils/lock_file.rs
@@ -2,18 +2,20 @@ use std::path::PathBuf;
 
 pub struct LockFile {
   path: PathBuf,
+  pid: u32,
 }
 
 impl LockFile {
-  pub fn new<P: Into<PathBuf>>(path: P) -> Self {
-    Self { path: path.into() }
+  pub fn new<P: Into<PathBuf>>(path: P, pid: u32) -> Self {
+    Self { path: path.into(), pid }
   }
 
   pub fn exists(&self) -> bool {
     self.path.exists()
   }
 
-  pub fn lock(&self, content: impl AsRef<[u8]>) -> anyhow::Result<()> {
+  pub fn lock(&self, content: &str) -> anyhow::Result<()> {
+    let content = format!("{}:{}", self.pid, content);
     std::fs::write(&self.path, content)?;
     Ok(())
   }

diff --git a/packaging/binary/Makefile.in b/packaging/binary/Makefile.in
@@ -10,6 +10,10 @@ install:
 		install -Dm755 artifacts/usr/bin/gpgui $(DESTDIR)/usr/bin/gpgui; \
 	fi
 
+	# Install the sleep/shutdown scripts
+	install -Dm755 artifacts/usr/lib/systemd/system-sleep/gpclient $(DESTDIR)/usr/lib/systemd/system-sleep/gpclient
+	install -Dm755 artifacts/usr/lib/systemd/system-shutdown/gpclient $(DESTDIR)/usr/lib/systemd/system-shutdown/gpclient
+
 	install -Dm644 artifacts/usr/share/applications/gpgui.desktop $(DESTDIR)/usr/share/applications/gpgui.desktop
 	install -Dm644 artifacts/usr/share/icons/hicolor/scalable/apps/gpgui.svg $(DESTDIR)/usr/share/icons/hicolor/scalable/apps/gpgui.svg
 	install -Dm644 artifacts/usr/share/icons/hicolor/32x32/apps/gpgui.png $(DESTDIR)/usr/share/icons/hicolor/32x32/apps/gpgui.png
@@ -26,6 +30,9 @@ uninstall:
 	rm -f $(DESTDIR)/usr/bin/gpgui-helper
 	rm -f $(DESTDIR)/usr/bin/gpgui
 
+	rm -f $(DESTDIR)/usr/lib/systemd/system-sleep/gpclient
+	rm -f $(DESTDIR)/usr/lib/systemd/system-shutdown/gpclient
+
 	rm -f $(DESTDIR)/usr/share/applications/gpgui.desktop
 	rm -f $(DESTDIR)/usr/share/icons/hicolor/scalable/apps/gpgui.svg
 	rm -f $(DESTDIR)/usr/share/icons/hicolor/32x32/apps/gpgui.png

diff --git a/packaging/files/usr/lib/systemd/system-shutdown/gpclient b/packaging/files/usr/lib/systemd/system-shutdown/gpclient
@@ -0,0 +1,3 @@
+#!/bin/sh
+
+/usr/bin/gpclient disconnect
diff --git a/packaging/files/usr/lib/systemd/system-sleep/gpclient b/packaging/files/usr/lib/systemd/system-sleep/gpclient
@@ -0,0 +1,5 @@
+#!/bin/sh
+
+if [ "${1}" == "pre" ]; then
+  /usr/bin/gpclient disconnect
+fi
diff --git a/packaging/rpm/globalprotect-openconnect.spec.in b/packaging/rpm/globalprotect-openconnect.spec.in
@@ -60,6 +60,13 @@ make build OFFLINE=@OFFLINE@ BUILD_FE=0
 %{_datadir}/icons/hicolor/scalable/apps/gpgui.svg
 %{_datadir}/polkit-1/actions/com.yuezk.gpgui.policy
 
+%dir /usr/lib/systemd
+%dir /usr/lib/systemd/system-sleep
+%dir /usr/lib/systemd/system-shutdown
+
+/usr/lib/systemd/system-sleep/gpclient
+/usr/lib/systemd/system-shutdown/gpclient
+
 %dir %{_datadir}/icons/hicolor
 %dir %{_datadir}/icons/hicolor/32x32
 %dir %{_datadir}/icons/hicolor/32x32/apps