use crate::audit::AuditLogEntry;
use crate::error::Result;
use dashmap::DashMap;
use std::sync::Arc;
pub trait AuditStorage: Send + Sync {
fn store(&self, entry: AuditLogEntry) -> Result<()>;
fn query(&self, query: &StorageQuery) -> Result<Vec<AuditLogEntry>>;
fn count(&self, query: &StorageQuery) -> Result<usize>;
}
pub struct InMemoryAuditStorage {
entries: Arc<DashMap<String, AuditLogEntry>>,
}
impl InMemoryAuditStorage {
pub fn new() -> Self {
Self {
entries: Arc::new(DashMap::new()),
}
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn clear(&self) {
self.entries.clear();
}
}
impl Default for InMemoryAuditStorage {
fn default() -> Self {
Self::new()
}
}
impl AuditStorage for InMemoryAuditStorage {
fn store(&self, entry: AuditLogEntry) -> Result<()> {
self.entries.insert(entry.id.clone(), entry);
Ok(())
}
fn query(&self, query: &StorageQuery) -> Result<Vec<AuditLogEntry>> {
let mut results: Vec<AuditLogEntry> = self
.entries
.iter()
.map(|e| e.value().clone())
.filter(|entry| query.matches(entry))
.collect();
results.sort_by_key(|x| std::cmp::Reverse(x.timestamp));
if let Some(limit) = query.limit {
results.truncate(limit);
}
Ok(results)
}
fn count(&self, query: &StorageQuery) -> Result<usize> {
Ok(self
.entries
.iter()
.filter(|e| query.matches(e.value()))
.count())
}
}
#[derive(Debug, Clone, Default)]
pub struct StorageQuery {
pub subject: Option<String>,
pub resource: Option<String>,
pub event_type: Option<crate::audit::AuditEventType>,
pub result: Option<crate::audit::AuditResult>,
pub tenant_id: Option<String>,
pub start_time: Option<chrono::DateTime<chrono::Utc>>,
pub end_time: Option<chrono::DateTime<chrono::Utc>>,
pub limit: Option<usize>,
}
impl StorageQuery {
pub fn new() -> Self {
Self::default()
}
pub fn with_subject(mut self, subject: String) -> Self {
self.subject = Some(subject);
self
}
pub fn with_resource(mut self, resource: String) -> Self {
self.resource = Some(resource);
self
}
pub fn with_event_type(mut self, event_type: crate::audit::AuditEventType) -> Self {
self.event_type = Some(event_type);
self
}
pub fn with_result(mut self, result: crate::audit::AuditResult) -> Self {
self.result = Some(result);
self
}
pub fn with_tenant_id(mut self, tenant_id: String) -> Self {
self.tenant_id = Some(tenant_id);
self
}
pub fn with_time_range(
mut self,
start: chrono::DateTime<chrono::Utc>,
end: chrono::DateTime<chrono::Utc>,
) -> Self {
self.start_time = Some(start);
self.end_time = Some(end);
self
}
pub fn with_limit(mut self, limit: usize) -> Self {
self.limit = Some(limit);
self
}
pub fn matches(&self, entry: &AuditLogEntry) -> bool {
if let Some(ref subject) = self.subject {
if entry.subject.as_ref() != Some(subject) {
return false;
}
}
if let Some(ref resource) = self.resource {
if entry.resource.as_ref() != Some(resource) {
return false;
}
}
if let Some(event_type) = self.event_type {
if entry.event_type != event_type {
return false;
}
}
if let Some(result) = self.result {
if entry.result != result {
return false;
}
}
if let Some(ref tenant_id) = self.tenant_id {
if entry.tenant_id.as_ref() != Some(tenant_id) {
return false;
}
}
if let Some(start) = self.start_time {
if entry.timestamp < start {
return false;
}
}
if let Some(end) = self.end_time {
if entry.timestamp > end {
return false;
}
}
true
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::audit::{AuditEventType, AuditResult};
#[test]
fn test_in_memory_storage() {
let storage = InMemoryAuditStorage::new();
let entry = AuditLogEntry::new(AuditEventType::Authentication, AuditResult::Success)
.with_subject("user-123".to_string());
storage.store(entry.clone()).expect("Failed to store");
assert_eq!(storage.len(), 1);
let query = StorageQuery::new().with_subject("user-123".to_string());
let results = storage.query(&query).expect("Query failed");
assert_eq!(results.len(), 1);
assert_eq!(results[0].id, entry.id);
}
#[test]
fn test_storage_query() {
let storage = InMemoryAuditStorage::new();
let entry1 = AuditLogEntry::new(AuditEventType::Authentication, AuditResult::Success)
.with_subject("user-1".to_string());
let entry2 = AuditLogEntry::new(AuditEventType::DataAccess, AuditResult::Success)
.with_subject("user-2".to_string());
storage.store(entry1).expect("Failed to store");
storage.store(entry2).expect("Failed to store");
let query = StorageQuery::new().with_event_type(AuditEventType::Authentication);
let results = storage.query(&query).expect("Query failed");
assert_eq!(results.len(), 1);
}
}