using System; using System.Collections.Concurrent; using System.Threading; using System.Threading.Tasks; namespace MarketAlly.AIPlugin.DevOps.Security { public class RateLimiter { private readonly ConcurrentDictionary _buckets; private readonly AuditLogger _auditLogger; public RateLimiter(AuditLogger auditLogger = null) { _buckets = new ConcurrentDictionary(); _auditLogger = auditLogger; } public async Task TryExecuteAsync(string clientId, int tokensRequired = 1, int maxTokens = 100, TimeSpan? refillInterval = null) { var interval = refillInterval ?? TimeSpan.FromMinutes(1); var bucket = _buckets.GetOrAdd(clientId, _ => new TokenBucket(maxTokens, interval)); var allowed = bucket.TryConsume(tokensRequired); await _auditLogger?.LogSecurityEventAsync(new SecurityAuditEvent { EventType = SecurityEventType.PermissionChecked, Severity = allowed ? SecuritySeverity.Low : SecuritySeverity.Medium, Source = nameof(RateLimiter), UserId = clientId, Details = $"Rate limit check: {(allowed ? "ALLOWED" : "DENIED")}", Metadata = new() { ["clientId"] = clientId, ["tokensRequired"] = tokensRequired, ["tokensAvailable"] = bucket.AvailableTokens, ["allowed"] = allowed } }); return allowed; } public void ClearClient(string clientId) { _buckets.TryRemove(clientId, out _); } public void ClearAll() { _buckets.Clear(); } } public class TokenBucket { private readonly int _maxTokens; private readonly TimeSpan _refillInterval; private readonly object _lock = new(); private int _availableTokens; private DateTime _lastRefill; public TokenBucket(int maxTokens, TimeSpan refillInterval) { _maxTokens = maxTokens; _refillInterval = refillInterval; _availableTokens = maxTokens; _lastRefill = DateTime.UtcNow; } public int AvailableTokens { get { lock (_lock) { Refill(); return _availableTokens; } } } public bool TryConsume(int tokens) { lock (_lock) { Refill(); if (_availableTokens >= tokens) { _availableTokens -= tokens; return true; } return false; } } private void Refill() { var now = DateTime.UtcNow; var timeSinceLastRefill = now - _lastRefill; if (timeSinceLastRefill >= _refillInterval) { var refillCycles = (int)(timeSinceLastRefill.TotalMilliseconds / _refillInterval.TotalMilliseconds); var tokensToAdd = refillCycles * _maxTokens; _availableTokens = Math.Min(_maxTokens, _availableTokens + tokensToAdd); _lastRefill = now; } } } }