graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Security hardening and performance improvements

Browse source
This commit is contained in:
Graham McIntire 2026-02-04 16:50:36 -06:00
parent 91913bd6e4
commit df334e4db8
No known key found for this signature in database
11 changed files with 367 additions and 163 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -2606,6 +2606,7 @@ dependencies = [
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
"zeroize",
]
[[package]]

1
Cargo.toml
View file

@ -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"

13
src/main.rs
View file

@ -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;

75
src/mikrotik/types.rs
View file

@ -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<String>) -> 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());
}
}

87
src/secret.rs Normal file
View file

@ -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<String>) -> 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());
}
}

53
src/snmp/client.rs
View file

@ -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<String>,
pub priv_password: Option<String>,
pub auth_password: Option<SecretString>,
pub priv_password: Option<SecretString>,
pub auth_protocol: Option<String>,
pub priv_protocol: Option<String>,
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,7 +303,11 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult<SyncSession> {
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(|| {
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,14 +319,22 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult<SyncSession> {
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(|| {
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(|| {
let priv_pass = config
.priv_password
.as_ref()
.map(|s| s.expose())
.ok_or_else(|| {
SnmpError::RequestFailed("Priv password required for authPriv".into())
})?;
@ -321,7 +353,12 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult<SyncSession> {
};
// 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"));
}

34
src/snmp/device_poller.rs
View file

@ -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<V3Config>,
}
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<SyncSession, Stri
.ok_or("v3_config required for SNMPv3")?;
create_v3_session(addr, timeout, v3_config)
} else {
create_v1v2c_session(addr, config.community.as_bytes(), timeout, version_num)
create_v1v2c_session(
addr,
config.community.expose().as_bytes(),
timeout,
version_num,
)
}
}
@ -227,7 +245,8 @@ fn create_v3_session(
)?;
let priv_pass = config
.priv_password
.as_deref()
.as_ref()
.map(|s| s.expose())
.ok_or("Priv password required for authPriv")?;
Auth::AuthPriv {
@ -244,7 +263,12 @@ fn create_v3_session(
};
let username = config.username.as_bytes();
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();
let needs_auth_protocol = !matches!(auth, Auth::NoAuthNoPriv);
let mut security = Security::new(username, auth_password).with_auth(auth);

6
src/snmp/poller_registry.rs
View file

@ -99,7 +99,7 @@ impl Default for PollerRegistry {
#[cfg(test)]
mod tests {
use super::*;
use crate::snmp::V3Config;
use crate::secret::SecretString;
#[test]
fn test_registry_remove() {
@ -110,7 +110,7 @@ mod tests {
ip: "127.0.0.1".to_string(),
port: 161,
version: "2c".to_string(),
community: "public".to_string(),
community: SecretString::new("public"),
v3_config: None,
};
@ -130,7 +130,7 @@ mod tests {
ip: "192.168.1.1".to_string(),
port: 161,
version: "2c".to_string(),
community: "public".to_string(),
community: SecretString::new("public"),
v3_config: None,
};

10
src/ssh/client.rs
View file

@ -1,3 +1,4 @@
use crate::secret::SecretString;
use russh::client;
use russh::keys::PublicKey;
use std::future::Future;
@ -42,7 +43,12 @@ pub struct SshClient {
impl SshClient {
/// Connect to an SSH server and authenticate with password
pub async fn connect(host: &str, port: u32, username: &str, password: &str) -> SshResult<Self> {
pub async fn connect(
host: &str,
port: u32,
username: &str,
password: &SecretString,
) -> SshResult<Self> {
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()))?;

2
src/version.rs
View file

@ -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]

236
src/websocket_client.rs
View file

@ -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<String>,
}
/// 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<MaybeTlsStream<TcpStream>>,
ws_read: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
ws_write_tx: mpsc::Sender<WsMessage>,
agent_id: String,
result_tx: mpsc::UnboundedSender<SnmpResult>,
result_rx: mpsc::UnboundedReceiver<SnmpResult>,
mikrotik_result_tx: mpsc::UnboundedSender<MikrotikResult>,
mikrotik_result_rx: mpsc::UnboundedReceiver<MikrotikResult>,
credential_test_tx: mpsc::UnboundedSender<CredentialTestResult>,
credential_test_rx: mpsc::UnboundedReceiver<CredentialTestResult>,
result_tx: mpsc::Sender<SnmpResult>,
result_rx: mpsc::Receiver<SnmpResult>,
mikrotik_result_tx: mpsc::Sender<MikrotikResult>,
mikrotik_result_rx: mpsc::Receiver<MikrotikResult>,
credential_test_tx: mpsc::Sender<CredentialTestResult>,
credential_test_rx: mpsc::Receiver<CredentialTestResult>,
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<Self> {
pub async fn connect(url: &str, token: &SecretString) -> Result<Self> {
// 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::<WsMessage>(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<bool>) -> 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<WebSocketStream<MaybeTlsStream<TcpStream>>, WsMessage>,
mut rx: mpsc::Receiver<WsMessage>,
) {
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<SnmpResult>,
result_tx: mpsc::Sender<SnmpResult>,
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<String, String> = 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<CredentialTestResult>,
result_tx: mpsc::Sender<CredentialTestResult>,
) -> 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<MikrotikResult>,
result_tx: mpsc::Sender<MikrotikResult>,
) -> 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<MikrotikResult>,
mut mikrotik_device: crate::proto::agent::MikrotikDevice,
result_tx: mpsc::Sender<MikrotikResult>,
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"), "***");
}
}