use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum PolicyError {
#[error("Policy not found: {0}")]
NotFound(String),
#[error("Invalid policy: {0}")]
Invalid(String),
#[error("Policy conflict: {0}")]
Conflict(String),
#[error("Evaluation error: {0}")]
Evaluation(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PolicyAction {
Allow,
Deny,
RateLimit,
Log,
Verify,
}
#[derive(Debug, Clone)]
pub struct PolicyResult {
pub allowed: bool,
pub action: PolicyAction,
pub policy_name: Option<String>,
pub reason: String,
pub confidence: f64,
}
impl PolicyResult {
pub fn allow(reason: String) -> Self {
Self {
allowed: true,
action: PolicyAction::Allow,
policy_name: None,
reason,
confidence: 1.0,
}
}
pub fn deny(reason: String, policy_name: Option<String>) -> Self {
Self {
allowed: false,
action: PolicyAction::Deny,
policy_name,
reason,
confidence: 1.0,
}
}
}
#[derive(Debug, Clone)]
pub struct ConnectionPolicy {
pub name: String,
pub action: PolicyAction,
pub priority: u32,
pub peer_whitelist: Vec<String>,
pub peer_blacklist: Vec<String>,
pub max_connections_per_peer: Option<usize>,
pub rate_limit: Option<f64>,
pub enabled: bool,
}
impl ConnectionPolicy {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
action: PolicyAction::Allow,
priority: 50,
peer_whitelist: Vec::new(),
peer_blacklist: Vec::new(),
max_connections_per_peer: None,
rate_limit: None,
enabled: true,
}
}
pub fn with_action(mut self, action: PolicyAction) -> Self {
self.action = action;
self
}
pub fn with_priority(mut self, priority: u32) -> Self {
self.priority = priority;
self
}
pub fn with_whitelist_peer(mut self, peer: impl Into<String>) -> Self {
self.peer_whitelist.push(peer.into());
self
}
pub fn with_blacklist_peer(mut self, peer: impl Into<String>) -> Self {
self.peer_blacklist.push(peer.into());
self
}
pub fn with_max_connections(mut self, max: usize) -> Self {
self.max_connections_per_peer = Some(max);
self
}
pub fn with_rate_limit(mut self, rate: f64) -> Self {
self.rate_limit = Some(rate);
self
}
pub fn evaluate(&self, peer_id: &str) -> Option<PolicyResult> {
if !self.enabled {
return None;
}
if self.peer_blacklist.iter().any(|p| p == peer_id) {
return Some(PolicyResult::deny(
format!("Peer {} is blacklisted", peer_id),
Some(self.name.clone()),
));
}
if !self.peer_whitelist.is_empty() && !self.peer_whitelist.iter().any(|p| p == peer_id) {
return Some(PolicyResult::deny(
format!("Peer {} not in whitelist", peer_id),
Some(self.name.clone()),
));
}
Some(PolicyResult {
allowed: self.action == PolicyAction::Allow,
action: self.action,
policy_name: Some(self.name.clone()),
reason: format!("Policy {} matched", self.name),
confidence: 1.0,
})
}
}
#[derive(Debug, Clone)]
pub struct BandwidthPolicy {
pub name: String,
pub max_upload_bps: Option<u64>,
pub max_download_bps: Option<u64>,
pub per_peer_limit_bps: Option<u64>,
pub priority: u32,
pub enabled: bool,
}
impl BandwidthPolicy {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
max_upload_bps: None,
max_download_bps: None,
per_peer_limit_bps: None,
priority: 50,
enabled: true,
}
}
pub fn with_max_upload(mut self, bps: u64) -> Self {
self.max_upload_bps = Some(bps);
self
}
pub fn with_max_download(mut self, bps: u64) -> Self {
self.max_download_bps = Some(bps);
self
}
pub fn with_per_peer_limit(mut self, bps: u64) -> Self {
self.per_peer_limit_bps = Some(bps);
self
}
}
#[derive(Debug, Clone)]
pub struct ContentPolicy {
pub name: String,
pub action: PolicyAction,
pub allowed_patterns: Vec<String>,
pub blocked_patterns: Vec<String>,
pub max_size: Option<u64>,
pub priority: u32,
pub enabled: bool,
}
impl ContentPolicy {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
action: PolicyAction::Allow,
allowed_patterns: Vec::new(),
blocked_patterns: Vec::new(),
max_size: None,
priority: 50,
enabled: true,
}
}
pub fn with_allowed_pattern(mut self, pattern: impl Into<String>) -> Self {
self.allowed_patterns.push(pattern.into());
self
}
pub fn with_blocked_pattern(mut self, pattern: impl Into<String>) -> Self {
self.blocked_patterns.push(pattern.into());
self
}
pub fn with_max_size(mut self, size: u64) -> Self {
self.max_size = Some(size);
self
}
}
#[derive(Debug, Clone)]
pub struct PolicyConfig {
pub enabled: bool,
pub default_action: PolicyAction,
pub log_decisions: bool,
pub max_policies_per_type: usize,
pub evaluation_timeout: Duration,
}
impl Default for PolicyConfig {
fn default() -> Self {
Self {
enabled: true,
default_action: PolicyAction::Allow,
log_decisions: true,
max_policies_per_type: 100,
evaluation_timeout: Duration::from_millis(100),
}
}
}
impl PolicyConfig {
pub fn strict() -> Self {
Self {
enabled: true,
default_action: PolicyAction::Deny,
log_decisions: true,
max_policies_per_type: 200,
evaluation_timeout: Duration::from_millis(50),
}
}
pub fn permissive() -> Self {
Self {
enabled: true,
default_action: PolicyAction::Allow,
log_decisions: false,
max_policies_per_type: 50,
evaluation_timeout: Duration::from_millis(200),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct PolicyStats {
pub evaluations: u64,
pub allowed: u64,
pub denied: u64,
pub rate_limited: u64,
pub policy_hits: HashMap<String, u64>,
pub avg_eval_time_ms: f64,
}
pub struct PolicyEngine {
config: PolicyConfig,
connection_policies: Arc<RwLock<Vec<ConnectionPolicy>>>,
bandwidth_policies: Arc<RwLock<Vec<BandwidthPolicy>>>,
content_policies: Arc<RwLock<Vec<ContentPolicy>>>,
stats: Arc<RwLock<PolicyStats>>,
connection_counts: Arc<DashMap<String, usize>>,
}
impl PolicyEngine {
pub fn new(config: PolicyConfig) -> Self {
Self {
config,
connection_policies: Arc::new(RwLock::new(Vec::new())),
bandwidth_policies: Arc::new(RwLock::new(Vec::new())),
content_policies: Arc::new(RwLock::new(Vec::new())),
stats: Arc::new(RwLock::new(PolicyStats::default())),
connection_counts: Arc::new(DashMap::new()),
}
}
pub fn add_connection_policy(&self, policy: ConnectionPolicy) -> Result<(), PolicyError> {
let mut policies = self.connection_policies.write();
if policies.len() >= self.config.max_policies_per_type {
return Err(PolicyError::Invalid(
"Maximum connection policies reached".to_string(),
));
}
policies.push(policy);
policies.sort_by_key(|p| std::cmp::Reverse(p.priority));
Ok(())
}
pub fn add_bandwidth_policy(&self, policy: BandwidthPolicy) -> Result<(), PolicyError> {
let mut policies = self.bandwidth_policies.write();
if policies.len() >= self.config.max_policies_per_type {
return Err(PolicyError::Invalid(
"Maximum bandwidth policies reached".to_string(),
));
}
policies.push(policy);
policies.sort_by_key(|p| std::cmp::Reverse(p.priority));
Ok(())
}
pub fn add_content_policy(&self, policy: ContentPolicy) -> Result<(), PolicyError> {
let mut policies = self.content_policies.write();
if policies.len() >= self.config.max_policies_per_type {
return Err(PolicyError::Invalid(
"Maximum content policies reached".to_string(),
));
}
policies.push(policy);
policies.sort_by_key(|p| std::cmp::Reverse(p.priority));
Ok(())
}
pub async fn evaluate_connection(&self, peer_id: &str) -> Result<bool, PolicyError> {
let start = Instant::now();
if !self.config.enabled {
return Ok(true);
}
let policies = self.connection_policies.read();
for policy in policies.iter() {
if let Some(result) = policy.evaluate(peer_id) {
self.record_evaluation(&result, start.elapsed());
return Ok(result.allowed);
}
}
let allowed = self.config.default_action == PolicyAction::Allow;
self.record_default_evaluation(allowed, start.elapsed());
Ok(allowed)
}
pub fn can_connect(&self, peer_id: &str) -> bool {
let count = self.connection_counts.get(peer_id).map(|c| *c).unwrap_or(0);
let policies = self.connection_policies.read();
for policy in policies.iter() {
if let Some(max) = policy.max_connections_per_peer {
if count >= max {
return false;
}
}
}
true
}
pub fn record_connection(&self, peer_id: &str) {
self.connection_counts
.entry(peer_id.to_string())
.and_modify(|c| *c += 1)
.or_insert(1);
}
pub fn record_disconnection(&self, peer_id: &str) {
if let Some(mut count) = self.connection_counts.get_mut(peer_id) {
if *count > 0 {
*count -= 1;
}
}
}
pub fn remove_connection_policy(&self, name: &str) -> Result<(), PolicyError> {
let mut policies = self.connection_policies.write();
let len_before = policies.len();
policies.retain(|p| p.name != name);
if policies.len() == len_before {
Err(PolicyError::NotFound(name.to_string()))
} else {
Ok(())
}
}
pub fn connection_policies(&self) -> Vec<ConnectionPolicy> {
self.connection_policies.read().clone()
}
pub fn bandwidth_policies(&self) -> Vec<BandwidthPolicy> {
self.bandwidth_policies.read().clone()
}
pub fn content_policies(&self) -> Vec<ContentPolicy> {
self.content_policies.read().clone()
}
pub fn stats(&self) -> PolicyStats {
self.stats.read().clone()
}
pub fn reset_stats(&self) {
let mut stats = self.stats.write();
*stats = PolicyStats::default();
}
fn record_evaluation(&self, result: &PolicyResult, duration: Duration) {
let mut stats = self.stats.write();
stats.evaluations += 1;
if result.allowed {
stats.allowed += 1;
} else {
stats.denied += 1;
}
if result.action == PolicyAction::RateLimit {
stats.rate_limited += 1;
}
if let Some(ref policy_name) = result.policy_name {
*stats.policy_hits.entry(policy_name.clone()).or_insert(0) += 1;
}
let eval_time_ms = duration.as_secs_f64() * 1000.0;
let alpha = 0.3;
stats.avg_eval_time_ms = alpha * eval_time_ms + (1.0 - alpha) * stats.avg_eval_time_ms;
}
fn record_default_evaluation(&self, allowed: bool, duration: Duration) {
let mut stats = self.stats.write();
stats.evaluations += 1;
if allowed {
stats.allowed += 1;
} else {
stats.denied += 1;
}
let eval_time_ms = duration.as_secs_f64() * 1000.0;
let alpha = 0.3;
stats.avg_eval_time_ms = alpha * eval_time_ms + (1.0 - alpha) * stats.avg_eval_time_ms;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_policy_creation() {
let policy = ConnectionPolicy::new("test")
.with_action(PolicyAction::Allow)
.with_priority(100);
assert_eq!(policy.name, "test");
assert_eq!(policy.action, PolicyAction::Allow);
assert_eq!(policy.priority, 100);
}
#[test]
fn test_policy_engine_creation() {
let config = PolicyConfig::default();
let engine = PolicyEngine::new(config);
assert_eq!(engine.connection_policies().len(), 0);
}
#[tokio::test]
async fn test_add_connection_policy() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy = ConnectionPolicy::new("test");
assert!(engine.add_connection_policy(policy).is_ok());
assert_eq!(engine.connection_policies().len(), 1);
}
#[tokio::test]
async fn test_blacklist_evaluation() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy = ConnectionPolicy::new("blacklist")
.with_action(PolicyAction::Allow)
.with_blacklist_peer("bad_peer");
engine
.add_connection_policy(policy)
.expect("test: should add blacklist policy");
let allowed = engine
.evaluate_connection("bad_peer")
.await
.expect("test: should evaluate bad_peer connection");
assert!(!allowed);
let allowed = engine
.evaluate_connection("good_peer")
.await
.expect("test: should evaluate good_peer connection");
assert!(allowed);
}
#[tokio::test]
async fn test_whitelist_evaluation() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy = ConnectionPolicy::new("whitelist")
.with_action(PolicyAction::Allow)
.with_whitelist_peer("good_peer");
engine
.add_connection_policy(policy)
.expect("test: should add whitelist policy");
let allowed = engine
.evaluate_connection("good_peer")
.await
.expect("test: should evaluate good_peer connection");
assert!(allowed);
let allowed = engine
.evaluate_connection("bad_peer")
.await
.expect("test: should evaluate bad_peer connection");
assert!(!allowed);
}
#[test]
fn test_connection_counting() {
let engine = PolicyEngine::new(PolicyConfig::default());
engine.record_connection("peer1");
engine.record_connection("peer1");
assert!(engine.can_connect("peer1"));
engine.record_disconnection("peer1");
assert!(engine.can_connect("peer1"));
}
#[tokio::test]
async fn test_policy_priority() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy1 = ConnectionPolicy::new("low")
.with_action(PolicyAction::Deny)
.with_priority(10);
let policy2 = ConnectionPolicy::new("high")
.with_action(PolicyAction::Allow)
.with_priority(100);
engine
.add_connection_policy(policy1)
.expect("test: should add low-priority deny policy");
engine
.add_connection_policy(policy2)
.expect("test: should add high-priority allow policy");
let allowed = engine
.evaluate_connection("test_peer")
.await
.expect("test: should evaluate test_peer connection");
assert!(allowed);
}
#[tokio::test]
async fn test_remove_policy() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy = ConnectionPolicy::new("test");
engine
.add_connection_policy(policy)
.expect("test: should add test policy");
assert!(engine.remove_connection_policy("test").is_ok());
assert_eq!(engine.connection_policies().len(), 0);
assert!(engine.remove_connection_policy("nonexistent").is_err());
}
#[tokio::test]
async fn test_statistics() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy = ConnectionPolicy::new("test").with_action(PolicyAction::Allow);
engine
.add_connection_policy(policy)
.expect("test: should add allow policy for stats");
engine
.evaluate_connection("peer1")
.await
.expect("test: should evaluate peer1 for stats");
engine
.evaluate_connection("peer2")
.await
.expect("test: should evaluate peer2 for stats");
let stats = engine.stats();
assert_eq!(stats.evaluations, 2);
assert_eq!(stats.allowed, 2);
}
#[test]
fn test_bandwidth_policy() {
let policy = BandwidthPolicy::new("test")
.with_max_upload(1_000_000)
.with_max_download(5_000_000)
.with_per_peer_limit(100_000);
assert_eq!(policy.max_upload_bps, Some(1_000_000));
assert_eq!(policy.max_download_bps, Some(5_000_000));
assert_eq!(policy.per_peer_limit_bps, Some(100_000));
}
#[test]
fn test_content_policy() {
let policy = ContentPolicy::new("test")
.with_allowed_pattern("^Qm.*")
.with_max_size(10_000_000);
assert_eq!(policy.allowed_patterns.len(), 1);
assert_eq!(policy.max_size, Some(10_000_000));
}
#[test]
fn test_policy_config_presets() {
let strict = PolicyConfig::strict();
assert_eq!(strict.default_action, PolicyAction::Deny);
let permissive = PolicyConfig::permissive();
assert_eq!(permissive.default_action, PolicyAction::Allow);
}
#[tokio::test]
async fn test_default_action() {
let config = PolicyConfig {
default_action: PolicyAction::Deny,
..Default::default()
};
let engine = PolicyEngine::new(config);
let allowed = engine
.evaluate_connection("peer1")
.await
.expect("test: should evaluate peer1 with deny default");
assert!(!allowed);
}
#[tokio::test]
async fn test_reset_stats() {
let engine = PolicyEngine::new(PolicyConfig::default());
let policy = ConnectionPolicy::new("test");
engine
.add_connection_policy(policy)
.expect("test: should add policy for reset stats test");
engine
.evaluate_connection("peer1")
.await
.expect("test: should evaluate peer1 before stats reset");
assert_eq!(engine.stats().evaluations, 1);
engine.reset_stats();
assert_eq!(engine.stats().evaluations, 0);
}
}