From da2dd5bbfcd4070e8976b7b8e19c10df7b9cfb5b Mon Sep 17 00:00:00 2001 From: Vantz Stockwell Date: Sun, 29 Mar 2026 16:40:10 -0400 Subject: [PATCH] fix: SEC-3/CONC-1/2/3 vault zeroize + async mutex + cancellation tokens - Vault key uses Zeroizing<[u8; 32]>, passwords zeroized after use - vault/credentials Mutex upgraded to tokio::sync::Mutex - CWD tracker + monitor use CancellationToken for clean shutdown - Monitor exec_command has 10s timeout, 3-strike dead connection heuristic Co-Authored-By: Claude Opus 4.6 (1M context) --- src-tauri/Cargo.lock | 2 + src-tauri/Cargo.toml | 2 + src-tauri/src/commands/credentials.rs | 77 +++++++++------------ src-tauri/src/commands/vault.rs | 99 +++++++++++++++------------ src-tauri/src/lib.rs | 13 ++-- src-tauri/src/ssh/cwd.rs | 18 +++-- src-tauri/src/ssh/monitor.rs | 62 ++++++++++++++--- src-tauri/src/ssh/session.rs | 11 ++- src-tauri/src/vault/mod.rs | 15 ++-- 9 files changed, 183 insertions(+), 116 deletions(-) diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 86237cc..978537a 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -8913,9 +8913,11 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-rustls", + "tokio-util", "ureq", "uuid", "x509-cert", + "zeroize", ] [[package]] diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 1563a57..0003c4b 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -33,6 +33,8 @@ uuid = { version = "1", features = ["v4"] } base64 = "0.22" dashmap = "6" tokio = { version = "1", features = ["full"] } +tokio-util = "0.7" +zeroize = { version = "1", features = ["derive"] } async-trait = "0.1" log = "0.4" env_logger = "0.11" diff --git a/src-tauri/src/commands/credentials.rs b/src-tauri/src/commands/credentials.rs index 8681c5f..0513931 100644 --- a/src-tauri/src/commands/credentials.rs +++ b/src-tauri/src/commands/credentials.rs @@ -3,34 +3,16 @@ use tauri::State; use crate::credentials::Credential; use crate::AppState; -/// Guard helper: lock the credentials mutex and return a ref to the inner -/// `CredentialService`, or a "Vault is locked" error if the vault has not -/// been unlocked for this session. -/// -/// This is a macro rather than a function because returning a `MutexGuard` -/// from a helper function would require lifetime annotations that complicate -/// the tauri command signatures unnecessarily. -macro_rules! require_unlocked { - ($state:expr) => {{ - let guard = $state - .credentials - .lock() - .map_err(|_| "Credentials mutex was poisoned".to_string())?; - if guard.is_none() { - return Err("Vault is locked — call unlock before accessing credentials".into()); - } - // SAFETY: we just checked `is_none` above, so `unwrap` cannot panic. - guard - }}; -} - /// Return all credentials ordered by name. /// /// Secret values (passwords, private keys) are never included — only metadata. #[tauri::command] -pub fn list_credentials(state: State<'_, AppState>) -> Result, String> { - let guard = require_unlocked!(state); - guard.as_ref().unwrap().list() +pub async fn list_credentials(state: State<'_, AppState>) -> Result, String> { + let guard = state.credentials.lock().await; + let svc = guard + .as_ref() + .ok_or_else(|| "Vault is locked — call unlock before accessing credentials".to_string())?; + svc.list() } /// Store a new username/password credential. @@ -39,18 +21,18 @@ pub fn list_credentials(state: State<'_, AppState>) -> Result, S /// Returns the created credential record (without the plaintext password). /// `domain` is `None` for non-domain credentials; `Some("")` is treated as NULL. #[tauri::command] -pub fn create_password( +pub async fn create_password( name: String, username: String, password: String, domain: Option, state: State<'_, AppState>, ) -> Result { - let guard = require_unlocked!(state); - guard + let guard = state.credentials.lock().await; + let svc = guard .as_ref() - .unwrap() - .create_password(name, username, password, domain) + .ok_or_else(|| "Vault is locked — call unlock before accessing credentials".to_string())?; + svc.create_password(name, username, password, domain) } /// Store a new SSH private key credential. @@ -59,18 +41,18 @@ pub fn create_password( /// Pass `None` for `passphrase` when the key has no passphrase. /// Returns the created credential record without any secret material. #[tauri::command] -pub fn create_ssh_key( +pub async fn create_ssh_key( name: String, username: String, private_key_pem: String, passphrase: Option, state: State<'_, AppState>, ) -> Result { - let guard = require_unlocked!(state); - guard + let guard = state.credentials.lock().await; + let svc = guard .as_ref() - .unwrap() - .create_ssh_key(name, username, private_key_pem, passphrase) + .ok_or_else(|| "Vault is locked — call unlock before accessing credentials".to_string())?; + svc.create_ssh_key(name, username, private_key_pem, passphrase) } /// Delete a credential by id. @@ -78,21 +60,30 @@ pub fn create_ssh_key( /// For SSH key credentials, the associated `ssh_keys` row is also deleted. /// Returns `Err` if the vault is locked or the id does not exist. #[tauri::command] -pub fn delete_credential(id: i64, state: State<'_, AppState>) -> Result<(), String> { - let guard = require_unlocked!(state); - guard.as_ref().unwrap().delete(id) +pub async fn delete_credential(id: i64, state: State<'_, AppState>) -> Result<(), String> { + let guard = state.credentials.lock().await; + let svc = guard + .as_ref() + .ok_or_else(|| "Vault is locked — call unlock before accessing credentials".to_string())?; + svc.delete(id) } /// Decrypt and return the password for a credential. #[tauri::command] -pub fn decrypt_password(credential_id: i64, state: State<'_, AppState>) -> Result { - let guard = require_unlocked!(state); - guard.as_ref().unwrap().decrypt_password(credential_id) +pub async fn decrypt_password(credential_id: i64, state: State<'_, AppState>) -> Result { + let guard = state.credentials.lock().await; + let svc = guard + .as_ref() + .ok_or_else(|| "Vault is locked — call unlock before accessing credentials".to_string())?; + svc.decrypt_password(credential_id) } /// Decrypt and return the SSH private key and passphrase. #[tauri::command] -pub fn decrypt_ssh_key(ssh_key_id: i64, state: State<'_, AppState>) -> Result<(String, String), String> { - let guard = require_unlocked!(state); - guard.as_ref().unwrap().decrypt_ssh_key(ssh_key_id) +pub async fn decrypt_ssh_key(ssh_key_id: i64, state: State<'_, AppState>) -> Result<(String, String), String> { + let guard = state.credentials.lock().await; + let svc = guard + .as_ref() + .ok_or_else(|| "Vault is locked — call unlock before accessing credentials".to_string())?; + svc.decrypt_ssh_key(ssh_key_id) } diff --git a/src-tauri/src/commands/vault.rs b/src-tauri/src/commands/vault.rs index 1bd896f..86a9c85 100644 --- a/src-tauri/src/commands/vault.rs +++ b/src-tauri/src/commands/vault.rs @@ -1,4 +1,5 @@ use tauri::State; +use zeroize::Zeroize; use crate::vault::{self, VaultService}; use crate::credentials::CredentialService; @@ -21,28 +22,33 @@ pub fn is_first_run(state: State<'_, AppState>) -> bool { /// Returns `Err` if the vault has already been set up or if any storage /// operation fails. #[tauri::command] -pub fn create_vault(password: String, state: State<'_, AppState>) -> Result<(), String> { - if !state.is_first_run() { - return Err("Vault already exists — use unlock instead of create".into()); - } +pub async fn create_vault(mut password: String, state: State<'_, AppState>) -> Result<(), String> { + let result = async { + if !state.is_first_run() { + return Err("Vault already exists — use unlock instead of create".into()); + } - let salt = vault::generate_salt(); - let key = vault::derive_key(&password, &salt); - let vs = VaultService::new(key); + let salt = vault::generate_salt(); + let key = vault::derive_key(&password, &salt); + let vs = VaultService::new(key.clone()); - // Persist the salt so we can re-derive the key on future unlocks. - state.settings.set("vault_salt", &hex::encode(salt))?; + // Persist the salt so we can re-derive the key on future unlocks. + state.settings.set("vault_salt", &hex::encode(salt))?; - // Persist a known-plaintext check so unlock can verify the password. - let check = vs.encrypt("wraith-vault-check")?; - state.settings.set("vault_check", &check)?; + // Persist a known-plaintext check so unlock can verify the password. + let check = vs.encrypt("wraith-vault-check")?; + state.settings.set("vault_check", &check)?; - // Activate the vault and credentials service for this session. - let cred_svc = CredentialService::new(state.db.clone(), VaultService::new(key)); - *state.credentials.lock().unwrap() = Some(cred_svc); - *state.vault.lock().unwrap() = Some(vs); + // Activate the vault and credentials service for this session. + let cred_svc = CredentialService::new(state.db.clone(), VaultService::new(key)); + *state.credentials.lock().await = Some(cred_svc); + *state.vault.lock().await = Some(vs); - Ok(()) + Ok(()) + }.await; + + password.zeroize(); + result } /// Unlock an existing vault using the master password. @@ -52,42 +58,47 @@ pub fn create_vault(password: String, state: State<'_, AppState>) -> Result<(), /// /// Returns `Err("Incorrect master password")` if the password is wrong. #[tauri::command] -pub fn unlock(password: String, state: State<'_, AppState>) -> Result<(), String> { - let salt_hex = state - .settings - .get("vault_salt") - .ok_or_else(|| "Vault has not been set up — call create_vault first".to_string())?; +pub async fn unlock(mut password: String, state: State<'_, AppState>) -> Result<(), String> { + let result = async { + let salt_hex = state + .settings + .get("vault_salt") + .ok_or_else(|| "Vault has not been set up — call create_vault first".to_string())?; - let salt = hex::decode(&salt_hex) - .map_err(|e| format!("Stored vault salt is corrupt: {e}"))?; + let salt = hex::decode(&salt_hex) + .map_err(|e| format!("Stored vault salt is corrupt: {e}"))?; - let key = vault::derive_key(&password, &salt); - let vs = VaultService::new(key); + let key = vault::derive_key(&password, &salt); + let vs = VaultService::new(key.clone()); - // Verify the password by decrypting the check value. - let check_blob = state - .settings - .get("vault_check") - .ok_or_else(|| "Vault check value is missing — vault may be corrupt".to_string())?; + // Verify the password by decrypting the check value. + let check_blob = state + .settings + .get("vault_check") + .ok_or_else(|| "Vault check value is missing — vault may be corrupt".to_string())?; - let check_plain = vs - .decrypt(&check_blob) - .map_err(|_| "Incorrect master password".to_string())?; + let check_plain = vs + .decrypt(&check_blob) + .map_err(|_| "Incorrect master password".to_string())?; - if check_plain != "wraith-vault-check" { - return Err("Incorrect master password".into()); - } + if check_plain != "wraith-vault-check" { + return Err("Incorrect master password".into()); + } - // Activate the vault and credentials service for this session. - let cred_svc = CredentialService::new(state.db.clone(), VaultService::new(key)); - *state.credentials.lock().unwrap() = Some(cred_svc); - *state.vault.lock().unwrap() = Some(vs); + // Activate the vault and credentials service for this session. + let cred_svc = CredentialService::new(state.db.clone(), VaultService::new(key)); + *state.credentials.lock().await = Some(cred_svc); + *state.vault.lock().await = Some(vs); - Ok(()) + Ok(()) + }.await; + + password.zeroize(); + result } /// Returns `true` if the vault is currently unlocked for this session. #[tauri::command] -pub fn is_unlocked(state: State<'_, AppState>) -> bool { - state.is_unlocked() +pub async fn is_unlocked(state: State<'_, AppState>) -> Result { + Ok(state.is_unlocked().await) } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 21a9a59..78b0e1f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -23,7 +23,6 @@ pub mod scanner; pub mod commands; use std::path::PathBuf; -use std::sync::Mutex; use db::Database; use vault::VaultService; @@ -41,10 +40,10 @@ use mcp::error_watcher::ErrorWatcher; pub struct AppState { pub db: Database, - pub vault: Mutex>, + pub vault: tokio::sync::Mutex>, pub settings: SettingsService, pub connections: ConnectionService, - pub credentials: Mutex>, + pub credentials: tokio::sync::Mutex>, pub ssh: SshService, pub sftp: SftpService, pub rdp: RdpService, @@ -62,10 +61,10 @@ impl AppState { database.migrate()?; Ok(Self { db: database.clone(), - vault: Mutex::new(None), + vault: tokio::sync::Mutex::new(None), settings: SettingsService::new(database.clone()), connections: ConnectionService::new(database.clone()), - credentials: Mutex::new(None), + credentials: tokio::sync::Mutex::new(None), ssh: SshService::new(database.clone()), sftp: SftpService::new(), rdp: RdpService::new(), @@ -85,8 +84,8 @@ impl AppState { self.settings.get("vault_salt").unwrap_or_default().is_empty() } - pub fn is_unlocked(&self) -> bool { - self.vault.lock().unwrap().is_some() + pub async fn is_unlocked(&self) -> bool { + self.vault.lock().await.is_some() } } diff --git a/src-tauri/src/ssh/cwd.rs b/src-tauri/src/ssh/cwd.rs index 75f1ac2..1e76dc9 100644 --- a/src-tauri/src/ssh/cwd.rs +++ b/src-tauri/src/ssh/cwd.rs @@ -16,6 +16,7 @@ use russh::ChannelMsg; use tauri::{AppHandle, Emitter}; use tokio::sync::watch; use tokio::sync::Mutex as TokioMutex; +use tokio_util::sync::CancellationToken; use crate::ssh::session::SshClient; @@ -39,13 +40,15 @@ impl CwdTracker { /// Spawn a background tokio task that polls `pwd` every 2 seconds on a /// separate exec channel. /// - /// The task runs until the SSH connection is closed or the channel cannot - /// be opened. CWD changes are emitted as `ssh:cwd:{session_id}` events. + /// The task runs until cancelled via the `CancellationToken`, or until the + /// SSH connection is closed or the channel cannot be opened. + /// CWD changes are emitted as `ssh:cwd:{session_id}` events. pub fn start( &self, handle: Arc>>, app_handle: AppHandle, session_id: String, + cancel: CancellationToken, ) { let sender = self._sender.clone(); @@ -56,6 +59,10 @@ impl CwdTracker { let mut previous_cwd = String::new(); loop { + if cancel.is_cancelled() { + break; + } + // Open a fresh exec channel for each `pwd` invocation. // Some SSH servers do not allow multiple exec requests on a // single channel, so we open a new one each time. @@ -119,8 +126,11 @@ impl CwdTracker { } } - // Wait 2 seconds before the next poll. - tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; + // Wait 2 seconds before the next poll, or cancel. + tokio::select! { + _ = tokio::time::sleep(tokio::time::Duration::from_secs(2)) => {} + _ = cancel.cancelled() => { break; } + } } debug!("CWD tracker for session {} stopped", session_id); diff --git a/src-tauri/src/ssh/monitor.rs b/src-tauri/src/ssh/monitor.rs index 336e216..5db6e7f 100644 --- a/src-tauri/src/ssh/monitor.rs +++ b/src-tauri/src/ssh/monitor.rs @@ -6,11 +6,13 @@ use std::sync::Arc; +use log::warn; use russh::client::Handle; use russh::ChannelMsg; use serde::Serialize; use tauri::{AppHandle, Emitter}; use tokio::sync::Mutex as TokioMutex; +use tokio_util::sync::CancellationToken; use crate::ssh::session::SshClient; @@ -30,26 +32,53 @@ pub struct SystemStats { } /// Spawn a background task that polls system stats every 5 seconds. +/// +/// The task runs until cancelled via the `CancellationToken`, or until the +/// SSH connection is closed. pub fn start_monitor( handle: Arc>>, app_handle: AppHandle, session_id: String, + cancel: CancellationToken, ) { tokio::spawn(async move { // Brief delay to let the shell start up tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; - loop { - let stats = collect_stats(&handle).await; + let mut consecutive_timeouts: u32 = 0; - if let Some(stats) = stats { - let _ = app_handle.emit( - &format!("ssh:monitor:{}", session_id), - &stats, - ); + loop { + if cancel.is_cancelled() { + break; } - tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; + let stats = collect_stats(&handle).await; + + match stats { + Some(stats) => { + consecutive_timeouts = 0; + let _ = app_handle.emit( + &format!("ssh:monitor:{}", session_id), + &stats, + ); + } + None => { + consecutive_timeouts += 1; + if consecutive_timeouts >= 3 { + warn!( + "SSH monitor for session {}: 3 consecutive failures, stopping", + session_id + ); + break; + } + } + } + + // Wait 5 seconds before the next poll, or cancel. + tokio::select! { + _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {} + _ = cancel.cancelled() => { break; } + } } }); } @@ -125,7 +154,24 @@ fn parse_stats(raw: &str) -> Option { }) } +/// Execute a command on a separate exec channel with a 10-second timeout. async fn exec_command(handle: &Arc>>, cmd: &str) -> Option { + let result = tokio::time::timeout( + std::time::Duration::from_secs(10), + exec_command_inner(handle, cmd), + ) + .await; + + match result { + Ok(output) => output, + Err(_) => { + warn!("SSH monitor exec_command timed out after 10s"); + None + } + } +} + +async fn exec_command_inner(handle: &Arc>>, cmd: &str) -> Option { let mut channel = { let h = handle.lock().await; h.channel_open_session().await.ok()? diff --git a/src-tauri/src/ssh/session.rs b/src-tauri/src/ssh/session.rs index 4e61d9a..8626fb1 100644 --- a/src-tauri/src/ssh/session.rs +++ b/src-tauri/src/ssh/session.rs @@ -17,6 +17,7 @@ use crate::mcp::error_watcher::ErrorWatcher; use crate::sftp::SftpService; use crate::ssh::cwd::CwdTracker; use crate::ssh::host_key::{HostKeyResult, HostKeyStore}; +use tokio_util::sync::CancellationToken; pub enum AuthMethod { Password(String), @@ -47,6 +48,7 @@ pub struct SshSession { pub handle: Arc>>, pub command_tx: mpsc::UnboundedSender, pub cwd_tracker: Option, + pub cancel_token: CancellationToken, } pub struct SshClient { @@ -135,10 +137,11 @@ impl SshService { let channel_id = channel.id(); let handle = Arc::new(TokioMutex::new(handle)); let (command_tx, mut command_rx) = mpsc::unbounded_channel::(); + let cancel_token = CancellationToken::new(); let cwd_tracker = CwdTracker::new(); - cwd_tracker.start(handle.clone(), app_handle.clone(), session_id.clone()); + cwd_tracker.start(handle.clone(), app_handle.clone(), session_id.clone(), cancel_token.clone()); - let session = Arc::new(SshSession { id: session_id.clone(), hostname: hostname.to_string(), port, username: username.to_string(), channel_id, handle: handle.clone(), command_tx: command_tx.clone(), cwd_tracker: Some(cwd_tracker) }); + let session = Arc::new(SshSession { id: session_id.clone(), hostname: hostname.to_string(), port, username: username.to_string(), channel_id, handle: handle.clone(), command_tx: command_tx.clone(), cwd_tracker: Some(cwd_tracker), cancel_token: cancel_token.clone() }); self.sessions.insert(session_id.clone(), session); { let h = handle.lock().await; @@ -158,7 +161,7 @@ impl SshService { error_watcher.watch(&session_id); // Start remote monitoring if enabled (runs on a separate exec channel) - crate::ssh::monitor::start_monitor(handle.clone(), app_handle.clone(), session_id.clone()); + crate::ssh::monitor::start_monitor(handle.clone(), app_handle.clone(), session_id.clone(), cancel_token.clone()); // Inject OSC 7 CWD reporting hook into the user's shell. // This enables SFTP CWD following on all platforms (Linux, macOS, FreeBSD). @@ -246,6 +249,8 @@ impl SshService { pub async fn disconnect(&self, session_id: &str, sftp_service: &SftpService) -> Result<(), String> { let (_, session) = self.sessions.remove(session_id).ok_or_else(|| format!("Session {} not found", session_id))?; + // Cancel background tasks (CWD tracker, monitor) before tearing down the connection. + session.cancel_token.cancel(); let _ = session.command_tx.send(ChannelCommand::Shutdown); { let handle = session.handle.lock().await; let _ = handle.disconnect(Disconnect::ByApplication, "", "en").await; } sftp_service.remove_client(session_id); diff --git a/src-tauri/src/vault/mod.rs b/src-tauri/src/vault/mod.rs index e712acd..3f30ab8 100644 --- a/src-tauri/src/vault/mod.rs +++ b/src-tauri/src/vault/mod.rs @@ -4,6 +4,7 @@ use aes_gcm::{ Aes256Gcm, Key, Nonce, }; use argon2::{Algorithm, Argon2, Params, Version}; +use zeroize::Zeroizing; // --------------------------------------------------------------------------- // VaultService @@ -21,18 +22,18 @@ use argon2::{Algorithm, Argon2, Params, Version}; /// The version prefix allows a future migration to a different algorithm /// without breaking existing stored blobs. pub struct VaultService { - key: [u8; 32], + key: Zeroizing<[u8; 32]>, } impl VaultService { - pub fn new(key: [u8; 32]) -> Self { + pub fn new(key: Zeroizing<[u8; 32]>) -> Self { Self { key } } /// Encrypt `plaintext` and return a `v1:{iv_hex}:{sealed_hex}` blob. pub fn encrypt(&self, plaintext: &str) -> Result { // Build the AES-256-GCM cipher from our key. - let key = Key::::from_slice(&self.key); + let key = Key::::from_slice(&*self.key); let cipher = Aes256Gcm::new(key); // Generate a random 12-byte nonce (96-bit is the GCM standard). @@ -71,7 +72,7 @@ impl VaultService { )); } - let key = Key::::from_slice(&self.key); + let key = Key::::from_slice(&*self.key); let cipher = Aes256Gcm::new(key); let nonce = Nonce::from_slice(&iv_bytes); @@ -95,7 +96,7 @@ impl VaultService { /// t = 3 iterations /// m = 65536 KiB (64 MiB) memory /// p = 4 parallelism lanes -pub fn derive_key(password: &str, salt: &[u8]) -> [u8; 32] { +pub fn derive_key(password: &str, salt: &[u8]) -> Zeroizing<[u8; 32]> { let params = Params::new( 65536, // m_cost: 64 MiB 3, // t_cost: iterations @@ -106,9 +107,9 @@ pub fn derive_key(password: &str, salt: &[u8]) -> [u8; 32] { let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params); - let mut output_key = [0u8; 32]; + let mut output_key = Zeroizing::new([0u8; 32]); argon2 - .hash_password_into(password.as_bytes(), salt, &mut output_key) + .hash_password_into(password.as_bytes(), salt, &mut *output_key) .expect("Argon2id key derivation failed"); output_key