use crate::events::emitter::EventEmitter;
use crate::events::types::{Event, EventHandler};
#[cfg(feature = "webhook")]
use crate::webhook_validator::validate_webhook_url;
use log::{debug, error, info, warn};
use std::sync::Arc;
use tokio::sync::broadcast;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
pub struct EventDispatcher {
emitter: EventEmitter,
handlers: Arc<RwLock<Vec<Arc<dyn EventHandler>>>>,
dispatch_handle: Arc<RwLock<Option<JoinHandle<()>>>>,
webhook_urls: Arc<RwLock<Vec<String>>>,
}
impl EventDispatcher {
pub fn new(emitter: EventEmitter) -> Self {
let webhook_urls = emitter.config().webhook_urls.clone();
Self {
emitter,
handlers: Arc::new(RwLock::new(Vec::new())),
dispatch_handle: Arc::new(RwLock::new(None)),
webhook_urls: Arc::new(RwLock::new(webhook_urls)),
}
}
pub async fn register_handler(&self, handler: Arc<dyn EventHandler>) {
let mut handlers = self.handlers.write().await;
info!("Registering event handler: {}", handler.name());
handlers.push(handler);
}
pub async fn clear_handlers(&self) {
let mut handlers = self.handlers.write().await;
handlers.clear();
info!("All event handlers cleared");
}
pub async fn add_webhook_url(&self, url: String) {
let mut urls = self.webhook_urls.write().await;
if !urls.contains(&url) {
info!("Adding webhook URL: {}", url);
urls.push(url);
}
}
pub async fn remove_webhook_url(&self, url: &str) {
let mut urls = self.webhook_urls.write().await;
urls.retain(|u| u != url);
info!("Removed webhook URL: {}", url);
}
pub async fn start(&self) {
let mut handle_guard = self.dispatch_handle.write().await;
if handle_guard.is_some() {
warn!("EventDispatcher is already running");
return;
}
let handlers = self.handlers.clone();
let webhook_urls = self.webhook_urls.clone();
let mut receiver = self.emitter.subscribe();
let handle = tokio::spawn(async move {
info!("EventDispatcher started");
loop {
match receiver.recv().await {
Ok(event) => {
debug!(
"Received event: name={}, severity={}",
event.name(),
event.severity()
);
let handlers_snapshot = handlers.read().await.clone();
for handler in &handlers_snapshot {
if handler.accepts(&event.event_type) {
if let Err(e) = handler.handle(&event).await {
error!(
"Handler {} failed to process event {}: {}",
handler.name(),
event.name(),
e
);
}
}
}
let urls = webhook_urls.read().await.clone();
for url in &urls {
if let Err(e) = send_webhook(url, &event).await {
error!("Failed to send webhook to {}: {}", url, e);
}
}
}
Err(broadcast::error::RecvError::Lagged(n)) => {
warn!("EventDispatcher lagged behind, dropped {} events", n);
}
Err(broadcast::error::RecvError::Closed) => {
info!("EventDispatcher channel closed, stopping");
break;
}
}
}
});
*handle_guard = Some(handle);
}
pub async fn stop(&self) {
let mut handle_guard = self.dispatch_handle.write().await;
if let Some(handle) = handle_guard.take() {
handle.abort();
info!("EventDispatcher stopped");
}
}
pub fn emitter(&self) -> &EventEmitter {
&self.emitter
}
pub async fn handler_count(&self) -> usize {
self.handlers.read().await.len()
}
}
impl Drop for EventDispatcher {
fn drop(&mut self) {
if let Ok(mut handle_guard) = self.dispatch_handle.try_write() {
if let Some(handle) = handle_guard.take() {
handle.abort();
debug!("EventDispatcher dispatch task aborted on drop");
}
}
}
}
#[cfg(feature = "webhook")]
async fn send_webhook(
url: &str,
event: &Event,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
validate_webhook_url(url, !cfg!(debug_assertions))
.map_err(Box::<dyn std::error::Error + Send + Sync>::from)?;
let client = reqwest::Client::new();
let response = client
.post(url)
.json(event)
.timeout(std::time::Duration::from_secs(5))
.send()
.await?;
if response.status().is_success() {
debug!("Webhook sent successfully to {}", url);
Ok(())
} else {
Err(format!("Webhook returned error status: {}", response.status()).into())
}
}
#[cfg(not(feature = "webhook"))]
async fn send_webhook(
_url: &str,
_event: &Event,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Err("Webhook feature is not enabled. Enable 'webhook' feature to use webhooks.".into())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::events::types::EventType;
use std::sync::atomic::{AtomicU64, Ordering};
struct TestHandler {
call_count: Arc<AtomicU64>,
}
#[async_trait::async_trait]
impl EventHandler for TestHandler {
async fn handle(
&self,
_event: &Event,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
self.call_count.fetch_add(1, Ordering::SeqCst);
Ok(())
}
fn name(&self) -> &str {
"TestHandler"
}
}
#[tokio::test]
async fn test_dispatcher_register_handler() {
let emitter = EventEmitter::with_default_config();
let dispatcher = EventDispatcher::new(emitter);
assert_eq!(dispatcher.handler_count().await, 0);
let handler = Arc::new(TestHandler {
call_count: Arc::new(AtomicU64::new(0)),
});
dispatcher.register_handler(handler).await;
assert_eq!(dispatcher.handler_count().await, 1);
}
#[tokio::test]
async fn test_dispatcher_clear_handlers() {
let emitter = EventEmitter::with_default_config();
let dispatcher = EventDispatcher::new(emitter);
let handler1 = Arc::new(TestHandler {
call_count: Arc::new(AtomicU64::new(0)),
});
let handler2 = Arc::new(TestHandler {
call_count: Arc::new(AtomicU64::new(0)),
});
dispatcher.register_handler(handler1).await;
dispatcher.register_handler(handler2).await;
assert_eq!(dispatcher.handler_count().await, 2);
dispatcher.clear_handlers().await;
assert_eq!(dispatcher.handler_count().await, 0);
}
#[tokio::test]
async fn test_dispatcher_webhook_urls() {
let emitter = EventEmitter::with_default_config();
let dispatcher = EventDispatcher::new(emitter);
dispatcher
.add_webhook_url("https://example.com/hook1".to_string())
.await;
dispatcher
.add_webhook_url("https://example.com/hook2".to_string())
.await;
dispatcher
.add_webhook_url("https://example.com/hook1".to_string())
.await;
let urls = dispatcher.webhook_urls.read().await;
assert_eq!(urls.len(), 2);
drop(urls);
dispatcher
.remove_webhook_url("https://example.com/hook1")
.await;
let urls = dispatcher.webhook_urls.read().await;
assert_eq!(urls.len(), 1);
assert_eq!(urls[0], "https://example.com/hook2");
}
#[tokio::test]
async fn test_dispatcher_event_distribution() {
let emitter = EventEmitter::with_default_config();
let dispatcher = EventDispatcher::new(emitter.clone());
let call_count = Arc::new(AtomicU64::new(0));
let handler = Arc::new(TestHandler {
call_count: call_count.clone(),
});
dispatcher.register_handler(handler).await;
dispatcher.start().await;
let event = Event::new(EventType::RateLimitTriggered {
key: "test".to_string(),
rule_id: "rule".to_string(),
decision: "Deny".to_string(),
});
emitter.emit(event).await.unwrap();
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(2);
let mut observed = 0u64;
while tokio::time::Instant::now() < deadline {
observed = call_count.load(Ordering::SeqCst);
if observed >= 1 {
break;
}
tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
}
dispatcher.stop().await;
assert_eq!(observed, 1, "handler should be called exactly once");
}
#[tokio::test]
async fn test_dispatcher_stop_and_restart() {
let emitter = EventEmitter::with_default_config();
let dispatcher = EventDispatcher::new(emitter);
dispatcher.start().await;
assert!(dispatcher.dispatch_handle.read().await.is_some());
dispatcher.start().await;
assert!(dispatcher.dispatch_handle.read().await.is_some());
dispatcher.stop().await;
assert!(dispatcher.dispatch_handle.read().await.is_none());
dispatcher.start().await;
assert!(dispatcher.dispatch_handle.read().await.is_some());
dispatcher.stop().await;
}
}