use crate::error::{AppError, ResourceLimitError};
use std::time::Duration;
pub const DEFAULT_MAX_INPUT: u64 = 16 * 1024 * 1024 * 1024;
pub const DEFAULT_MAX_OUTPUT: u64 = 16 * 1024 * 1024 * 1024;
pub const DEFAULT_BLOCK_SIZE: u32 = 1024 * 1024;
pub const HARD_MAX_BLOCK_SIZE: u32 = 64 * 1024 * 1024;
pub const DEFAULT_MAX_PAYLOAD: u32 = 1024 * 1024;
pub const DEFAULT_MAX_MODEL: u32 = 2048;
pub const DEFAULT_MAX_BLOCKS: u64 = 1_000_000;
pub const DEFAULT_MAX_TRACE_SYMBOLS: u64 = 256;
pub const DEFAULT_MAX_ORACLE_OUTPUT: u64 = 64 * 1024 * 1024;
pub const DEFAULT_ORACLE_TIMEOUT: Duration = Duration::from_secs(60);
#[derive(Debug, Clone)]
pub struct Limits {
pub max_input_bytes: u64,
pub max_output_bytes: u64,
pub max_block_bytes: u32,
pub max_payload_bytes: u32,
pub max_model_bytes: u32,
pub max_blocks: u64,
pub max_trace_symbols: u64,
pub max_oracle_output_bytes: u64,
pub oracle_timeout: Duration,
}
impl Default for Limits {
fn default() -> Self {
Self {
max_input_bytes: DEFAULT_MAX_INPUT,
max_output_bytes: DEFAULT_MAX_OUTPUT,
max_block_bytes: DEFAULT_BLOCK_SIZE,
max_payload_bytes: DEFAULT_MAX_PAYLOAD,
max_model_bytes: DEFAULT_MAX_MODEL,
max_blocks: DEFAULT_MAX_BLOCKS,
max_trace_symbols: DEFAULT_MAX_TRACE_SYMBOLS,
max_oracle_output_bytes: DEFAULT_MAX_ORACLE_OUTPUT,
oracle_timeout: DEFAULT_ORACLE_TIMEOUT,
}
}
}
impl Limits {
pub fn check_block_size(&self, size: u32) -> Result<(), AppError> {
if size > self.max_block_bytes {
return Err(AppError::ResourceLimit(ResourceLimitError {
detail: "block size exceeds limit".into(),
limit: self.max_block_bytes as u64,
requested: size as u64,
}));
}
Ok(())
}
pub fn check_payload_size(&self, size: u32) -> Result<(), AppError> {
if size > self.max_payload_bytes {
return Err(AppError::ResourceLimit(ResourceLimitError {
detail: "payload size exceeds limit".into(),
limit: self.max_payload_bytes as u64,
requested: size as u64,
}));
}
Ok(())
}
pub fn check_model_size(&self, size: u32) -> Result<(), AppError> {
if size > self.max_model_bytes {
return Err(AppError::ResourceLimit(ResourceLimitError {
detail: "model size exceeds limit".into(),
limit: self.max_model_bytes as u64,
requested: size as u64,
}));
}
Ok(())
}
pub fn check_block_count(&self, count: u64) -> Result<(), AppError> {
if count > self.max_blocks {
return Err(AppError::ResourceLimit(ResourceLimitError {
detail: "block count exceeds limit".into(),
limit: self.max_blocks,
requested: count,
}));
}
Ok(())
}
pub fn check_output_total(&self, current: u64, additional: u64) -> Result<u64, AppError> {
let new_total = current.checked_add(additional).ok_or_else(|| {
AppError::ResourceLimit(ResourceLimitError {
detail: "output total overflow".into(),
limit: self.max_output_bytes,
requested: current.saturating_add(additional),
})
})?;
if new_total > self.max_output_bytes {
return Err(AppError::ResourceLimit(ResourceLimitError {
detail: "output total exceeds limit".into(),
limit: self.max_output_bytes,
requested: new_total,
}));
}
Ok(new_total)
}
pub fn parse_size(s: &str) -> Result<u64, String> {
let s = s.trim();
if s.is_empty() {
return Err("empty size string".into());
}
let num_end = s.find(|c: char| !c.is_ascii_digit()).unwrap_or(s.len());
let num_part = &s[..num_end];
let suffix = s[num_end..].trim();
let value: u64 = num_part
.parse()
.map_err(|_| format!("invalid number: '{}'", num_part))?;
match suffix.to_lowercase().as_str() {
"" | "b" => Ok(value),
"kib" | "ki" => value.checked_mul(1024).ok_or("overflow".into()),
"mib" | "mi" => value.checked_mul(1024 * 1024).ok_or("overflow".into()),
"gib" | "gi" => value
.checked_mul(1024 * 1024 * 1024)
.ok_or("overflow".into()),
_ => Err(format!("unknown size suffix: '{}'", suffix)),
}
}
}