fix: Address security and quality issues in audit infrastructure

Author: inureyes

src/server/audit/event.rs

@@ -59,6 +59,19 @@ pub struct AuditEvent {
     pub result: EventResult,
 
     /// Additional details about the event
+    ///
+    /// # Security Warning
+    ///
+    /// This field may contain sensitive information such as:
+    /// - Error messages with file paths or system information
+    /// - Command arguments that may include passwords or tokens
+    /// - User-supplied data that hasn't been sanitized
+    ///
+    /// When implementing exporters, ensure this field is handled securely:
+    /// - Apply appropriate access controls to audit logs
+    /// - Consider redacting or filtering sensitive patterns
+    /// - Use encryption when transmitting over networks
+    /// - Comply with data retention and privacy policies
     pub details: Option<String>,
 
     /// Protocol used (ssh, sftp, scp)
@@ -217,6 +230,11 @@ impl AuditEvent {
     }
 
     /// Set additional details.
+    ///
+    /// # Security Warning
+    ///
+    /// Be cautious when including sensitive information in this field.
+    /// See the `details` field documentation for security considerations.
     pub fn with_details(mut self, details: String) -> Self {
         self.details = Some(details);
         self
@@ -229,11 +247,9 @@ impl AuditEvent {
     }
 }
 
-impl Default for AuditEvent {
-    fn default() -> Self {
-        Self::new(EventType::CommandExecuted, String::new(), String::new())
-    }
-}
+// Note: Default implementation removed as it creates sentinel values with empty
+// user and session_id fields, which are semantically invalid for audit events.
+// Use AuditEvent::new() to create audit events with meaningful values.
 
 #[cfg(test)]
 mod tests {
@@ -319,15 +335,6 @@ mod tests {
         assert_eq!(deserialized.protocol, event.protocol);
     }
 
-    #[test]
-    fn test_default_event() {
-        let event = AuditEvent::default();
-        assert_eq!(event.event_type, EventType::CommandExecuted);
-        assert_eq!(event.result, EventResult::Success);
-        assert!(event.user.is_empty());
-        assert!(event.session_id.is_empty());
-    }
-
     #[test]
     fn test_event_with_dest_path() {
         let event = AuditEvent::new(

src/server/audit/exporter.rs

@@ -46,14 +46,14 @@ pub trait AuditExporter: Send + Sync {
     ///
     /// # Arguments
     ///
-    /// * `events` - Vector of audit events to export
+    /// * `events` - Slice of audit events to export
     ///
     /// # Errors
     ///
     /// Returns an error if any event fails to export.
-    async fn export_batch(&self, events: Vec<AuditEvent>) -> Result<()> {
+    async fn export_batch(&self, events: &[AuditEvent]) -> Result<()> {
         for event in events {
-            self.export(event).await?;
+            self.export(event.clone()).await?;
         }
         Ok(())
     }
@@ -98,7 +98,7 @@ impl AuditExporter for NullExporter {
         Ok(())
     }
 
-    async fn export_batch(&self, _events: Vec<AuditEvent>) -> Result<()> {
+    async fn export_batch(&self, _events: &[AuditEvent]) -> Result<()> {
         // Discard all events
         Ok(())
     }
@@ -149,7 +149,7 @@ mod tests {
             .with_result(EventResult::Failure),
         ];
 
-        let result = exporter.export_batch(events).await;
+        let result = exporter.export_batch(&events).await;
         assert!(result.is_ok());
     }
 
@@ -185,7 +185,7 @@ mod tests {
             "bob".to_string(),
             "session-456".to_string(),
         )];
-        exporter.export_batch(events).await.unwrap();
+        exporter.export_batch(&events).await.unwrap();
 
         // Flush and close
         exporter.flush().await.unwrap();

src/server/audit/mod.rs

@@ -53,6 +53,7 @@ use anyhow::Result;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
 
 pub use event::{AuditEvent, EventResult, EventType};
 pub use exporter::{AuditExporter, NullExporter};
@@ -128,19 +129,34 @@ impl AuditConfig {
     }
 
     /// Set the buffer size.
+    ///
+    /// # Panics
+    ///
+    /// Panics if size is 0.
     pub fn with_buffer_size(mut self, size: usize) -> Self {
+        assert!(size >= 1, "buffer_size must be at least 1");
         self.buffer_size = size;
         self
     }
 
     /// Set the batch size.
+    ///
+    /// # Panics
+    ///
+    /// Panics if size is 0.
     pub fn with_batch_size(mut self, size: usize) -> Self {
+        assert!(size >= 1, "batch_size must be at least 1");
         self.batch_size = size;
         self
     }
 
     /// Set the flush interval.
+    ///
+    /// # Panics
+    ///
+    /// Panics if secs is 0.
     pub fn with_flush_interval(mut self, secs: u64) -> Self {
+        assert!(secs >= 1, "flush_interval_secs must be at least 1");
         self.flush_interval_secs = secs;
         self
     }
@@ -166,6 +182,9 @@ pub struct AuditManager {
 
     /// Whether audit logging is enabled
     enabled: bool,
+
+    /// Handle to the background worker task
+    worker_handle: Option<JoinHandle<()>>,
 }
 
 impl AuditManager {
@@ -208,23 +227,26 @@ impl AuditManager {
             exporters.push(exporter);
         }
 
-        let manager = Self {
-            exporters: exporters.clone(),
-            sender,
-            enabled: config.enabled,
-        };
-
         // Start background worker
-        if config.enabled {
+        let worker_handle = if config.enabled {
             let batch_size = config.batch_size;
             let flush_interval = Duration::from_secs(config.flush_interval_secs);
-            tokio::spawn(Self::worker(
+            Some(tokio::spawn(Self::worker(
                 receiver,
-                exporters,
+                exporters.clone(),
                 batch_size,
                 flush_interval,
-            ));
-        }
+            )))
+        } else {
+            None
+        };
+
+        let manager = Self {
+            exporters,
+            sender,
+            enabled: config.enabled,
+            worker_handle,
+        };
 
         Ok(manager)
     }
@@ -263,26 +285,32 @@ impl AuditManager {
 
         loop {
             tokio::select! {
-                Some(event) = receiver.recv() => {
-                    buffer.push(event);
-
-                    // Flush if buffer is full
-                    if buffer.len() >= batch_size {
-                        Self::flush_buffer(&exporters, &mut buffer).await;
+                biased;
+
+                event_opt = receiver.recv() => {
+                    match event_opt {
+                        Some(event) => {
+                            buffer.push(event);
+
+                            // Flush if buffer is full
+                            if buffer.len() >= batch_size {
+                                Self::flush_buffer(&exporters, &mut buffer).await;
+                            }
+                        }
+                        None => {
+                            // Channel closed, flush remaining events and exit
+                            if !buffer.is_empty() {
+                                Self::flush_buffer(&exporters, &mut buffer).await;
+                            }
+                            break;
+                        }
                     }
                 }
                 _ = flush_timer.tick() => {
                     if !buffer.is_empty() {
                         Self::flush_buffer(&exporters, &mut buffer).await;
                     }
                 }
-                else => {
-                    // Channel closed, flush remaining events and exit
-                    if !buffer.is_empty() {
-                        Self::flush_buffer(&exporters, &mut buffer).await;
-                    }
-                    break;
-                }
             }
         }
 
@@ -297,7 +325,7 @@ impl AuditManager {
     /// Flush the event buffer to all exporters.
     async fn flush_buffer(exporters: &[Arc<dyn AuditExporter>], buffer: &mut Vec<AuditEvent>) {
         for exporter in exporters {
-            if let Err(e) = exporter.export_batch(buffer.clone()).await {
+            if let Err(e) = exporter.export_batch(buffer).await {
                 tracing::error!("Audit export failed: {}", e);
             }
         }
@@ -319,6 +347,33 @@ impl AuditManager {
     pub fn is_enabled(&self) -> bool {
         self.enabled
     }
+
+    /// Gracefully shut down the audit manager.
+    ///
+    /// This method:
+    /// 1. Drops the sender to signal the worker to stop accepting new events
+    /// 2. Waits for the worker to finish processing buffered events
+    /// 3. Ensures all exporters are properly closed
+    ///
+    /// After calling this method, the AuditManager should not be used.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the worker task panicked or if there was an issue
+    /// waiting for the worker to complete.
+    pub async fn shutdown(mut self) -> Result<()> {
+        // Drop the sender to signal the worker to exit
+        drop(self.sender);
+
+        // Wait for the worker to finish
+        if let Some(handle) = self.worker_handle.take() {
+            handle
+                .await
+                .map_err(|e| anyhow::anyhow!("Worker task panicked: {}", e))?;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -448,4 +503,63 @@ mod tests {
         // Wait for flush interval
         tokio::time::sleep(Duration::from_millis(1100)).await;
     }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_audit_manager_shutdown() {
+        let config = AuditConfig::new()
+            .with_enabled(true)
+            .with_buffer_size(10)
+            .with_batch_size(5)
+            .with_flush_interval(1);
+
+        let manager = AuditManager::new(&config).unwrap();
+
+        // Send some events
+        for i in 0..3 {
+            let event = AuditEvent::new(
+                EventType::FileUploaded,
+                format!("user{}", i),
+                format!("session-{}", i),
+            );
+            manager.log(event).await;
+        }
+
+        // Give a small amount of time for events to be queued
+        tokio::time::sleep(Duration::from_millis(50)).await;
+
+        // Shutdown should wait for all events to be processed
+        let result = tokio::time::timeout(Duration::from_secs(10), manager.shutdown()).await;
+        assert!(result.is_ok(), "Shutdown timed out");
+        assert!(result.unwrap().is_ok(), "Shutdown failed");
+    }
+
+    #[test]
+    #[should_panic(expected = "buffer_size must be at least 1")]
+    fn test_audit_config_invalid_buffer_size() {
+        let _config = AuditConfig::new().with_buffer_size(0);
+    }
+
+    #[test]
+    #[should_panic(expected = "batch_size must be at least 1")]
+    fn test_audit_config_invalid_batch_size() {
+        let _config = AuditConfig::new().with_batch_size(0);
+    }
+
+    #[test]
+    #[should_panic(expected = "flush_interval_secs must be at least 1")]
+    fn test_audit_config_invalid_flush_interval() {
+        let _config = AuditConfig::new().with_flush_interval(0);
+    }
+
+    #[test]
+    fn test_audit_config_valid_minimum_values() {
+        let config = AuditConfig::new()
+            .with_buffer_size(1)
+            .with_batch_size(1)
+            .with_flush_interval(1);
+
+        assert_eq!(config.buffer_size, 1);
+        assert_eq!(config.batch_size, 1);
+        assert_eq!(config.flush_interval_secs, 1);
+    }
 }