use super::event::{AuditEvent, QueryLogEntry, SecurityEventType};
use once_cell::sync::Lazy;
use std::collections::HashSet;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::RwLock;
use tracing::{info, trace};
pub static AUDIT_LOGGER: Lazy<AuditLogger> = Lazy::new(AuditLogger::new);
pub struct AuditLogger {
enabled_events: RwLock<HashSet<SecurityEventType>>,
queries_logged: AtomicU64,
queries_sampled_out: AtomicU64,
security_events_logged: AtomicU64,
}
impl AuditLogger {
pub fn new() -> Self {
Self {
enabled_events: RwLock::new(HashSet::new()),
queries_logged: AtomicU64::new(0),
queries_sampled_out: AtomicU64::new(0),
security_events_logged: AtomicU64::new(0),
}
}
pub fn log_query(&self, entry: QueryLogEntry) {
let subscribers = super::event_bus::publish_query(entry);
if subscribers > 0 {
self.queries_logged.fetch_add(1, Ordering::Relaxed);
trace!(subscribers, "Query published to event bus");
}
}
pub async fn log_security(&self, event: AuditEvent) {
if let AuditEvent::Security { event_type, .. } = &event {
let enabled = self.enabled_events.read().await;
if !enabled.is_empty() && !enabled.contains(event_type) {
return;
}
}
let subscribers = super::event_bus::publish_security(event);
if subscribers > 0 {
self.security_events_logged.fetch_add(1, Ordering::Relaxed);
trace!(subscribers, "Security event published to event bus");
}
}
pub async fn log_security_event(
&self,
event_type: SecurityEventType,
message: impl Into<String>,
client_ip: Option<std::net::IpAddr>,
qname: Option<String>,
) {
let event = AuditEvent::security_with_client(event_type, message, client_ip, qname);
self.log_security(event).await;
}
pub fn stats(&self) -> AuditStats {
let bus_stats = super::event_bus::event_bus()
.map(|b| b.stats())
.unwrap_or_default();
AuditStats {
queries_logged: self.queries_logged.load(Ordering::Relaxed),
queries_sampled_out: self.queries_sampled_out.load(Ordering::Relaxed),
security_events_logged: self.security_events_logged.load(Ordering::Relaxed),
events_dropped: bus_stats.events_dropped,
active_subscribers: bus_stats.active_subscribers,
}
}
pub async fn shutdown(&self) {
info!(
queries = self.queries_logged.load(Ordering::Relaxed),
sampled_out = self.queries_sampled_out.load(Ordering::Relaxed),
security_events = self.security_events_logged.load(Ordering::Relaxed),
"Audit logger shutdown"
);
}
}
impl Default for AuditLogger {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Default)]
pub struct AuditStats {
pub queries_logged: u64,
pub queries_sampled_out: u64,
pub security_events_logged: u64,
pub events_dropped: u64,
pub active_subscribers: usize,
}