use std::time::Duration;
#[derive(Debug, Clone)]
pub struct QuerySecurityConfig {
timeout: Duration,
result_cap: usize,
memory_limit: usize,
audit_enabled: bool,
cost_limit: usize,
}
impl QuerySecurityConfig {
pub const HARD_CEILING_TIMEOUT: Duration = Duration::from_secs(30);
pub const DEFAULT_RESULT_CAP: usize = 10_000;
pub const DEFAULT_MEMORY_LIMIT: usize = 512 * 1024 * 1024;
pub const DEFAULT_COST_LIMIT: usize = 1_000_000;
#[must_use]
pub fn interactive() -> Self {
Self {
timeout: Duration::from_secs(10), ..Default::default()
}
}
#[must_use]
pub fn with_timeout(self, timeout: Duration) -> Self {
Self {
timeout: timeout.min(Self::HARD_CEILING_TIMEOUT),
..self
}
}
#[must_use]
pub fn with_result_cap(self, cap: usize) -> Self {
Self {
result_cap: cap,
..self
}
}
#[must_use]
pub fn with_memory_limit(self, limit: usize) -> Self {
Self {
memory_limit: limit,
..self
}
}
#[must_use]
pub fn with_cost_limit(self, limit: usize) -> Self {
Self {
cost_limit: limit,
..self
}
}
#[must_use]
pub fn without_audit(self) -> Self {
Self {
audit_enabled: false,
..self
}
}
#[must_use]
pub fn timeout(&self) -> Duration {
self.timeout
}
#[must_use]
pub fn result_cap(&self) -> usize {
self.result_cap
}
#[must_use]
pub fn memory_limit(&self) -> usize {
self.memory_limit
}
#[must_use]
pub fn audit_enabled(&self) -> bool {
self.audit_enabled
}
#[must_use]
pub fn cost_limit(&self) -> usize {
self.cost_limit
}
}
impl Default for QuerySecurityConfig {
fn default() -> Self {
Self {
timeout: Self::HARD_CEILING_TIMEOUT,
result_cap: Self::DEFAULT_RESULT_CAP,
memory_limit: Self::DEFAULT_MEMORY_LIMIT,
audit_enabled: true,
cost_limit: Self::DEFAULT_COST_LIMIT,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config() {
let config = QuerySecurityConfig::default();
assert_eq!(config.timeout(), Duration::from_secs(30));
assert_eq!(config.result_cap(), 10_000);
assert_eq!(config.memory_limit(), 512 * 1024 * 1024);
assert!(config.audit_enabled());
assert_eq!(config.cost_limit(), 1_000_000);
}
#[test]
fn test_interactive_config() {
let config = QuerySecurityConfig::interactive();
assert_eq!(config.timeout(), Duration::from_secs(10));
assert_eq!(config.result_cap(), 10_000);
}
#[test]
fn test_timeout_capped_at_30s() {
let config = QuerySecurityConfig::default().with_timeout(Duration::from_secs(60));
assert_eq!(config.timeout(), Duration::from_secs(30));
}
#[test]
fn test_timeout_under_ceiling() {
let config = QuerySecurityConfig::default().with_timeout(Duration::from_secs(5));
assert_eq!(config.timeout(), Duration::from_secs(5));
}
#[test]
fn test_builder_chain() {
let config = QuerySecurityConfig::default()
.with_timeout(Duration::from_secs(15))
.with_result_cap(500)
.with_memory_limit(64 * 1024 * 1024)
.without_audit();
assert_eq!(config.timeout(), Duration::from_secs(15));
assert_eq!(config.result_cap(), 500);
assert_eq!(config.memory_limit(), 64 * 1024 * 1024);
assert!(!config.audit_enabled());
}
#[test]
fn test_hard_ceiling_constant() {
assert_eq!(
QuerySecurityConfig::HARD_CEILING_TIMEOUT,
Duration::from_secs(30)
);
}
}