Refactor using Tauri (#278)

apps/gpservice/src/cli.rs

@@ -0,0 +1,182 @@
+use std::sync::Arc;
+use std::{collections::HashMap, io::Write};
+
+use anyhow::bail;
+use clap::Parser;
+use gpapi::{
+  process::gui_launcher::GuiLauncher,
+  service::{request::WsRequest, vpn_state::VpnState},
+  utils::{
+    crypto::generate_key, env_file, lock_file::LockFile, redact::Redaction, shutdown_signal,
+  },
+  GP_SERVICE_LOCK_FILE,
+};
+use log::{info, warn, LevelFilter};
+use tokio::sync::{mpsc, watch};
+
+use crate::{vpn_task::VpnTask, ws_server::WsServer};
+
+const VERSION: &str = concat!(
+  env!("CARGO_PKG_VERSION"),
+  " (",
+  compile_time::date_str!(),
+  ")"
+);
+
+#[derive(Parser)]
+#[command(version = VERSION)]
+struct Cli {
+  #[clap(long)]
+  minimized: bool,
+  #[clap(long)]
+  env_file: Option<String>,
+  #[cfg(debug_assertions)]
+  #[clap(long)]
+  no_gui: bool,
+}
+
+impl Cli {
+  async fn run(&mut self, redaction: Arc<Redaction>) -> anyhow::Result<()> {
+    let lock_file = Arc::new(LockFile::new(GP_SERVICE_LOCK_FILE));
+
+    if lock_file.check_health().await {
+      bail!("Another instance of the service is already running");
+    }
+
+    let api_key = self.prepare_api_key();
+
+    // Channel for sending requests to the VPN task
+    let (ws_req_tx, ws_req_rx) = mpsc::channel::<WsRequest>(32);
+    // Channel for receiving the VPN state from the VPN task
+    let (vpn_state_tx, vpn_state_rx) = watch::channel(VpnState::Disconnected);
+
+    let mut vpn_task = VpnTask::new(ws_req_rx, vpn_state_tx);
+    let ws_server = WsServer::new(
+      api_key.clone(),
+      ws_req_tx,
+      vpn_state_rx,
+      lock_file.clone(),
+      redaction,
+    );
+
+    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 server_token = ws_server.cancel_token();
+
+    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 });
+
+    #[cfg(debug_assertions)]
+    let no_gui = self.no_gui;
+
+    #[cfg(not(debug_assertions))]
+    let no_gui = false;
+
+    if no_gui {
+      info!("GUI is disabled");
+    } else {
+      let envs = self
+        .env_file
+        .as_ref()
+        .map(env_file::load_env_vars)
+        .transpose()?;
+
+      let minimized = self.minimized;
+
+      tokio::spawn(async move {
+        launch_gui(envs, api_key, minimized).await;
+        let _ = shutdown_tx.send(()).await;
+      });
+    }
+
+    tokio::select! {
+        _ = shutdown_signal() => {
+            info!("Shutdown signal received");
+        }
+        _ = shutdown_rx.recv() => {
+            info!("Shutdown request received, shutting down");
+        }
+    }
+
+    vpn_task_token.cancel();
+    let _ = tokio::join!(vpn_task_handle, ws_server_handle);
+
+    lock_file.unlock()?;
+
+    info!("gpservice stopped");
+
+    Ok(())
+  }
+
+  fn prepare_api_key(&self) -> Vec<u8> {
+    #[cfg(debug_assertions)]
+    if self.no_gui {
+      return gpapi::GP_API_KEY.to_vec();
+    }
+
+    generate_key().to_vec()
+  }
+}
+
+fn init_logger() -> Arc<Redaction> {
+  let redaction = Arc::new(Redaction::new());
+  let redaction_clone = Arc::clone(&redaction);
+  // let target = Box::new(File::create("log.txt").expect("Can't create file"));
+  env_logger::builder()
+    .filter_level(LevelFilter::Info)
+    .format(move |buf, record| {
+      let timestamp = buf.timestamp();
+      writeln!(
+        buf,
+        "[{} {} {}] {}",
+        timestamp,
+        record.level(),
+        record.module_path().unwrap_or_default(),
+        redaction_clone.redact_str(&record.args().to_string())
+      )
+    })
+    // .target(env_logger::Target::Pipe(target))
+    .init();
+
+  redaction
+}
+
+async fn launch_gui(envs: Option<HashMap<String, String>>, api_key: Vec<u8>, mut minimized: bool) {
+  loop {
+    let api_key_clone = api_key.clone();
+    let gui_launcher = GuiLauncher::new()
+      .envs(envs.clone())
+      .api_key(api_key_clone)
+      .minimized(minimized);
+
+    match gui_launcher.launch().await {
+      Ok(exit_status) => {
+        // Exit code 99 means that the GUI needs to be restarted
+        if exit_status.code() != Some(99) {
+          info!("GUI exited with code {:?}", exit_status.code());
+          break;
+        }
+
+        info!("GUI exited with code 99, restarting");
+        minimized = false;
+      }
+      Err(err) => {
+        warn!("Failed to launch GUI: {}", err);
+        break;
+      }
+    }
+  }
+}
+
+pub async fn run() {
+  let mut cli = Cli::parse();
+
+  let redaction = init_logger();
+  info!("gpservice started: {}", VERSION);
+
+  if let Err(e) = cli.run(redaction).await {
+    eprintln!("Error: {}", e);
+    std::process::exit(1);
+  }
+}

apps/gpservice/src/handlers.rs

@@ -0,0 +1,94 @@
+use std::{borrow::Cow, ops::ControlFlow, sync::Arc};
+
+use axum::{
+  extract::{
+    ws::{self, CloseFrame, Message, WebSocket},
+    State, WebSocketUpgrade,
+  },
+  response::IntoResponse,
+};
+use futures::{SinkExt, StreamExt};
+use gpapi::service::event::WsEvent;
+use log::{info, warn};
+
+use crate::ws_server::WsServerContext;
+
+pub(crate) async fn health() -> impl IntoResponse {
+  "OK"
+}
+
+pub(crate) async fn active_gui(State(ctx): State<Arc<WsServerContext>>) -> impl IntoResponse {
+  ctx.send_event(WsEvent::ActiveGui).await;
+}
+
+pub(crate) async fn ws_handler(
+  ws: WebSocketUpgrade,
+  State(ctx): State<Arc<WsServerContext>>,
+) -> impl IntoResponse {
+  ws.on_upgrade(move |socket| handle_socket(socket, ctx))
+}
+
+async fn handle_socket(mut socket: WebSocket, ctx: Arc<WsServerContext>) {
+  // Send ping message
+  if let Err(err) = socket.send(Message::Ping("Hi".into())).await {
+    warn!("Failed to send ping: {}", err);
+    return;
+  }
+
+  // Wait for pong message
+  if socket.recv().await.is_none() {
+    warn!("Failed to receive pong");
+    return;
+  }
+
+  info!("New client connected");
+
+  let (mut sender, mut receiver) = socket.split();
+  let (connection, mut msg_rx) = ctx.add_connection().await;
+
+  let send_task = tokio::spawn(async move {
+    while let Some(msg) = msg_rx.recv().await {
+      if let Err(err) = sender.send(msg).await {
+        info!("Failed to send message: {}", err);
+        break;
+      }
+    }
+
+    let close_msg = Message::Close(Some(CloseFrame {
+      code: ws::close_code::NORMAL,
+      reason: Cow::from("Goodbye"),
+    }));
+
+    if let Err(err) = sender.send(close_msg).await {
+      warn!("Failed to close socket: {}", err);
+    }
+  });
+
+  let conn = Arc::clone(&connection);
+  let ctx_clone = Arc::clone(&ctx);
+  let recv_task = tokio::spawn(async move {
+    while let Some(Ok(msg)) = receiver.next().await {
+      let ControlFlow::Continue(ws_req) = conn.recv_msg(msg) else {
+        break;
+      };
+
+      if let Err(err) = ctx_clone.forward_req(ws_req).await {
+        info!("Failed to forward request: {}", err);
+        break;
+      }
+    }
+  });
+
+  tokio::select! {
+    _ = send_task => {
+        info!("WS server send task completed");
+    },
+    _ = recv_task => {
+        info!("WS server recv task completed");
+    }
+  }
+
+  info!("Client disconnected");
+
+  ctx.remove_connection(connection).await;
+}

apps/gpservice/src/main.rs

@@ -0,0 +1,11 @@
+mod cli;
+mod handlers;
+mod routes;
+mod vpn_task;
+mod ws_server;
+mod ws_connection;
+
+#[tokio::main]
+async fn main() {
+  cli::run().await;
+}

apps/gpservice/src/routes.rs

@@ -0,0 +1,13 @@
+use std::sync::Arc;
+
+use axum::{routing::{get, post}, Router};
+
+use crate::{handlers, ws_server::WsServerContext};
+
+pub(crate) fn routes(ctx: Arc<WsServerContext>) -> Router {
+  Router::new()
+    .route("/health", get(handlers::health))
+    .route("/active-gui", post(handlers::active_gui))
+    .route("/ws", get(handlers::ws_handler))
+    .with_state(ctx)
+}

apps/gpservice/src/vpn_task.rs

@@ -0,0 +1,144 @@
+use std::{sync::Arc, thread};
+
+use gpapi::service::{
+  request::{ConnectRequest, WsRequest},
+  vpn_state::VpnState,
+};
+use log::info;
+use openconnect::Vpn;
+use tokio::sync::{mpsc, oneshot, watch, RwLock};
+use tokio_util::sync::CancellationToken;
+
+pub(crate) struct VpnTaskContext {
+  vpn_handle: Arc<RwLock<Option<Vpn>>>,
+  vpn_state_tx: Arc<watch::Sender<VpnState>>,
+  disconnect_rx: RwLock<Option<oneshot::Receiver<()>>>,
+}
+
+impl VpnTaskContext {
+  pub fn new(vpn_state_tx: watch::Sender<VpnState>) -> Self {
+    Self {
+      vpn_handle: Default::default(),
+      vpn_state_tx: Arc::new(vpn_state_tx),
+      disconnect_rx: Default::default(),
+    }
+  }
+
+  pub async fn connect(&self, req: ConnectRequest) {
+    let vpn_state = self.vpn_state_tx.borrow().clone();
+    if !matches!(vpn_state, VpnState::Disconnected) {
+      info!("VPN is not disconnected, ignore the request");
+      return;
+    }
+
+    let info = req.info().clone();
+    let vpn_handle = self.vpn_handle.clone();
+    let args = req.args();
+    let vpn = Vpn::builder(req.gateway().server(), args.cookie())
+      .user_agent(args.user_agent())
+      .script(args.vpnc_script())
+      .os(args.openconnect_os())
+      .build();
+
+    // Save the VPN handle
+    vpn_handle.write().await.replace(vpn);
+
+    let vpn_state_tx = self.vpn_state_tx.clone();
+    let connect_info = Box::new(info.clone());
+    vpn_state_tx.send(VpnState::Connecting(connect_info)).ok();
+
+    let (disconnect_tx, disconnect_rx) = oneshot::channel::<()>();
+    self.disconnect_rx.write().await.replace(disconnect_rx);
+
+    // Spawn a new thread to process the VPN connection, cannot use tokio::spawn here.
+    // Otherwise, it will block the tokio runtime and cannot send the VPN state to the channel
+    thread::spawn(move || {
+      let vpn_state_tx_clone = vpn_state_tx.clone();
+
+      vpn_handle.blocking_read().as_ref().map(|vpn| {
+        vpn.connect(move || {
+          let connect_info = Box::new(info.clone());
+          vpn_state_tx.send(VpnState::Connected(connect_info)).ok();
+        })
+      });
+
+      // Notify the VPN is disconnected
+      vpn_state_tx_clone.send(VpnState::Disconnected).ok();
+      // Remove the VPN handle
+      vpn_handle.blocking_write().take();
+
+      disconnect_tx.send(()).ok();
+    });
+  }
+
+  pub async fn disconnect(&self) {
+    if let Some(disconnect_rx) = self.disconnect_rx.write().await.take() {
+      if let Some(vpn) = self.vpn_handle.read().await.as_ref() {
+        self.vpn_state_tx.send(VpnState::Disconnecting).ok();
+        vpn.disconnect()
+      }
+      // Wait for the VPN to be disconnected
+      disconnect_rx.await.ok();
+      info!("VPN disconnected");
+    } else {
+      info!("VPN is not connected, skip disconnect");
+      self.vpn_state_tx.send(VpnState::Disconnected).ok();
+    }
+  }
+}
+
+pub(crate) struct VpnTask {
+  ws_req_rx: mpsc::Receiver<WsRequest>,
+  ctx: Arc<VpnTaskContext>,
+  cancel_token: CancellationToken,
+}
+
+impl VpnTask {
+  pub fn new(ws_req_rx: mpsc::Receiver<WsRequest>, vpn_state_tx: watch::Sender<VpnState>) -> Self {
+    let ctx = Arc::new(VpnTaskContext::new(vpn_state_tx));
+    let cancel_token = CancellationToken::new();
+
+    Self {
+      ws_req_rx,
+      ctx,
+      cancel_token,
+    }
+  }
+
+  pub fn cancel_token(&self) -> CancellationToken {
+    self.cancel_token.clone()
+  }
+
+  pub async fn start(&mut self, server_cancel_token: CancellationToken) {
+    let cancel_token = self.cancel_token.clone();
+
+    tokio::select! {
+        _ = self.recv() => {
+            info!("VPN task stopped");
+        }
+        _ = cancel_token.cancelled() => {
+            info!("VPN task cancelled");
+            self.ctx.disconnect().await;
+        }
+    }
+
+    server_cancel_token.cancel();
+  }
+
+  async fn recv(&mut self) {
+    while let Some(req) = self.ws_req_rx.recv().await {
+      tokio::spawn(process_ws_req(req, self.ctx.clone()));
+    }
+  }
+}
+
+async fn process_ws_req(req: WsRequest, ctx: Arc<VpnTaskContext>) {
+  match req {
+    WsRequest::Connect(req) => {
+      ctx.connect(*req).await;
+    }
+    WsRequest::Disconnect(_) => {
+      ctx.disconnect().await;
+    }
+  }
+}

apps/gpservice/src/ws_connection.rs

@@ -0,0 +1,53 @@
+use std::{ops::ControlFlow, sync::Arc};
+
+use axum::extract::ws::{CloseFrame, Message};
+use gpapi::{
+  service::{event::WsEvent, request::WsRequest},
+  utils::crypto::Crypto,
+};
+use log::{info, warn};
+use tokio::sync::mpsc;
+
+pub(crate) struct WsConnection {
+  crypto: Arc<Crypto>,
+  tx: mpsc::Sender<Message>,
+}
+
+impl WsConnection {
+  pub fn new(crypto: Arc<Crypto>, tx: mpsc::Sender<Message>) -> Self {
+    Self { crypto, tx }
+  }
+
+  pub async fn send_event(&self, event: &WsEvent) -> anyhow::Result<()> {
+    let encrypted = self.crypto.encrypt(event)?;
+    let msg = Message::Binary(encrypted);
+
+    self.tx.send(msg).await?;
+
+    Ok(())
+  }
+
+  pub fn recv_msg(&self, msg: Message) -> ControlFlow<(), WsRequest> {
+    match msg {
+      Message::Binary(data) => match self.crypto.decrypt(data) {
+        Ok(ws_req) => ControlFlow::Continue(ws_req),
+        Err(err) => {
+          info!("Failed to decrypt message: {}", err);
+          ControlFlow::Break(())
+        }
+      },
+      Message::Close(cf) => {
+        if let Some(CloseFrame { code, reason }) = cf {
+          info!("Client sent close, code {} and reason `{}`", code, reason);
+        } else {
+          info!("Client somehow sent close message without CloseFrame");
+        }
+        ControlFlow::Break(())
+      }
+      _ => {
+        warn!("WS server received unexpected message: {:?}", msg);
+        ControlFlow::Break(())
+      }
+    }
+  }
+}

apps/gpservice/src/ws_server.rs

@@ -0,0 +1,158 @@
+use std::sync::Arc;
+
+use axum::extract::ws::Message;
+use gpapi::{
+  service::{event::WsEvent, request::WsRequest, vpn_state::VpnState},
+  utils::{crypto::Crypto, lock_file::LockFile, redact::Redaction},
+};
+use log::{info, warn};
+use tokio::{
+  net::TcpListener,
+  sync::{mpsc, watch, RwLock},
+};
+use tokio_util::sync::CancellationToken;
+
+use crate::{routes, ws_connection::WsConnection};
+
+pub(crate) struct WsServerContext {
+  crypto: Arc<Crypto>,
+  ws_req_tx: mpsc::Sender<WsRequest>,
+  vpn_state_rx: watch::Receiver<VpnState>,
+  redaction: Arc<Redaction>,
+  connections: RwLock<Vec<Arc<WsConnection>>>,
+}
+
+impl WsServerContext {
+  pub fn new(
+    api_key: Vec<u8>,
+    ws_req_tx: mpsc::Sender<WsRequest>,
+    vpn_state_rx: watch::Receiver<VpnState>,
+    redaction: Arc<Redaction>,
+  ) -> Self {
+    Self {
+      crypto: Arc::new(Crypto::new(api_key)),
+      ws_req_tx,
+      vpn_state_rx,
+      redaction,
+      connections: Default::default(),
+    }
+  }
+
+  pub async fn send_event(&self, event: WsEvent) {
+    let connections = self.connections.read().await;
+
+    for conn in connections.iter() {
+      let _ = conn.send_event(&event).await;
+    }
+  }
+
+  pub async fn add_connection(&self) -> (Arc<WsConnection>, mpsc::Receiver<Message>) {
+    let (tx, rx) = mpsc::channel::<Message>(32);
+    let conn = Arc::new(WsConnection::new(Arc::clone(&self.crypto), tx));
+
+    // Send current VPN state to new client
+    info!("Sending current VPN state to new client");
+    let vpn_state = self.vpn_state_rx.borrow().clone();
+    if let Err(err) = conn.send_event(&WsEvent::VpnState(vpn_state)).await {
+      warn!("Failed to send VPN state to new client: {}", err);
+    }
+
+    self.connections.write().await.push(Arc::clone(&conn));
+
+    (conn, rx)
+  }
+
+  pub async fn remove_connection(&self, conn: Arc<WsConnection>) {
+    let mut connections = self.connections.write().await;
+    connections.retain(|c| !Arc::ptr_eq(c, &conn));
+  }
+
+  fn vpn_state_rx(&self) -> watch::Receiver<VpnState> {
+    self.vpn_state_rx.clone()
+  }
+
+  pub async fn forward_req(&self, req: WsRequest) -> anyhow::Result<()> {
+    if let WsRequest::Connect(ref req) = req {
+      self
+        .redaction
+        .add_values(&[req.gateway().server(), req.args().cookie()])?
+    }
+
+    self.ws_req_tx.send(req).await?;
+
+    Ok(())
+  }
+}
+
+pub(crate) struct WsServer {
+  ctx: Arc<WsServerContext>,
+  cancel_token: CancellationToken,
+  lock_file: Arc<LockFile>,
+}
+
+impl WsServer {
+  pub fn new(
+    api_key: Vec<u8>,
+    ws_req_tx: mpsc::Sender<WsRequest>,
+    vpn_state_rx: watch::Receiver<VpnState>,
+    lock_file: Arc<LockFile>,
+    redaction: Arc<Redaction>,
+  ) -> Self {
+    let ctx = Arc::new(WsServerContext::new(
+      api_key,
+      ws_req_tx,
+      vpn_state_rx,
+      redaction,
+    ));
+    let cancel_token = CancellationToken::new();
+
+    Self {
+      ctx,
+      cancel_token,
+      lock_file,
+    }
+  }
+
+  pub fn cancel_token(&self) -> CancellationToken {
+    self.cancel_token.clone()
+  }
+
+  pub async fn start(&self, shutdown_tx: mpsc::Sender<()>) {
+    if let Ok(listener) = TcpListener::bind("127.0.0.1:0").await {
+      let local_addr = listener.local_addr().unwrap();
+
+      self.lock_file.lock(local_addr.port().to_string()).unwrap();
+
+      info!("WS server listening on port: {}", local_addr.port());
+
+      tokio::select! {
+        _ = watch_vpn_state(self.ctx.vpn_state_rx(), Arc::clone(&self.ctx)) => {
+          info!("VPN state watch task completed");
+        }
+        _ = start_server(listener, self.ctx.clone()) => {
+            info!("WS server stopped");
+        }
+        _ = self.cancel_token.cancelled() => {
+          info!("WS server cancelled");
+        }
+      }
+    }
+
+    let _ = shutdown_tx.send(()).await;
+  }
+}
+
+async fn watch_vpn_state(mut vpn_state_rx: watch::Receiver<VpnState>, ctx: Arc<WsServerContext>) {
+  while vpn_state_rx.changed().await.is_ok() {
+    let vpn_state = vpn_state_rx.borrow().clone();
+    ctx.send_event(WsEvent::VpnState(vpn_state)).await;
+  }
+}
+
+async fn start_server(listener: TcpListener, ctx: Arc<WsServerContext>) -> anyhow::Result<()> {
+  let routes = routes::routes(ctx);
+
+  axum::serve(listener, routes).await?;
+
+  Ok(())
+}