use std::{
num::NonZeroUsize,
sync::{Arc, Mutex},
};
use lru::LruCache;
use crate::{
error::FraiseQLError,
utils::clock::{Clock, SystemClock},
};
const MAX_RATE_LIMITER_ENTRIES: usize = 100_000;
#[derive(Debug, Clone)]
pub struct RateLimitDimension {
pub max_requests: u32,
pub window_secs: u64,
}
impl RateLimitDimension {
const fn is_rate_limited(&self) -> bool {
self.max_requests == 0
}
}
#[derive(Debug, Clone)]
pub struct ValidationRateLimitingConfig {
pub enabled: bool,
pub validation_errors_max_requests: u32,
pub validation_errors_window_secs: u64,
pub depth_errors_max_requests: u32,
pub depth_errors_window_secs: u64,
pub complexity_errors_max_requests: u32,
pub complexity_errors_window_secs: u64,
pub malformed_errors_max_requests: u32,
pub malformed_errors_window_secs: u64,
pub async_validation_errors_max_requests: u32,
pub async_validation_errors_window_secs: u64,
}
impl Default for ValidationRateLimitingConfig {
fn default() -> Self {
Self {
enabled: true,
validation_errors_max_requests: 100,
validation_errors_window_secs: 60,
depth_errors_max_requests: 50,
depth_errors_window_secs: 60,
complexity_errors_max_requests: 30,
complexity_errors_window_secs: 60,
malformed_errors_max_requests: 40,
malformed_errors_window_secs: 60,
async_validation_errors_max_requests: 60,
async_validation_errors_window_secs: 60,
}
}
}
impl ValidationRateLimitingConfig {
#[must_use = "builder does nothing until .build() is called"]
pub fn builder() -> ValidationRateLimitingConfigBuilder {
ValidationRateLimitingConfigBuilder::default()
}
}
#[derive(Debug, Default)]
pub struct ValidationRateLimitingConfigBuilder {
inner: ValidationRateLimitingConfig,
}
impl ValidationRateLimitingConfigBuilder {
#[must_use = "builder method returns modified builder"]
pub const fn enabled(mut self, enabled: bool) -> Self {
self.inner.enabled = enabled;
self
}
#[must_use = "builder method returns modified builder"]
pub const fn validation_errors(mut self, max_requests: u32, window_secs: u64) -> Self {
self.inner.validation_errors_max_requests = max_requests;
self.inner.validation_errors_window_secs = window_secs;
self
}
#[must_use = "builder method returns modified builder"]
pub const fn depth_errors(mut self, max_requests: u32, window_secs: u64) -> Self {
self.inner.depth_errors_max_requests = max_requests;
self.inner.depth_errors_window_secs = window_secs;
self
}
#[must_use = "builder method returns modified builder"]
pub const fn complexity_errors(mut self, max_requests: u32, window_secs: u64) -> Self {
self.inner.complexity_errors_max_requests = max_requests;
self.inner.complexity_errors_window_secs = window_secs;
self
}
#[must_use = "builder method returns modified builder"]
pub const fn malformed_errors(mut self, max_requests: u32, window_secs: u64) -> Self {
self.inner.malformed_errors_max_requests = max_requests;
self.inner.malformed_errors_window_secs = window_secs;
self
}
#[must_use = "builder method returns modified builder"]
pub const fn async_validation_errors(mut self, max_requests: u32, window_secs: u64) -> Self {
self.inner.async_validation_errors_max_requests = max_requests;
self.inner.async_validation_errors_window_secs = window_secs;
self
}
#[must_use = "building a config that is not used has no effect"]
pub const fn build(self) -> ValidationRateLimitingConfig {
self.inner
}
}
#[derive(Debug, Clone)]
struct RequestRecord {
count: u32,
window_start: u64,
}
pub(crate) struct DimensionRateLimiter {
records: Arc<Mutex<LruCache<String, RequestRecord>>>,
dimension: RateLimitDimension,
clock: Arc<dyn Clock>,
}
impl DimensionRateLimiter {
#[cfg(test)]
pub(crate) fn new(max_requests: u32, window_secs: u64) -> Self {
Self::new_with_clock(max_requests, window_secs, Arc::new(SystemClock))
}
pub(crate) fn new_with_clock(
max_requests: u32,
window_secs: u64,
clock: Arc<dyn Clock>,
) -> Self {
#[allow(clippy::expect_used)]
let cap = NonZeroUsize::new(MAX_RATE_LIMITER_ENTRIES)
.expect("MAX_RATE_LIMITER_ENTRIES must be > 0");
Self {
records: Arc::new(Mutex::new(LruCache::new(cap))),
dimension: RateLimitDimension {
max_requests,
window_secs,
},
clock,
}
}
pub(crate) fn check(&self, key: &str) -> Result<(), FraiseQLError> {
if self.dimension.is_rate_limited() {
return Ok(());
}
let mut records = self.records.lock().expect("records mutex poisoned");
let now = self.clock.now_secs();
let record = records.get_or_insert_mut(key.to_string(), || RequestRecord {
count: 0,
window_start: now,
});
if now >= record.window_start + self.dimension.window_secs {
record.count = 1;
record.window_start = now;
Ok(())
} else if record.count < self.dimension.max_requests {
record.count += 1;
Ok(())
} else {
Err(FraiseQLError::RateLimited {
message: "Rate limit exceeded for validation errors".to_string(),
retry_after_secs: self.dimension.window_secs,
})
}
}
pub(crate) fn clear(&self) {
let mut records = self.records.lock().expect("records mutex poisoned");
records.clear();
}
}
impl Clone for DimensionRateLimiter {
fn clone(&self) -> Self {
Self {
records: Arc::clone(&self.records),
dimension: self.dimension.clone(),
clock: Arc::clone(&self.clock),
}
}
}
#[derive(Clone)]
#[allow(clippy::module_name_repetitions, clippy::struct_field_names)] pub struct ValidationRateLimiter {
validation_errors: DimensionRateLimiter,
depth_errors: DimensionRateLimiter,
complexity_errors: DimensionRateLimiter,
malformed_errors: DimensionRateLimiter,
async_validation_errors: DimensionRateLimiter,
}
impl ValidationRateLimiter {
#[must_use]
pub fn new(config: &ValidationRateLimitingConfig) -> Self {
Self::new_with_clock(config, Arc::new(SystemClock))
}
pub fn new_with_clock(config: &ValidationRateLimitingConfig, clock: Arc<dyn Clock>) -> Self {
Self {
validation_errors: DimensionRateLimiter::new_with_clock(
config.validation_errors_max_requests,
config.validation_errors_window_secs,
Arc::clone(&clock),
),
depth_errors: DimensionRateLimiter::new_with_clock(
config.depth_errors_max_requests,
config.depth_errors_window_secs,
Arc::clone(&clock),
),
complexity_errors: DimensionRateLimiter::new_with_clock(
config.complexity_errors_max_requests,
config.complexity_errors_window_secs,
Arc::clone(&clock),
),
malformed_errors: DimensionRateLimiter::new_with_clock(
config.malformed_errors_max_requests,
config.malformed_errors_window_secs,
Arc::clone(&clock),
),
async_validation_errors: DimensionRateLimiter::new_with_clock(
config.async_validation_errors_max_requests,
config.async_validation_errors_window_secs,
clock,
),
}
}
pub fn check_validation_errors(&self, key: &str) -> Result<(), FraiseQLError> {
self.validation_errors.check(key)
}
pub fn check_depth_errors(&self, key: &str) -> Result<(), FraiseQLError> {
self.depth_errors.check(key)
}
pub fn check_complexity_errors(&self, key: &str) -> Result<(), FraiseQLError> {
self.complexity_errors.check(key)
}
pub fn check_malformed_errors(&self, key: &str) -> Result<(), FraiseQLError> {
self.malformed_errors.check(key)
}
pub fn check_async_validation_errors(&self, key: &str) -> Result<(), FraiseQLError> {
self.async_validation_errors.check(key)
}
pub fn clear(&self) {
self.validation_errors.clear();
self.depth_errors.clear();
self.complexity_errors.clear();
self.malformed_errors.clear();
self.async_validation_errors.clear();
}
}