use anyhow::{Result, anyhow};
use clap::Parser;
use serde::Deserialize;
#[derive(Debug, Clone, Deserialize)]
pub struct Config {
pub transports: TransportConfig,
pub store: StoreConfig,
pub buffer_size: usize,
pub max_denied_keys: u32,
pub log_level: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct TransportConfig {
pub http: Option<HttpConfig>,
pub grpc: Option<GrpcConfig>,
pub redis: Option<RedisConfig>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct HttpConfig {
pub host: String,
pub port: u16,
}
#[derive(Debug, Clone, Deserialize)]
pub struct GrpcConfig {
pub host: String,
pub port: u16,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RedisConfig {
pub host: String,
pub port: u16,
}
#[derive(Debug, Clone, Deserialize)]
pub struct StoreConfig {
pub store_type: StoreType,
pub capacity: usize,
pub cleanup_interval: u64,
pub cleanup_probability: u64,
pub min_interval: u64,
pub max_interval: u64,
pub max_operations: usize,
}
#[derive(Debug, Clone, Copy, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum StoreType {
Periodic,
Probabilistic,
Adaptive,
}
impl std::str::FromStr for StoreType {
type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self> {
match s.to_lowercase().as_str() {
"periodic" => Ok(StoreType::Periodic),
"probabilistic" => Ok(StoreType::Probabilistic),
"adaptive" => Ok(StoreType::Adaptive),
_ => Err(anyhow!(
"Invalid store type: {}. Valid options are: periodic, probabilistic, adaptive",
s
)),
}
}
}
#[derive(Parser, Debug)]
#[command(
name = "throttlecrab-server",
version = env!("CARGO_PKG_VERSION"),
about = "High-performance rate limiting server",
long_about = "A high-performance rate limiting server with multiple protocol support.\n\nAt least one transport must be specified.\n\nEnvironment variables with THROTTLECRAB_ prefix are supported. CLI arguments take precedence over environment variables."
)]
pub struct Args {
#[arg(long, help = "Enable HTTP transport", env = "THROTTLECRAB_HTTP")]
pub http: bool,
#[arg(
long,
value_name = "HOST",
help = "HTTP host",
default_value = "0.0.0.0",
env = "THROTTLECRAB_HTTP_HOST"
)]
pub http_host: String,
#[arg(
long,
value_name = "PORT",
help = "HTTP port",
default_value_t = 8080,
env = "THROTTLECRAB_HTTP_PORT"
)]
pub http_port: u16,
#[arg(long, help = "Enable gRPC transport", env = "THROTTLECRAB_GRPC")]
pub grpc: bool,
#[arg(
long,
value_name = "HOST",
help = "gRPC host",
default_value = "0.0.0.0",
env = "THROTTLECRAB_GRPC_HOST"
)]
pub grpc_host: String,
#[arg(
long,
value_name = "PORT",
help = "gRPC port",
default_value_t = 8070,
env = "THROTTLECRAB_GRPC_PORT"
)]
pub grpc_port: u16,
#[arg(
long,
help = "Enable Redis protocol transport",
env = "THROTTLECRAB_REDIS"
)]
pub redis: bool,
#[arg(
long,
value_name = "HOST",
help = "Redis host",
default_value = "0.0.0.0",
env = "THROTTLECRAB_REDIS_HOST"
)]
pub redis_host: String,
#[arg(
long,
value_name = "PORT",
help = "Redis port",
default_value_t = 6379,
env = "THROTTLECRAB_REDIS_PORT"
)]
pub redis_port: u16,
#[arg(
long,
value_name = "TYPE",
help = "Store type: periodic, probabilistic, adaptive",
default_value = "periodic",
env = "THROTTLECRAB_STORE"
)]
pub store: StoreType,
#[arg(
long,
value_name = "SIZE",
help = "Initial store capacity",
default_value_t = 100_000,
env = "THROTTLECRAB_STORE_CAPACITY"
)]
pub store_capacity: usize,
#[arg(
long,
value_name = "SECS",
help = "Cleanup interval for periodic store (seconds)",
default_value_t = 300,
env = "THROTTLECRAB_STORE_CLEANUP_INTERVAL"
)]
pub store_cleanup_interval: u64,
#[arg(
long,
value_name = "N",
help = "Cleanup probability for probabilistic store (1 in N)",
default_value_t = 10_000,
env = "THROTTLECRAB_STORE_CLEANUP_PROBABILITY"
)]
pub store_cleanup_probability: u64,
#[arg(
long,
value_name = "SECS",
help = "Minimum cleanup interval for adaptive store (seconds)",
default_value_t = 5,
env = "THROTTLECRAB_STORE_MIN_INTERVAL"
)]
pub store_min_interval: u64,
#[arg(
long,
value_name = "SECS",
help = "Maximum cleanup interval for adaptive store (seconds)",
default_value_t = 300,
env = "THROTTLECRAB_STORE_MAX_INTERVAL"
)]
pub store_max_interval: u64,
#[arg(
long,
value_name = "N",
help = "Maximum operations before cleanup for adaptive store",
default_value_t = 1_000_000,
env = "THROTTLECRAB_STORE_MAX_OPERATIONS"
)]
pub store_max_operations: usize,
#[arg(
long,
value_name = "SIZE",
help = "Channel buffer size",
default_value_t = 100_000,
env = "THROTTLECRAB_BUFFER_SIZE"
)]
pub buffer_size: usize,
#[arg(
long,
value_name = "COUNT",
help = "Maximum number of denied keys to track in metrics (0 to disable, max: 10000)",
default_value_t = 100,
env = "THROTTLECRAB_MAX_DENIED_KEYS",
value_parser = clap::value_parser!(u32).range(0..=10000)
)]
pub max_denied_keys: u32,
#[arg(
long,
value_name = "LEVEL",
help = "Log level: error, warn, info, debug, trace",
default_value = "info",
env = "THROTTLECRAB_LOG_LEVEL"
)]
pub log_level: String,
#[arg(
long,
help = "List all environment variables and exit",
action = clap::ArgAction::SetTrue
)]
pub list_env_vars: bool,
}
impl Config {
pub fn from_env_and_args() -> Result<Self> {
let args = Args::parse();
if args.list_env_vars {
Self::print_env_vars();
std::process::exit(0);
}
let mut config = Config {
transports: TransportConfig {
http: None,
grpc: None,
redis: None,
},
store: StoreConfig {
store_type: args.store,
capacity: args.store_capacity,
cleanup_interval: args.store_cleanup_interval,
cleanup_probability: args.store_cleanup_probability,
min_interval: args.store_min_interval,
max_interval: args.store_max_interval,
max_operations: args.store_max_operations,
},
buffer_size: args.buffer_size,
max_denied_keys: args.max_denied_keys,
log_level: args.log_level,
};
if args.http {
config.transports.http = Some(HttpConfig {
host: args.http_host,
port: args.http_port,
});
}
if args.grpc {
config.transports.grpc = Some(GrpcConfig {
host: args.grpc_host,
port: args.grpc_port,
});
}
if args.redis {
config.transports.redis = Some(RedisConfig {
host: args.redis_host,
port: args.redis_port,
});
}
config.validate()?;
Ok(config)
}
pub fn has_any_transport(&self) -> bool {
self.transports.http.is_some()
|| self.transports.grpc.is_some()
|| self.transports.redis.is_some()
}
fn validate(&self) -> Result<()> {
if !self.has_any_transport() {
return Err(anyhow!(
"At least one transport must be specified.\n\n\
Available transports:\n \
--http Enable HTTP transport\n \
--grpc Enable gRPC transport\n \
--redis Enable Redis protocol transport\n \
Example:\n \
throttlecrab-server --http --http-port 7070\n \
throttlecrab-server --http --grpc --redis\n\n\
For more information, try '--help'"
));
}
Ok(())
}
fn print_env_vars() {
println!("ThrottleCrab Environment Variables");
println!("==================================");
println!();
println!("All environment variables use the THROTTLECRAB_ prefix.");
println!("CLI arguments take precedence over environment variables.");
println!();
println!("Transport Configuration:");
println!(" THROTTLECRAB_HTTP=true|false Enable HTTP transport");
println!(" THROTTLECRAB_HTTP_HOST=<host> HTTP host [default: 0.0.0.0]");
println!(" THROTTLECRAB_HTTP_PORT=<port> HTTP port [default: 8080]");
println!();
println!(" THROTTLECRAB_GRPC=true|false Enable gRPC transport");
println!(" THROTTLECRAB_GRPC_HOST=<host> gRPC host [default: 0.0.0.0]");
println!(" THROTTLECRAB_GRPC_PORT=<port> gRPC port [default: 8070]");
println!();
println!(" THROTTLECRAB_REDIS=true|false Enable Redis protocol transport");
println!(" THROTTLECRAB_REDIS_HOST=<host> Redis host [default: 0.0.0.0]");
println!(" THROTTLECRAB_REDIS_PORT=<port> Redis port [default: 6379]");
println!();
println!("Store Configuration:");
println!(
" THROTTLECRAB_STORE=<type> Store type: periodic, probabilistic, adaptive [default: periodic]"
);
println!(
" THROTTLECRAB_STORE_CAPACITY=<size> Initial store capacity [default: 100000]"
);
println!();
println!(" For periodic store:");
println!(
" THROTTLECRAB_STORE_CLEANUP_INTERVAL=<secs> Cleanup interval in seconds [default: 300]"
);
println!();
println!(" For probabilistic store:");
println!(
" THROTTLECRAB_STORE_CLEANUP_PROBABILITY=<n> Cleanup probability (1 in N) [default: 10000]"
);
println!();
println!(" For adaptive store:");
println!(
" THROTTLECRAB_STORE_MIN_INTERVAL=<secs> Minimum cleanup interval [default: 5]"
);
println!(
" THROTTLECRAB_STORE_MAX_INTERVAL=<secs> Maximum cleanup interval [default: 300]"
);
println!(
" THROTTLECRAB_STORE_MAX_OPERATIONS=<n> Max operations before cleanup [default: 1000000]"
);
println!();
println!("General Configuration:");
println!(" THROTTLECRAB_BUFFER_SIZE=<size> Channel buffer size [default: 100000]");
println!(
" THROTTLECRAB_MAX_DENIED_KEYS=<count> Maximum denied keys to track (0=disabled, max: 10000) [default: 100]"
);
println!(
" THROTTLECRAB_LOG_LEVEL=<level> Log level: error, warn, info, debug, trace [default: info]"
);
println!();
println!("Examples:");
println!(" # Enable HTTP transport on port 8080");
println!(" export THROTTLECRAB_HTTP=true");
println!(" export THROTTLECRAB_HTTP_PORT=8080");
println!();
println!(" # Use adaptive store with custom settings");
println!(" export THROTTLECRAB_STORE=adaptive");
println!(" export THROTTLECRAB_STORE_MIN_INTERVAL=10");
println!(" export THROTTLECRAB_STORE_MAX_INTERVAL=600");
println!();
println!(" # Run server (CLI args override env vars)");
println!(" throttlecrab-server --http-port 9090 # Will use port 9090, not 8080");
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
#[test]
fn test_store_type_from_str() {
assert_eq!(
StoreType::from_str("periodic").unwrap(),
StoreType::Periodic
);
assert_eq!(
StoreType::from_str("PERIODIC").unwrap(),
StoreType::Periodic
);
assert_eq!(
StoreType::from_str("probabilistic").unwrap(),
StoreType::Probabilistic
);
assert_eq!(
StoreType::from_str("adaptive").unwrap(),
StoreType::Adaptive
);
assert!(StoreType::from_str("invalid").is_err());
}
#[test]
fn test_config_validation_no_transport() {
let config = Config {
transports: TransportConfig {
http: None,
grpc: None,
redis: None,
},
store: StoreConfig {
store_type: StoreType::Periodic,
capacity: 100_000,
cleanup_interval: 300,
cleanup_probability: 10_000,
min_interval: 5,
max_interval: 300,
max_operations: 1_000_000,
},
buffer_size: 100_000,
max_denied_keys: 100,
log_level: "info".to_string(),
};
assert!(config.validate().is_err());
assert!(!config.has_any_transport());
}
#[test]
fn test_config_validation_with_transport() {
let config = Config {
transports: TransportConfig {
http: Some(HttpConfig {
host: "0.0.0.0".to_string(),
port: 8080,
}),
grpc: None,
redis: None,
},
store: StoreConfig {
store_type: StoreType::Periodic,
capacity: 100_000,
cleanup_interval: 300,
cleanup_probability: 10_000,
min_interval: 5,
max_interval: 300,
max_operations: 1_000_000,
},
buffer_size: 100_000,
max_denied_keys: 100,
log_level: "info".to_string(),
};
assert!(config.validate().is_ok());
assert!(config.has_any_transport());
}
#[test]
fn test_config_multiple_transports() {
let config = Config {
transports: TransportConfig {
http: Some(HttpConfig {
host: "0.0.0.0".to_string(),
port: 8080,
}),
grpc: Some(GrpcConfig {
host: "0.0.0.0".to_string(),
port: 50051,
}),
redis: None,
},
store: StoreConfig {
store_type: StoreType::Adaptive,
capacity: 200_000,
cleanup_interval: 300,
cleanup_probability: 10_000,
min_interval: 10,
max_interval: 600,
max_operations: 2_000_000,
},
buffer_size: 50_000,
max_denied_keys: 100,
log_level: "debug".to_string(),
};
assert!(config.validate().is_ok());
assert!(config.has_any_transport());
}
}