From df334e4db8813761761e6fbdaab83bcd653d9060 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Wed, 4 Feb 2026 16:50:36 -0600 Subject: [PATCH] Security hardening and performance improvements --- Cargo.lock | 1 + Cargo.toml | 1 + src/main.rs | 13 +- src/mikrotik/types.rs | 75 +----------- src/secret.rs | 87 +++++++++++++ src/snmp/client.rs | 65 +++++++--- src/snmp/device_poller.rs | 34 +++++- src/snmp/poller_registry.rs | 6 +- src/ssh/client.rs | 10 +- src/version.rs | 2 +- src/websocket_client.rs | 236 +++++++++++++++++++++++++++--------- 11 files changed, 367 insertions(+), 163 deletions(-) create mode 100644 src/secret.rs diff --git a/Cargo.lock b/Cargo.lock index 5cbf379..9c40c12 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2606,6 +2606,7 @@ dependencies = [ "tokio-tungstenite", "tracing", "tracing-subscriber", + "zeroize", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 3ecd6ed..f56c8b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } russh = "0.57" async-trait = "0.1" home = "=0.5.12" +zeroize = { version = "1", features = ["derive"] } [build-dependencies] prost-build = "0.14" diff --git a/src/main.rs b/src/main.rs index 554d4ce..21d98c8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod mikrotik; mod proto; +pub mod secret; mod snmp; mod ssh; mod version; @@ -191,6 +192,8 @@ async fn async_main() { let max_retry_delay = Duration::from_secs(60); let mut attempt = 0; + let token = secret::SecretString::new(args.token.as_ref().unwrap().clone()); + loop { // Check if shutdown was requested if *shutdown_rx.borrow() { @@ -213,7 +216,7 @@ async fn async_main() { } // Connect to Towerops server via WebSocket - let mut client = match AgentClient::connect(&ws_url, args.token.as_ref().unwrap()).await { + let mut client = match AgentClient::connect(&ws_url, &token).await { Ok(client) => { tracing::info!("Successfully connected to server"); // Mark as connected for health check @@ -255,6 +258,7 @@ async fn async_main() { /// Run SNMPv3 test async fn run_snmpv3_test(args: &Args) { + use secret::SecretString; use snmp::V3Config; let ip = args.snmpv3_ip.as_ref().expect("--snmpv3-ip required"); @@ -284,12 +288,12 @@ async fn run_snmpv3_test(args: &Args) { let v3_config = V3Config { username: username.clone(), auth_password: if !auth_pass.is_empty() { - Some(auth_pass.clone()) + Some(SecretString::new(auth_pass.clone())) } else { None }, priv_password: if !priv_pass.is_empty() { - Some(priv_pass.clone()) + Some(SecretString::new(priv_pass.clone())) } else { None }, @@ -353,7 +357,8 @@ async fn run_snmpv3_test(args: &Args) { /// Run MikroTik API test async fn run_mikrotik_test(args: &Args) { - use mikrotik::{MikrotikClient, SecretString}; + use mikrotik::MikrotikClient; + use secret::SecretString; let ip = args.mikrotik_ip.as_ref().expect("--mikrotik-ip required"); let port = args.mikrotik_port; diff --git a/src/mikrotik/types.rs b/src/mikrotik/types.rs index 8877416..38b406d 100644 --- a/src/mikrotik/types.rs +++ b/src/mikrotik/types.rs @@ -1,44 +1,6 @@ use std::collections::HashMap; -/// A wrapper for sensitive strings (passwords, tokens) that prevents accidental logging. -/// - Debug and Display show "[REDACTED]" instead of the actual value -/// - The inner value is zeroed on drop -#[derive(Clone)] -pub struct SecretString(String); - -impl SecretString { - pub fn new(value: impl Into) -> Self { - Self(value.into()) - } - - /// Access the secret value. Use sparingly and never log the result. - pub fn expose(&self) -> &str { - &self.0 - } -} - -impl std::fmt::Debug for SecretString { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "[REDACTED]") - } -} - -impl std::fmt::Display for SecretString { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "[REDACTED]") - } -} - -impl Drop for SecretString { - fn drop(&mut self) { - // Zero out the memory before dropping - // SAFETY: We're replacing the string contents with zeros - unsafe { - let bytes = self.0.as_bytes_mut(); - std::ptr::write_bytes(bytes.as_mut_ptr(), 0, bytes.len()); - } - } -} +pub use crate::secret::SecretString; /// Error types for MikroTik operations #[derive(Debug)] @@ -142,39 +104,4 @@ mod tests { assert!(response.sentences.is_empty()); assert!(response.error.is_none()); } - - #[test] - fn test_secret_string_expose() { - let secret = SecretString::new("my_password"); - assert_eq!(secret.expose(), "my_password"); - } - - #[test] - fn test_secret_string_debug_is_redacted() { - let secret = SecretString::new("my_password"); - let debug_output = format!("{:?}", secret); - assert_eq!(debug_output, "[REDACTED]"); - assert!(!debug_output.contains("my_password")); - } - - #[test] - fn test_secret_string_display_is_redacted() { - let secret = SecretString::new("my_password"); - let display_output = format!("{}", secret); - assert_eq!(display_output, "[REDACTED]"); - assert!(!display_output.contains("my_password")); - } - - #[test] - fn test_secret_string_clone() { - let secret = SecretString::new("my_password"); - let cloned = secret.clone(); - assert_eq!(cloned.expose(), "my_password"); - } - - #[test] - fn test_secret_string_empty() { - let secret = SecretString::new(""); - assert!(secret.expose().is_empty()); - } } diff --git a/src/secret.rs b/src/secret.rs new file mode 100644 index 0000000..90f36a3 --- /dev/null +++ b/src/secret.rs @@ -0,0 +1,87 @@ +use zeroize::Zeroize; + +/// A wrapper for sensitive strings (passwords, tokens) that prevents accidental logging. +/// - Debug and Display show "[REDACTED]" instead of the actual value +/// - The inner value is zeroized on drop using volatile writes (cannot be optimized away) +#[derive(Clone)] +pub struct SecretString(String); + +impl SecretString { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + /// Access the secret value. Use sparingly and never log the result. + pub fn expose(&self) -> &str { + &self.0 + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +impl std::fmt::Debug for SecretString { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[REDACTED]") + } +} + +impl std::fmt::Display for SecretString { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[REDACTED]") + } +} + +impl Drop for SecretString { + fn drop(&mut self) { + self.0.zeroize(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_secret_string_expose() { + let secret = SecretString::new("my_password"); + assert_eq!(secret.expose(), "my_password"); + } + + #[test] + fn test_secret_string_debug_is_redacted() { + let secret = SecretString::new("my_password"); + let debug_output = format!("{:?}", secret); + assert_eq!(debug_output, "[REDACTED]"); + assert!(!debug_output.contains("my_password")); + } + + #[test] + fn test_secret_string_display_is_redacted() { + let secret = SecretString::new("my_password"); + let display_output = format!("{}", secret); + assert_eq!(display_output, "[REDACTED]"); + assert!(!display_output.contains("my_password")); + } + + #[test] + fn test_secret_string_clone() { + let secret = SecretString::new("my_password"); + let cloned = secret.clone(); + assert_eq!(cloned.expose(), "my_password"); + } + + #[test] + fn test_secret_string_empty() { + let secret = SecretString::new(""); + assert!(secret.is_empty()); + assert!(secret.expose().is_empty()); + } + + #[test] + fn test_secret_string_is_empty() { + let secret = SecretString::new("not_empty"); + assert!(!secret.is_empty()); + } +} diff --git a/src/snmp/client.rs b/src/snmp/client.rs index 39bfab4..4b94ab0 100644 --- a/src/snmp/client.rs +++ b/src/snmp/client.rs @@ -1,4 +1,5 @@ use super::types::{SnmpError, SnmpResult, SnmpValue}; +use crate::secret::SecretString; use snmp2::SyncSession; use std::time::Duration; @@ -7,16 +8,35 @@ use std::time::Duration; const SNMP_TIMEOUT_SECS: u64 = 30; /// SNMPv3 configuration bundle -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct V3Config { pub username: String, - pub auth_password: Option, - pub priv_password: Option, + pub auth_password: Option, + pub priv_password: Option, pub auth_protocol: Option, pub priv_protocol: Option, pub security_level: String, } +impl std::fmt::Debug for V3Config { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("V3Config") + .field("username", &self.username) + .field( + "auth_password", + &self.auth_password.as_ref().map(|_| "[REDACTED]"), + ) + .field( + "priv_password", + &self.priv_password.as_ref().map(|_| "[REDACTED]"), + ) + .field("auth_protocol", &self.auth_protocol) + .field("priv_protocol", &self.priv_protocol) + .field("security_level", &self.security_level) + .finish() + } +} + /// SNMP client for polling devices #[derive(Debug, Clone, Copy)] pub struct SnmpClient; @@ -283,9 +303,13 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult { parse_auth_protocol(config.auth_protocol.as_deref().ok_or_else(|| { SnmpError::RequestFailed("Auth protocol required for authNoPriv".into()) })?)?; - let _auth_pass = config.auth_password.as_deref().ok_or_else(|| { - SnmpError::RequestFailed("Auth password required for authNoPriv".into()) - })?; + let _auth_pass = config + .auth_password + .as_ref() + .map(|s| s.expose()) + .ok_or_else(|| { + SnmpError::RequestFailed("Auth password required for authNoPriv".into()) + })?; Auth::AuthNoPriv } @@ -295,16 +319,24 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult { parse_auth_protocol(config.auth_protocol.as_deref().ok_or_else(|| { SnmpError::RequestFailed("Auth protocol required for authPriv".into()) })?)?; - let _auth_pass = config.auth_password.as_deref().ok_or_else(|| { - SnmpError::RequestFailed("Auth password required for authPriv".into()) - })?; + let _auth_pass = config + .auth_password + .as_ref() + .map(|s| s.expose()) + .ok_or_else(|| { + SnmpError::RequestFailed("Auth password required for authPriv".into()) + })?; let cipher = parse_priv_protocol(config.priv_protocol.as_deref().ok_or_else(|| { SnmpError::RequestFailed("Priv protocol required for authPriv".into()) })?)?; - let priv_pass = config.priv_password.as_deref().ok_or_else(|| { - SnmpError::RequestFailed("Priv password required for authPriv".into()) - })?; + let priv_pass = config + .priv_password + .as_ref() + .map(|s| s.expose()) + .ok_or_else(|| { + SnmpError::RequestFailed("Priv password required for authPriv".into()) + })?; Auth::AuthPriv { cipher, @@ -321,7 +353,12 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult { }; // Get auth password (empty for noAuthNoPriv) - let auth_password = config.auth_password.as_deref().unwrap_or("").as_bytes(); + let auth_password = config + .auth_password + .as_ref() + .map(|s| s.expose()) + .unwrap_or("") + .as_bytes(); // Determine if we need auth protocol (before auth is moved) let needs_auth_protocol = !matches!(auth, Auth::NoAuthNoPriv); @@ -401,7 +438,7 @@ mod tests { #[test] fn test_snmp_client_default() { - let client = SnmpClient::default(); + let client = SnmpClient; assert!(format!("{:?}", client).contains("SnmpClient")); } diff --git a/src/snmp/device_poller.rs b/src/snmp/device_poller.rs index 1b777c7..707f920 100644 --- a/src/snmp/device_poller.rs +++ b/src/snmp/device_poller.rs @@ -1,5 +1,6 @@ use super::types::{SnmpError, SnmpResult, SnmpValue}; use super::V3Config; +use crate::secret::SecretString; use snmp2::{Oid, SyncSession}; use std::str::FromStr; use std::time::Duration; @@ -22,15 +23,27 @@ pub enum SnmpRequest { } /// Configuration for a device poller -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct DeviceConfig { pub ip: String, pub port: u16, pub version: String, - pub community: String, + pub community: SecretString, pub v3_config: Option, } +impl std::fmt::Debug for DeviceConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DeviceConfig") + .field("ip", &self.ip) + .field("port", &self.port) + .field("version", &self.version) + .field("community", &"[REDACTED]") + .field("v3_config", &self.v3_config) + .finish() + } +} + /// Per-device polling thread that maintains a persistent SNMP session pub struct DevicePoller { device_id: String, @@ -184,7 +197,12 @@ fn create_session(addr: &str, config: &DeviceConfig) -> Result SshResult { + pub async fn connect( + host: &str, + port: u32, + username: &str, + password: &SecretString, + ) -> SshResult { let config = client::Config::default(); let sh = Client; @@ -53,7 +59,7 @@ impl SshClient { .map_err(|e| SshError::ConnectionFailed(e.to_string()))?; let auth_result = session - .authenticate_password(username, password) + .authenticate_password(username, password.expose()) .await .map_err(|e| SshError::ConnectionFailed(e.to_string()))?; diff --git a/src/version.rs b/src/version.rs index ae75929..923d04d 100644 --- a/src/version.rs +++ b/src/version.rs @@ -27,7 +27,7 @@ mod tests { let version = current_version(); // Version should be in semver-like format (e.g., "0.1.0" or custom BUILD_VERSION) // At minimum, should have some content - assert!(version.len() >= 1); + assert!(!version.is_empty()); } #[test] diff --git a/src/websocket_client.rs b/src/websocket_client.rs index 09f0b58..a910411 100644 --- a/src/websocket_client.rs +++ b/src/websocket_client.rs @@ -6,6 +6,8 @@ /// /// Connection URL: {url}/socket/agent/websocket /// Authentication: Token sent in Phoenix channel join payload +use crate::secret::SecretString; +use futures::stream::SplitStream; use futures::{SinkExt, StreamExt}; use prost::Message; use std::collections::HashMap; @@ -15,6 +17,7 @@ use tokio::time::{interval, timeout, Duration}; use tokio_tungstenite::{ connect_async, tungstenite::protocol::Message as WsMessage, MaybeTlsStream, WebSocketStream, }; +use zeroize::Zeroize; /// Connection timeout for WebSocket establishment (30 seconds) const CONNECTION_TIMEOUT: Duration = Duration::from_secs(30); @@ -37,17 +40,33 @@ struct PhoenixMessage { reference: Option, } +/// Channel capacity for result backpressure. If the WebSocket write side +/// falls behind, job tasks will slow down rather than consuming unbounded memory. +const RESULT_CHANNEL_CAPACITY: usize = 1000; + +/// Channel capacity for outgoing WebSocket messages routed through the writer task. +const WS_WRITE_CHANNEL_CAPACITY: usize = 500; + /// WebSocket client for agent communication. +/// +/// The WebSocket stream is split into a read half (owned here) and a write half +/// (owned by a dedicated writer task). All outgoing messages are sent through +/// `ws_write_tx`, allowing reads and writes to proceed concurrently. pub struct AgentClient { - ws_stream: WebSocketStream>, + ws_read: SplitStream>>, + ws_write_tx: mpsc::Sender, agent_id: String, - result_tx: mpsc::UnboundedSender, - result_rx: mpsc::UnboundedReceiver, - mikrotik_result_tx: mpsc::UnboundedSender, - mikrotik_result_rx: mpsc::UnboundedReceiver, - credential_test_tx: mpsc::UnboundedSender, - credential_test_rx: mpsc::UnboundedReceiver, + result_tx: mpsc::Sender, + result_rx: mpsc::Receiver, + mikrotik_result_tx: mpsc::Sender, + mikrotik_result_rx: mpsc::Receiver, + credential_test_tx: mpsc::Sender, + credential_test_rx: mpsc::Receiver, poller_registry: PollerRegistry, + /// Counter for Phoenix transport heartbeat refs + phx_heartbeat_ref: u64, + /// Cached hostname (computed once at startup, avoids blocking /proc reads) + cached_hostname: String, } impl AgentClient { @@ -61,7 +80,7 @@ impl AgentClient { /// ```no_run /// let client = AgentClient::connect("wss://towerops.net", "token123").await?; /// ``` - pub async fn connect(url: &str, token: &str) -> Result { + pub async fn connect(url: &str, token: &SecretString) -> Result { // Strip trailing slash from base URL to avoid double slashes let base_url = url.trim_end_matches('/'); let ws_url = format!("{}/socket/agent/websocket", base_url); @@ -72,7 +91,7 @@ impl AgentClient { ); // Wrap connection in timeout to avoid hanging indefinitely on bad network - let (mut ws_stream, _) = match timeout(CONNECTION_TIMEOUT, connect_async(&ws_url)).await { + let (ws_stream, _) = match timeout(CONNECTION_TIMEOUT, connect_async(&ws_url)).await { Ok(Ok(result)) => result, Ok(Err(e)) => { tracing::error!("WebSocket connection failed: {}", e); @@ -94,27 +113,38 @@ impl AgentClient { tracing::info!("Connected to Towerops server at {}", url); let agent_id = generate_agent_id(); - let (result_tx, result_rx) = mpsc::unbounded_channel(); - let (mikrotik_result_tx, mikrotik_result_rx) = mpsc::unbounded_channel(); - let (credential_test_tx, credential_test_rx) = mpsc::unbounded_channel(); + let (result_tx, result_rx) = mpsc::channel(RESULT_CHANNEL_CAPACITY); + let (mikrotik_result_tx, mikrotik_result_rx) = mpsc::channel(RESULT_CHANNEL_CAPACITY); + let (credential_test_tx, credential_test_rx) = mpsc::channel(RESULT_CHANNEL_CAPACITY); + + // Split the WebSocket stream so reads and writes can proceed concurrently. + // The write half is owned by a dedicated writer task. + let (ws_write, ws_read) = ws_stream.split(); + let (ws_write_tx, ws_write_rx) = mpsc::channel::(WS_WRITE_CHANNEL_CAPACITY); + + tokio::spawn(ws_writer_task(ws_write, ws_write_rx)); // Join Phoenix channel with token in payload let join_msg = PhoenixMessage { topic: format!("agent:{}", agent_id), event: "phx_join".to_string(), - payload: serde_json::json!({"token": token}), + payload: serde_json::json!({"token": token.expose()}), reference: Some("1".to_string()), }; let join_text = serde_json::to_string(&join_msg)?; - ws_stream.send(WsMessage::Text(join_text.into())).await?; + ws_write_tx + .send(WsMessage::Text(join_text.into())) + .await + .map_err(|e| format!("Failed to send join message: {}", e))?; tracing::info!( "Sent channel join request with token for agent:{}", agent_id ); Ok(Self { - ws_stream, + ws_read, + ws_write_tx, agent_id, result_tx, result_rx, @@ -123,6 +153,8 @@ impl AgentClient { credential_test_tx, credential_test_rx, poller_registry: PollerRegistry::new(), + phx_heartbeat_ref: 0, + cached_hostname: get_hostname(), }) } @@ -136,6 +168,7 @@ impl AgentClient { /// - Graceful shutdown on SIGTERM pub async fn run(&mut self, mut shutdown_rx: watch::Receiver) -> Result<()> { let mut heartbeat_interval = interval(Duration::from_secs(60)); + let mut phx_heartbeat_interval = interval(Duration::from_secs(25)); loop { tokio::select! { @@ -143,16 +176,14 @@ impl AgentClient { _ = shutdown_rx.changed() => { if *shutdown_rx.borrow() { tracing::info!("Shutdown signal received, closing WebSocket connection gracefully"); - // Send close frame to server - if let Err(e) = self.ws_stream.close(None).await { - tracing::warn!("Error sending WebSocket close frame: {}", e); - } + // Send close frame through the writer task + let _ = self.ws_write_tx.send(WsMessage::Close(None)).await; break Ok(()); } } // Receive messages from server - msg = self.ws_stream.next() => { + msg = self.ws_read.next() => { match msg { Some(Ok(WsMessage::Binary(data))) => { if let Err(e) = self.handle_message(&data).await { @@ -215,6 +246,13 @@ impl AgentClient { tracing::debug!("Active device pollers: {}", count); } } + + // Send Phoenix transport heartbeats to keep connection alive + _ = phx_heartbeat_interval.tick() => { + if let Err(e) = self.send_phx_heartbeat().await { + tracing::error!("Error sending Phoenix heartbeat: {}", e); + } + } } } } @@ -325,7 +363,7 @@ impl AgentClient { async fn send_heartbeat(&mut self) -> Result<()> { let heartbeat = AgentHeartbeat { version: env!("CARGO_PKG_VERSION").to_string(), - hostname: get_hostname(), + hostname: self.cached_hostname.clone(), uptime_seconds: get_uptime_seconds(), ip_address: get_local_ip().unwrap_or_else(|| "unknown".to_string()), }; @@ -341,12 +379,42 @@ impl AgentClient { }; let text = serde_json::to_string(&msg)?; - self.ws_stream.send(WsMessage::Text(text.into())).await?; + self.ws_write_tx + .send(WsMessage::Text(text.into())) + .await + .map_err(|e| format!("Writer task closed: {}", e))?; tracing::debug!("Sent heartbeat"); Ok(()) } + /// Send Phoenix transport heartbeat to keep the WebSocket connection alive. + /// + /// This is separate from the application heartbeat. Phoenix's transport layer + /// expects periodic messages on the "phoenix" topic to detect dead connections. + async fn send_phx_heartbeat(&mut self) -> Result<()> { + self.phx_heartbeat_ref += 1; + + let msg = PhoenixMessage { + topic: "phoenix".to_string(), + event: "heartbeat".to_string(), + payload: serde_json::json!({}), + reference: Some(self.phx_heartbeat_ref.to_string()), + }; + + let text = serde_json::to_string(&msg)?; + self.ws_write_tx + .send(WsMessage::Text(text.into())) + .await + .map_err(|e| format!("Writer task closed: {}", e))?; + + tracing::debug!( + "Sent Phoenix transport heartbeat (ref: {})", + self.phx_heartbeat_ref + ); + Ok(()) + } + /// Send SNMP results to server. async fn send_snmp_result(&mut self, result: SnmpResult) -> Result<()> { let binary = result.encode_to_vec(); @@ -359,7 +427,10 @@ impl AgentClient { }; let text = serde_json::to_string(&msg)?; - self.ws_stream.send(WsMessage::Text(text.into())).await?; + self.ws_write_tx + .send(WsMessage::Text(text.into())) + .await + .map_err(|e| format!("Writer task closed: {}", e))?; tracing::debug!("Sent SNMP result for device {}", result.device_id); Ok(()) @@ -377,7 +448,10 @@ impl AgentClient { }; let text = serde_json::to_string(&msg)?; - self.ws_stream.send(WsMessage::Text(text.into())).await?; + self.ws_write_tx + .send(WsMessage::Text(text.into())) + .await + .map_err(|e| format!("Writer task closed: {}", e))?; tracing::debug!( "Sent MikroTik result for device {} (job: {})", @@ -399,7 +473,10 @@ impl AgentClient { }; let text = serde_json::to_string(&msg)?; - self.ws_stream.send(WsMessage::Text(text.into())).await?; + self.ws_write_tx + .send(WsMessage::Text(text.into())) + .await + .map_err(|e| format!("Writer task closed: {}", e))?; tracing::info!( "Sent credential test result (test_id: {}, success: {})", @@ -410,34 +487,55 @@ impl AgentClient { } } -/// Redact SNMP community string for logging, showing first 2 chars only -fn redact_community(community: &str) -> String { - if community.is_empty() { - return "**".to_string(); +/// Dedicated writer task that owns the WebSocket write half. +/// +/// All outgoing messages are funnelled through an mpsc channel, allowing the +/// main event loop to continue reading while writes are in progress. +async fn ws_writer_task( + mut ws_sink: futures::stream::SplitSink>, WsMessage>, + mut rx: mpsc::Receiver, +) { + while let Some(msg) = rx.recv().await { + let is_close = matches!(msg, WsMessage::Close(_)); + if let Err(e) = ws_sink.send(msg).await { + tracing::error!("WebSocket write error: {}", e); + break; + } + if is_close { + break; + } + } + tracing::debug!("WebSocket writer task stopped"); +} + +/// Redact SNMP community string for logging. +fn redact_community(community: &str) -> &'static str { + if community.is_empty() { + "(empty)" + } else { + "***" } - let visible = std::cmp::min(community.len(), 2); - format!("{}**", &community[..visible]) } /// Execute an SNMP job and collect results. async fn execute_snmp_job( job: AgentJob, - result_tx: mpsc::UnboundedSender, + result_tx: mpsc::Sender, poller_registry: PollerRegistry, ) -> Result<()> { - let snmp_device = job.snmp_device.ok_or("Job missing SNMP device info")?; + let mut snmp_device = job.snmp_device.ok_or("Job missing SNMP device info")?; // Build v3 config if version is "3" let v3_config = if snmp_device.version == "3" { let config = crate::snmp::V3Config { username: snmp_device.v3_username.clone(), auth_password: if !snmp_device.v3_auth_password.is_empty() { - Some(snmp_device.v3_auth_password.clone()) + Some(SecretString::new(snmp_device.v3_auth_password.clone())) } else { None }, priv_password: if !snmp_device.v3_priv_password.is_empty() { - Some(snmp_device.v3_priv_password.clone()) + Some(SecretString::new(snmp_device.v3_priv_password.clone())) } else { None }, @@ -476,10 +574,15 @@ async fn execute_snmp_job( ip: snmp_device.ip.clone(), port: snmp_device.port as u16, version: snmp_device.version.clone(), - community: snmp_device.community.clone(), + community: SecretString::new(snmp_device.community.clone()), v3_config, }; + // Zeroize credentials in protobuf message after extraction + snmp_device.community.zeroize(); + snmp_device.v3_auth_password.zeroize(); + snmp_device.v3_priv_password.zeroize(); + let poller = poller_registry.get_or_create(job.device_id.clone(), device_config); let mut oid_values: HashMap = HashMap::new(); @@ -554,7 +657,7 @@ async fn execute_snmp_job( ); // Send result back to main client task - if let Err(e) = result_tx.send(result) { + if let Err(e) = result_tx.send(result).await { tracing::warn!( "Failed to send SNMP result for job {}: channel closed (connection may have dropped)", job.job_id @@ -571,9 +674,9 @@ async fn execute_snmp_job( /// Returns success with system description or failure with error message. async fn execute_credential_test( job: AgentJob, - result_tx: mpsc::UnboundedSender, + result_tx: mpsc::Sender, ) -> Result<()> { - let snmp_device = job.snmp_device.ok_or("Job missing SNMP device info")?; + let mut snmp_device = job.snmp_device.ok_or("Job missing SNMP device info")?; tracing::info!( "Testing SNMP credentials for {}:{} (version: {})", @@ -591,12 +694,12 @@ async fn execute_credential_test( Some(crate::snmp::V3Config { username: snmp_device.v3_username.clone(), auth_password: if !snmp_device.v3_auth_password.is_empty() { - Some(snmp_device.v3_auth_password.clone()) + Some(SecretString::new(snmp_device.v3_auth_password.clone())) } else { None }, priv_password: if !snmp_device.v3_priv_password.is_empty() { - Some(snmp_device.v3_priv_password.clone()) + Some(SecretString::new(snmp_device.v3_priv_password.clone())) } else { None }, @@ -616,6 +719,11 @@ async fn execute_credential_test( None }; + // Zeroize credentials in protobuf message after extraction + snmp_device.community.zeroize(); + snmp_device.v3_auth_password.zeroize(); + snmp_device.v3_priv_password.zeroize(); + // Create a temporary SNMP client for testing (don't use persistent poller) let snmp_client = crate::snmp::SnmpClient::new(); @@ -660,7 +768,7 @@ async fn execute_credential_test( }; // Send result back to main client task - if let Err(e) = result_tx.send(result) { + if let Err(e) = result_tx.send(result).await { tracing::warn!( "Failed to send credential test result for job {}: channel closed", job.job_id @@ -674,9 +782,9 @@ async fn execute_credential_test( /// Execute a MikroTik API job and collect results. async fn execute_mikrotik_job( job: AgentJob, - result_tx: mpsc::UnboundedSender, + result_tx: mpsc::Sender, ) -> Result<()> { - use crate::mikrotik::{MikrotikClient, SecretString}; + use crate::mikrotik::MikrotikClient; let timestamp = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH)? @@ -692,7 +800,7 @@ async fn execute_mikrotik_job( return execute_mikrotik_backup_via_ssh(job, mikrotik_device, result_tx, timestamp).await; } - let mikrotik_device = job + let mut mikrotik_device = job .mikrotik_device .ok_or("Job missing MikroTik device info")?; @@ -726,7 +834,7 @@ async fn execute_mikrotik_job( error: format!("Connection failed: {}", e), timestamp, }; - let _ = result_tx.send(result); + let _ = result_tx.send(result).await; return Err(format!("MikroTik connection failed: {}", e).into()); } } @@ -748,12 +856,15 @@ async fn execute_mikrotik_job( error: format!("Connection failed: {}", e), timestamp, }; - let _ = result_tx.send(result); + let _ = result_tx.send(result).await; return Err(format!("MikroTik connection failed: {}", e).into()); } } }; + // Zeroize credentials in protobuf message after extraction + mikrotik_device.password.zeroize(); + // Execute each command and collect results let mut all_sentences = Vec::new(); let mut error_message = String::new(); @@ -847,7 +958,7 @@ async fn execute_mikrotik_job( ); let job_id_for_error = result.job_id.clone(); - if let Err(e) = result_tx.send(result) { + if let Err(e) = result_tx.send(result).await { tracing::warn!( "Failed to send MikroTik result for job {}: channel closed", job_id_for_error @@ -861,8 +972,8 @@ async fn execute_mikrotik_job( /// Execute a MikroTik backup job via SSH (because /export doesn't work via API). async fn execute_mikrotik_backup_via_ssh( job: AgentJob, - mikrotik_device: crate::proto::agent::MikrotikDevice, - result_tx: mpsc::UnboundedSender, + mut mikrotik_device: crate::proto::agent::MikrotikDevice, + result_tx: mpsc::Sender, timestamp: i64, ) -> Result<()> { use crate::ssh::SshClient; @@ -875,12 +986,14 @@ async fn execute_mikrotik_backup_via_ssh( job.job_id ); + let password = SecretString::new(mikrotik_device.password.clone()); + // Connect via SSH let mut ssh_client = match SshClient::connect( &mikrotik_device.ip, mikrotik_device.ssh_port, &mikrotik_device.username, - &mikrotik_device.password, + &password, ) .await { @@ -895,7 +1008,7 @@ async fn execute_mikrotik_backup_via_ssh( error: error_msg, timestamp, }; - let _ = result_tx.send(result); + let _ = result_tx.send(result).await; return Err(format!("SSH connection failed: {}", e).into()); } }; @@ -913,7 +1026,7 @@ async fn execute_mikrotik_backup_via_ssh( error: error_msg, timestamp, }; - let _ = result_tx.send(result); + let _ = result_tx.send(result).await; let _ = ssh_client.close().await; return Err(format!("SSH command failed: {}", e).into()); } @@ -922,6 +1035,9 @@ async fn execute_mikrotik_backup_via_ssh( // Close SSH connection let _ = ssh_client.close().await; + // Zeroize credentials in protobuf message after use + mikrotik_device.password.zeroize(); + tracing::info!( "Backup completed: {} bytes, {} lines", config.len(), @@ -947,7 +1063,7 @@ async fn execute_mikrotik_backup_via_ssh( result.job_id ); - if let Err(e) = result_tx.send(result) { + if let Err(e) = result_tx.send(result).await { tracing::warn!( "Failed to send MikroTik backup result for job {}: channel closed", job_id_for_log @@ -1249,27 +1365,27 @@ mod tests { #[test] fn test_redact_community_normal() { - assert_eq!(redact_community("public"), "pu**"); + assert_eq!(redact_community("public"), "***"); } #[test] fn test_redact_community_short() { - assert_eq!(redact_community("ab"), "ab**"); - assert_eq!(redact_community("a"), "a**"); + assert_eq!(redact_community("ab"), "***"); + assert_eq!(redact_community("a"), "***"); } #[test] fn test_redact_community_empty() { - assert_eq!(redact_community(""), "**"); + assert_eq!(redact_community(""), "(empty)"); } #[test] fn test_redact_community_three_chars() { - assert_eq!(redact_community("abc"), "ab**"); + assert_eq!(redact_community("abc"), "***"); } #[test] fn test_redact_community_long() { - assert_eq!(redact_community("mysecretcommunity"), "my**"); + assert_eq!(redact_community("mysecretcommunity"), "***"); } }