use anyhow::{Context, Result};
use prometheus::{
register_counter_vec, register_gauge, register_histogram_vec, register_int_counter_vec,
register_int_gauge, register_int_gauge_vec, CounterVec, Gauge, HistogramVec, IntCounterVec,
IntGauge, IntGaugeVec,
};
use std::time::Duration;
use tracing::{error, info};
use tracing_subscriber::{fmt, prelude::*, EnvFilter};
pub fn init_tracing() -> Result<()> {
let json_layer =
if std::env::var("ZENTINEL_LOG_FORMAT").unwrap_or_else(|_| "json".to_string()) == "json" {
Some(
fmt::layer()
.json()
.with_target(true)
.with_thread_ids(true)
.with_thread_names(true)
.with_file(true)
.with_line_number(true),
)
} else {
None
};
let pretty_layer = if std::env::var("ZENTINEL_LOG_FORMAT")
.unwrap_or_else(|_| "json".to_string())
== "pretty"
{
Some(
fmt::layer()
.pretty()
.with_target(true)
.with_thread_ids(true)
.with_thread_names(true)
.with_file(true)
.with_line_number(true),
)
} else {
None
};
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
tracing_subscriber::registry()
.with(env_filter)
.with(json_layer)
.with(pretty_layer)
.init();
info!("Tracing initialized");
Ok(())
}
pub struct RequestMetrics {
request_duration: HistogramVec,
request_count: IntCounterVec,
active_requests: IntGauge,
upstream_attempts: IntCounterVec,
upstream_failures: IntCounterVec,
circuit_breaker_state: IntGaugeVec,
agent_latency: HistogramVec,
agent_timeouts: IntCounterVec,
blocked_requests: CounterVec,
request_body_size: HistogramVec,
response_body_size: HistogramVec,
tls_handshake_duration: HistogramVec,
connection_pool_size: IntGaugeVec,
connection_pool_idle: IntGaugeVec,
connection_pool_acquired: IntCounterVec,
memory_usage: IntGauge,
cpu_usage: Gauge,
open_connections: IntGauge,
websocket_frames_total: IntCounterVec,
websocket_connections_total: IntCounterVec,
websocket_inspection_duration: HistogramVec,
websocket_frame_size: HistogramVec,
decompression_total: IntCounterVec,
decompression_ratio: HistogramVec,
shadow_requests_total: IntCounterVec,
shadow_errors_total: IntCounterVec,
shadow_latency_seconds: HistogramVec,
pii_detected_total: IntCounterVec,
}
fn status_str(status: u16) -> &'static str {
match status {
200 => "200",
201 => "201",
204 => "204",
301 => "301",
302 => "302",
304 => "304",
307 => "307",
308 => "308",
400 => "400",
401 => "401",
403 => "403",
404 => "404",
405 => "405",
408 => "408",
409 => "409",
413 => "413",
429 => "429",
500 => "500",
502 => "502",
503 => "503",
504 => "504",
_ => Box::leak(status.to_string().into_boxed_str()),
}
}
impl RequestMetrics {
pub fn new() -> Result<Self> {
let latency_buckets = vec![
0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
];
let size_buckets = vec![
100.0,
1_000.0,
10_000.0,
100_000.0,
1_000_000.0,
10_000_000.0,
100_000_000.0,
];
let request_duration = register_histogram_vec!(
"zentinel_request_duration_seconds",
"Request duration in seconds",
&["route", "method"],
latency_buckets.clone()
)
.context("Failed to register request_duration metric")?;
let request_count = register_int_counter_vec!(
"zentinel_requests_total",
"Total number of requests",
&["route", "method", "status"]
)
.context("Failed to register request_count metric")?;
let active_requests = register_int_gauge!(
"zentinel_active_requests",
"Number of currently active requests"
)
.context("Failed to register active_requests metric")?;
let upstream_attempts = register_int_counter_vec!(
"zentinel_upstream_attempts_total",
"Total upstream connection attempts",
&["upstream", "route"]
)
.context("Failed to register upstream_attempts metric")?;
let upstream_failures = register_int_counter_vec!(
"zentinel_upstream_failures_total",
"Total upstream connection failures",
&["upstream", "route", "reason"]
)
.context("Failed to register upstream_failures metric")?;
let circuit_breaker_state = register_int_gauge_vec!(
"zentinel_circuit_breaker_state",
"Circuit breaker state (0=closed, 1=open)",
&["component", "route"]
)
.context("Failed to register circuit_breaker_state metric")?;
let agent_latency = register_histogram_vec!(
"zentinel_agent_latency_seconds",
"Agent call latency in seconds",
&["agent", "event"],
latency_buckets.clone()
)
.context("Failed to register agent_latency metric")?;
let agent_timeouts = register_int_counter_vec!(
"zentinel_agent_timeouts_total",
"Total agent call timeouts",
&["agent", "event"]
)
.context("Failed to register agent_timeouts metric")?;
let blocked_requests = register_counter_vec!(
"zentinel_blocked_requests_total",
"Total blocked requests by reason",
&["reason"]
)
.context("Failed to register blocked_requests metric")?;
let request_body_size = register_histogram_vec!(
"zentinel_request_body_size_bytes",
"Request body size in bytes",
&["route"],
size_buckets.clone()
)
.context("Failed to register request_body_size metric")?;
let response_body_size = register_histogram_vec!(
"zentinel_response_body_size_bytes",
"Response body size in bytes",
&["route"],
size_buckets.clone()
)
.context("Failed to register response_body_size metric")?;
let tls_handshake_duration = register_histogram_vec!(
"zentinel_tls_handshake_duration_seconds",
"TLS handshake duration in seconds",
&["version"],
latency_buckets
)
.context("Failed to register tls_handshake_duration metric")?;
let connection_pool_size = register_int_gauge_vec!(
"zentinel_connection_pool_size",
"Total connections in pool",
&["upstream"]
)
.context("Failed to register connection_pool_size metric")?;
let connection_pool_idle = register_int_gauge_vec!(
"zentinel_connection_pool_idle",
"Idle connections in pool",
&["upstream"]
)
.context("Failed to register connection_pool_idle metric")?;
let connection_pool_acquired = register_int_counter_vec!(
"zentinel_connection_pool_acquired_total",
"Total connections acquired from pool",
&["upstream"]
)
.context("Failed to register connection_pool_acquired metric")?;
let memory_usage = register_int_gauge!(
"zentinel_memory_usage_bytes",
"Current memory usage in bytes"
)
.context("Failed to register memory_usage metric")?;
let cpu_usage =
register_gauge!("zentinel_cpu_usage_percent", "Current CPU usage percentage")
.context("Failed to register cpu_usage metric")?;
let open_connections =
register_int_gauge!("zentinel_open_connections", "Number of open connections")
.context("Failed to register open_connections metric")?;
let websocket_frames_total = register_int_counter_vec!(
"zentinel_websocket_frames_total",
"Total WebSocket frames processed",
&["route", "direction", "opcode", "decision"]
)
.context("Failed to register websocket_frames_total metric")?;
let websocket_connections_total = register_int_counter_vec!(
"zentinel_websocket_connections_total",
"Total WebSocket connections with inspection enabled",
&["route"]
)
.context("Failed to register websocket_connections_total metric")?;
let frame_latency_buckets = vec![
0.0001, 0.0005, 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5,
];
let websocket_inspection_duration = register_histogram_vec!(
"zentinel_websocket_inspection_duration_seconds",
"WebSocket frame inspection duration in seconds",
&["route"],
frame_latency_buckets
)
.context("Failed to register websocket_inspection_duration metric")?;
let frame_size_buckets = vec![
64.0, 256.0, 1024.0, 4096.0, 16384.0, 65536.0, 262144.0, 1048576.0,
];
let websocket_frame_size = register_histogram_vec!(
"zentinel_websocket_frame_size_bytes",
"WebSocket frame payload size in bytes",
&["route", "direction", "opcode"],
frame_size_buckets
)
.context("Failed to register websocket_frame_size metric")?;
let decompression_total = register_int_counter_vec!(
"zentinel_decompression_total",
"Total body decompression operations",
&["encoding", "result"]
)
.context("Failed to register decompression_total metric")?;
let ratio_buckets = vec![1.0, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0, 200.0, 500.0, 1000.0];
let decompression_ratio = register_histogram_vec!(
"zentinel_decompression_ratio",
"Decompression ratio (decompressed_size / compressed_size)",
&["encoding"],
ratio_buckets
)
.context("Failed to register decompression_ratio metric")?;
let shadow_requests_total = register_int_counter_vec!(
"zentinel_shadow_requests_total",
"Total shadow requests sent to mirror upstream",
&["route", "upstream", "result"]
)
.context("Failed to register shadow_requests_total metric")?;
let shadow_errors_total = register_int_counter_vec!(
"zentinel_shadow_errors_total",
"Total shadow request errors",
&["route", "upstream", "error_type"]
)
.context("Failed to register shadow_errors_total metric")?;
let shadow_latency_buckets = vec![
0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
];
let shadow_latency_seconds = register_histogram_vec!(
"zentinel_shadow_latency_seconds",
"Shadow request latency in seconds",
&["route", "upstream"],
shadow_latency_buckets
)
.context("Failed to register shadow_latency_seconds metric")?;
let pii_detected_total = register_int_counter_vec!(
"zentinel_pii_detected_total",
"Total PII detections in inference responses",
&["route", "category"]
)
.context("Failed to register pii_detected_total metric")?;
Ok(Self {
request_duration,
request_count,
active_requests,
upstream_attempts,
upstream_failures,
circuit_breaker_state,
agent_latency,
agent_timeouts,
blocked_requests,
request_body_size,
response_body_size,
tls_handshake_duration,
connection_pool_size,
connection_pool_idle,
connection_pool_acquired,
memory_usage,
cpu_usage,
open_connections,
websocket_frames_total,
websocket_connections_total,
websocket_inspection_duration,
websocket_frame_size,
decompression_total,
decompression_ratio,
shadow_requests_total,
shadow_errors_total,
shadow_latency_seconds,
pii_detected_total,
})
}
pub fn record_request(&self, route: &str, method: &str, status: u16, duration: Duration) {
self.request_duration
.with_label_values(&[route, method])
.observe(duration.as_secs_f64());
self.request_count
.with_label_values(&[route, method, status_str(status)])
.inc();
}
pub fn inc_active_requests(&self) {
self.active_requests.inc();
}
pub fn dec_active_requests(&self) {
self.active_requests.dec();
}
pub fn record_upstream_attempt(&self, upstream: &str, route: &str) {
self.upstream_attempts
.with_label_values(&[upstream, route])
.inc();
}
pub fn record_upstream_failure(&self, upstream: &str, route: &str, reason: &str) {
self.upstream_failures
.with_label_values(&[upstream, route, reason])
.inc();
}
pub fn set_circuit_breaker_state(&self, component: &str, route: &str, is_open: bool) {
let state = if is_open { 1 } else { 0 };
self.circuit_breaker_state
.with_label_values(&[component, route])
.set(state);
}
pub fn record_agent_latency(&self, agent: &str, event: &str, duration: Duration) {
self.agent_latency
.with_label_values(&[agent, event])
.observe(duration.as_secs_f64());
}
pub fn record_agent_timeout(&self, agent: &str, event: &str) {
self.agent_timeouts.with_label_values(&[agent, event]).inc();
}
pub fn record_blocked_request(&self, reason: &str) {
self.blocked_requests.with_label_values(&[reason]).inc();
}
pub fn record_pii_detected(&self, route: &str, category: &str) {
self.pii_detected_total
.with_label_values(&[route, category])
.inc();
}
pub fn record_request_body_size(&self, route: &str, size_bytes: usize) {
self.request_body_size
.with_label_values(&[route])
.observe(size_bytes as f64);
}
pub fn record_response_body_size(&self, route: &str, size_bytes: usize) {
self.response_body_size
.with_label_values(&[route])
.observe(size_bytes as f64);
}
pub fn record_tls_handshake(&self, version: &str, duration: Duration) {
self.tls_handshake_duration
.with_label_values(&[version])
.observe(duration.as_secs_f64());
}
pub fn update_connection_pool(&self, upstream: &str, size: i64, idle: i64) {
self.connection_pool_size
.with_label_values(&[upstream])
.set(size);
self.connection_pool_idle
.with_label_values(&[upstream])
.set(idle);
}
pub fn record_connection_acquired(&self, upstream: &str) {
self.connection_pool_acquired
.with_label_values(&[upstream])
.inc();
}
pub fn update_system_metrics(&self) {
use sysinfo::{CpuRefreshKind, MemoryRefreshKind, RefreshKind, System};
let mut system = System::new_with_specifics(
RefreshKind::nothing()
.with_cpu(CpuRefreshKind::everything())
.with_memory(MemoryRefreshKind::everything()),
);
self.memory_usage.set(system.total_memory() as i64);
system.refresh_cpu_usage();
self.cpu_usage.set(system.global_cpu_usage() as f64);
}
pub fn set_open_connections(&self, count: i64) {
self.open_connections.set(count);
}
pub fn record_websocket_frame(
&self,
route: &str,
direction: &str,
opcode: &str,
decision: &str,
) {
self.websocket_frames_total
.with_label_values(&[route, direction, opcode, decision])
.inc();
}
pub fn record_websocket_connection(&self, route: &str) {
self.websocket_connections_total
.with_label_values(&[route])
.inc();
}
pub fn record_websocket_inspection_duration(&self, route: &str, duration: Duration) {
self.websocket_inspection_duration
.with_label_values(&[route])
.observe(duration.as_secs_f64());
}
pub fn record_websocket_frame_size(
&self,
route: &str,
direction: &str,
opcode: &str,
size_bytes: usize,
) {
self.websocket_frame_size
.with_label_values(&[route, direction, opcode])
.observe(size_bytes as f64);
}
pub fn record_decompression_success(&self, encoding: &str, ratio: f64) {
self.decompression_total
.with_label_values(&[encoding, "success"])
.inc();
self.decompression_ratio
.with_label_values(&[encoding])
.observe(ratio);
}
pub fn record_decompression_failure(&self, encoding: &str, reason: &str) {
self.decompression_total
.with_label_values(&[encoding, reason])
.inc();
}
pub fn record_shadow_success(&self, route: &str, upstream: &str, duration: Duration) {
self.shadow_requests_total
.with_label_values(&[route, upstream, "success"])
.inc();
self.shadow_latency_seconds
.with_label_values(&[route, upstream])
.observe(duration.as_secs_f64());
}
pub fn record_shadow_error(&self, route: &str, upstream: &str, error_type: &str) {
self.shadow_requests_total
.with_label_values(&[route, upstream, "error"])
.inc();
self.shadow_errors_total
.with_label_values(&[route, upstream, error_type])
.inc();
}
pub fn record_shadow_timeout(&self, route: &str, upstream: &str, duration: Duration) {
self.shadow_requests_total
.with_label_values(&[route, upstream, "timeout"])
.inc();
self.shadow_errors_total
.with_label_values(&[route, upstream, "timeout"])
.inc();
self.shadow_latency_seconds
.with_label_values(&[route, upstream])
.observe(duration.as_secs_f64());
}
}
#[derive(Debug, serde::Serialize)]
pub struct AuditLogEntry {
pub timestamp: String,
pub correlation_id: String,
pub event_type: String,
pub route: Option<String>,
pub client_addr: Option<String>,
pub user_agent: Option<String>,
pub method: String,
pub path: String,
pub status: Option<u16>,
pub duration_ms: u64,
pub upstream: Option<String>,
pub waf_decision: Option<WafDecision>,
pub agent_decisions: Vec<AgentDecision>,
pub error: Option<String>,
pub tags: Vec<String>,
}
#[derive(Debug, serde::Serialize)]
pub struct WafDecision {
pub action: String,
pub rule_ids: Vec<String>,
pub confidence: f32,
pub reason: String,
pub matched_data: Option<String>,
}
#[derive(Debug, serde::Serialize)]
pub struct AgentDecision {
pub agent_name: String,
pub event: String,
pub action: String,
pub latency_ms: u64,
pub metadata: serde_json::Value,
}
impl AuditLogEntry {
pub fn new(correlation_id: String, method: String, path: String) -> Self {
Self {
timestamp: chrono::Utc::now().to_rfc3339(),
correlation_id,
event_type: "request".to_string(),
route: None,
client_addr: None,
user_agent: None,
method,
path,
status: None,
duration_ms: 0,
upstream: None,
waf_decision: None,
agent_decisions: vec![],
error: None,
tags: vec![],
}
}
pub fn write(&self) {
match serde_json::to_string(self) {
Ok(json) => println!("AUDIT: {}", json),
Err(e) => error!("Failed to serialize audit log: {}", e),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub enum HealthStatus {
Healthy,
Degraded,
Unhealthy,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct ComponentHealth {
pub name: String,
pub status: HealthStatus,
pub last_check: chrono::DateTime<chrono::Utc>,
pub consecutive_failures: u32,
pub error_message: Option<String>,
}
pub struct ComponentHealthTracker {
components: parking_lot::RwLock<Vec<ComponentHealth>>,
}
impl Default for ComponentHealthTracker {
fn default() -> Self {
Self::new()
}
}
impl ComponentHealthTracker {
pub fn new() -> Self {
Self {
components: parking_lot::RwLock::new(vec![]),
}
}
pub fn update_component(&self, name: String, status: HealthStatus, error: Option<String>) {
let mut components = self.components.write();
if let Some(component) = components.iter_mut().find(|c| c.name == name) {
component.status = status;
component.last_check = chrono::Utc::now();
component.error_message = error;
if status != HealthStatus::Healthy {
component.consecutive_failures += 1;
} else {
component.consecutive_failures = 0;
}
} else {
components.push(ComponentHealth {
name,
status,
last_check: chrono::Utc::now(),
consecutive_failures: if status != HealthStatus::Healthy {
1
} else {
0
},
error_message: error,
});
}
}
pub fn get_status(&self) -> HealthStatus {
let components = self.components.read();
if components.is_empty() {
return HealthStatus::Healthy;
}
let unhealthy_count = components
.iter()
.filter(|c| c.status == HealthStatus::Unhealthy)
.count();
let degraded_count = components
.iter()
.filter(|c| c.status == HealthStatus::Degraded)
.count();
if unhealthy_count > 0 {
HealthStatus::Unhealthy
} else if degraded_count > 0 {
HealthStatus::Degraded
} else {
HealthStatus::Healthy
}
}
pub fn get_report(&self) -> Vec<ComponentHealth> {
self.components.read().clone()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_metrics_creation() {
let metrics = RequestMetrics::new().expect("Failed to create metrics");
metrics.record_request("test_route", "GET", 200, Duration::from_millis(100));
metrics.inc_active_requests();
metrics.dec_active_requests();
metrics.record_upstream_attempt("backend1", "test_route");
}
#[test]
fn test_audit_log() {
let mut entry = AuditLogEntry::new(
"test-correlation-id".to_string(),
"GET".to_string(),
"/api/test".to_string(),
);
entry.status = Some(200);
entry.duration_ms = 150;
entry.tags.push("test".to_string());
let json = serde_json::to_string(&entry).expect("Failed to serialize audit log");
assert!(json.contains("test-correlation-id"));
}
#[test]
fn test_health_checker() {
let checker = ComponentHealthTracker::new();
assert_eq!(checker.get_status(), HealthStatus::Healthy);
checker.update_component("upstream1".to_string(), HealthStatus::Healthy, None);
assert_eq!(checker.get_status(), HealthStatus::Healthy);
checker.update_component(
"agent1".to_string(),
HealthStatus::Degraded,
Some("Slow response".to_string()),
);
assert_eq!(checker.get_status(), HealthStatus::Degraded);
checker.update_component(
"upstream2".to_string(),
HealthStatus::Unhealthy,
Some("Connection refused".to_string()),
);
assert_eq!(checker.get_status(), HealthStatus::Unhealthy);
}
}