1use anyhow::{Result, anyhow};
33use clap::Parser;
34use serde::Deserialize;
35
36#[derive(Debug, Clone, Deserialize)]
41pub struct Config {
42 pub transports: TransportConfig,
44 pub store: StoreConfig,
46 pub buffer_size: usize,
48 pub log_level: String,
50}
51
52#[derive(Debug, Clone, Deserialize)]
57pub struct TransportConfig {
58 pub http: Option<HttpConfig>,
60 pub grpc: Option<GrpcConfig>,
62 pub redis: Option<RedisConfig>,
64}
65
66#[derive(Debug, Clone, Deserialize)]
68pub struct HttpConfig {
69 pub host: String,
71 pub port: u16,
73}
74
75#[derive(Debug, Clone, Deserialize)]
77pub struct GrpcConfig {
78 pub host: String,
80 pub port: u16,
82}
83
84#[derive(Debug, Clone, Deserialize)]
86pub struct RedisConfig {
87 pub host: String,
89 pub port: u16,
91}
92
93#[derive(Debug, Clone, Deserialize)]
100pub struct StoreConfig {
101 pub store_type: StoreType,
103 pub capacity: usize,
105 pub cleanup_interval: u64,
108 pub cleanup_probability: u64,
110 pub min_interval: u64,
112 pub max_interval: u64,
114 pub max_operations: usize,
116}
117
118#[derive(Debug, Clone, Copy, Deserialize, PartialEq)]
125#[serde(rename_all = "lowercase")]
126pub enum StoreType {
127 Periodic,
129 Probabilistic,
131 Adaptive,
133}
134
135impl std::str::FromStr for StoreType {
136 type Err = anyhow::Error;
137
138 fn from_str(s: &str) -> Result<Self> {
139 match s.to_lowercase().as_str() {
140 "periodic" => Ok(StoreType::Periodic),
141 "probabilistic" => Ok(StoreType::Probabilistic),
142 "adaptive" => Ok(StoreType::Adaptive),
143 _ => Err(anyhow!(
144 "Invalid store type: {}. Valid options are: periodic, probabilistic, adaptive",
145 s
146 )),
147 }
148 }
149}
150
151#[derive(Parser, Debug)]
173#[command(
174 name = "throttlecrab-server",
175 version = env!("CARGO_PKG_VERSION"),
176 about = "High-performance rate limiting server",
177 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."
178)]
179pub struct Args {
180 #[arg(long, help = "Enable HTTP transport", env = "THROTTLECRAB_HTTP")]
182 pub http: bool,
183 #[arg(
184 long,
185 value_name = "HOST",
186 help = "HTTP host",
187 default_value = "127.0.0.1",
188 env = "THROTTLECRAB_HTTP_HOST"
189 )]
190 pub http_host: String,
191 #[arg(
192 long,
193 value_name = "PORT",
194 help = "HTTP port",
195 default_value_t = 8080,
196 env = "THROTTLECRAB_HTTP_PORT"
197 )]
198 pub http_port: u16,
199
200 #[arg(long, help = "Enable gRPC transport", env = "THROTTLECRAB_GRPC")]
202 pub grpc: bool,
203 #[arg(
204 long,
205 value_name = "HOST",
206 help = "gRPC host",
207 default_value = "127.0.0.1",
208 env = "THROTTLECRAB_GRPC_HOST"
209 )]
210 pub grpc_host: String,
211 #[arg(
212 long,
213 value_name = "PORT",
214 help = "gRPC port",
215 default_value_t = 8070,
216 env = "THROTTLECRAB_GRPC_PORT"
217 )]
218 pub grpc_port: u16,
219
220 #[arg(
222 long,
223 help = "Enable Redis protocol transport",
224 env = "THROTTLECRAB_REDIS"
225 )]
226 pub redis: bool,
227 #[arg(
228 long,
229 value_name = "HOST",
230 help = "Redis host",
231 default_value = "127.0.0.1",
232 env = "THROTTLECRAB_REDIS_HOST"
233 )]
234 pub redis_host: String,
235 #[arg(
236 long,
237 value_name = "PORT",
238 help = "Redis port",
239 default_value_t = 6379,
240 env = "THROTTLECRAB_REDIS_PORT"
241 )]
242 pub redis_port: u16,
243
244 #[arg(
246 long,
247 value_name = "TYPE",
248 help = "Store type: periodic, probabilistic, adaptive",
249 default_value = "periodic",
250 env = "THROTTLECRAB_STORE"
251 )]
252 pub store: StoreType,
253 #[arg(
254 long,
255 value_name = "SIZE",
256 help = "Initial store capacity",
257 default_value_t = 100_000,
258 env = "THROTTLECRAB_STORE_CAPACITY"
259 )]
260 pub store_capacity: usize,
261
262 #[arg(
264 long,
265 value_name = "SECS",
266 help = "Cleanup interval for periodic store (seconds)",
267 default_value_t = 300,
268 env = "THROTTLECRAB_STORE_CLEANUP_INTERVAL"
269 )]
270 pub store_cleanup_interval: u64,
271 #[arg(
272 long,
273 value_name = "N",
274 help = "Cleanup probability for probabilistic store (1 in N)",
275 default_value_t = 10_000,
276 env = "THROTTLECRAB_STORE_CLEANUP_PROBABILITY"
277 )]
278 pub store_cleanup_probability: u64,
279 #[arg(
280 long,
281 value_name = "SECS",
282 help = "Minimum cleanup interval for adaptive store (seconds)",
283 default_value_t = 5,
284 env = "THROTTLECRAB_STORE_MIN_INTERVAL"
285 )]
286 pub store_min_interval: u64,
287 #[arg(
288 long,
289 value_name = "SECS",
290 help = "Maximum cleanup interval for adaptive store (seconds)",
291 default_value_t = 300,
292 env = "THROTTLECRAB_STORE_MAX_INTERVAL"
293 )]
294 pub store_max_interval: u64,
295 #[arg(
296 long,
297 value_name = "N",
298 help = "Maximum operations before cleanup for adaptive store",
299 default_value_t = 1_000_000,
300 env = "THROTTLECRAB_STORE_MAX_OPERATIONS"
301 )]
302 pub store_max_operations: usize,
303
304 #[arg(
306 long,
307 value_name = "SIZE",
308 help = "Channel buffer size",
309 default_value_t = 100_000,
310 env = "THROTTLECRAB_BUFFER_SIZE"
311 )]
312 pub buffer_size: usize,
313 #[arg(
314 long,
315 value_name = "LEVEL",
316 help = "Log level: error, warn, info, debug, trace",
317 default_value = "info",
318 env = "THROTTLECRAB_LOG_LEVEL"
319 )]
320 pub log_level: String,
321
322 #[arg(
324 long,
325 help = "List all environment variables and exit",
326 action = clap::ArgAction::SetTrue
327 )]
328 pub list_env_vars: bool,
329}
330
331impl Config {
332 pub fn from_env_and_args() -> Result<Self> {
346 let args = Args::parse();
351
352 if args.list_env_vars {
354 Self::print_env_vars();
355 std::process::exit(0);
356 }
357
358 let mut config = Config {
360 transports: TransportConfig {
361 http: None,
362 grpc: None,
363 redis: None,
364 },
365 store: StoreConfig {
366 store_type: args.store,
367 capacity: args.store_capacity,
368 cleanup_interval: args.store_cleanup_interval,
369 cleanup_probability: args.store_cleanup_probability,
370 min_interval: args.store_min_interval,
371 max_interval: args.store_max_interval,
372 max_operations: args.store_max_operations,
373 },
374 buffer_size: args.buffer_size,
375 log_level: args.log_level,
376 };
377
378 if args.http {
380 config.transports.http = Some(HttpConfig {
381 host: args.http_host,
382 port: args.http_port,
383 });
384 }
385
386 if args.grpc {
387 config.transports.grpc = Some(GrpcConfig {
388 host: args.grpc_host,
389 port: args.grpc_port,
390 });
391 }
392
393 if args.redis {
394 config.transports.redis = Some(RedisConfig {
395 host: args.redis_host,
396 port: args.redis_port,
397 });
398 }
399
400 config.validate()?;
402
403 Ok(config)
404 }
405
406 pub fn has_any_transport(&self) -> bool {
410 self.transports.http.is_some()
411 || self.transports.grpc.is_some()
412 || self.transports.redis.is_some()
413 }
414
415 fn validate(&self) -> Result<()> {
424 if !self.has_any_transport() {
425 return Err(anyhow!(
426 "At least one transport must be specified.\n\n\
427 Available transports:\n \
428 --http Enable HTTP transport\n \
429 --grpc Enable gRPC transport\n \
430 --redis Enable Redis protocol transport\n \
431 Example:\n \
432 throttlecrab-server --http --http-port 7070\n \
433 throttlecrab-server --http --grpc --redis\n\n\
434 For more information, try '--help'"
435 ));
436 }
437
438 Ok(())
442 }
443
444 fn print_env_vars() {
450 println!("ThrottleCrab Environment Variables");
451 println!("==================================");
452 println!();
453 println!("All environment variables use the THROTTLECRAB_ prefix.");
454 println!("CLI arguments take precedence over environment variables.");
455 println!();
456
457 println!("Transport Configuration:");
458 println!(" THROTTLECRAB_HTTP=true|false Enable HTTP transport");
459 println!(" THROTTLECRAB_HTTP_HOST=<host> HTTP host [default: 127.0.0.1]");
460 println!(" THROTTLECRAB_HTTP_PORT=<port> HTTP port [default: 8080]");
461 println!();
462 println!(" THROTTLECRAB_GRPC=true|false Enable gRPC transport");
463 println!(" THROTTLECRAB_GRPC_HOST=<host> gRPC host [default: 127.0.0.1]");
464 println!(" THROTTLECRAB_GRPC_PORT=<port> gRPC port [default: 8070]");
465 println!();
466 println!(" THROTTLECRAB_REDIS=true|false Enable Redis protocol transport");
467 println!(" THROTTLECRAB_REDIS_HOST=<host> Redis host [default: 127.0.0.1]");
468 println!(" THROTTLECRAB_REDIS_PORT=<port> Redis port [default: 6379]");
469 println!();
470
471 println!("Store Configuration:");
472 println!(
473 " THROTTLECRAB_STORE=<type> Store type: periodic, probabilistic, adaptive [default: periodic]"
474 );
475 println!(
476 " THROTTLECRAB_STORE_CAPACITY=<size> Initial store capacity [default: 100000]"
477 );
478 println!();
479 println!(" For periodic store:");
480 println!(
481 " THROTTLECRAB_STORE_CLEANUP_INTERVAL=<secs> Cleanup interval in seconds [default: 300]"
482 );
483 println!();
484 println!(" For probabilistic store:");
485 println!(
486 " THROTTLECRAB_STORE_CLEANUP_PROBABILITY=<n> Cleanup probability (1 in N) [default: 10000]"
487 );
488 println!();
489 println!(" For adaptive store:");
490 println!(
491 " THROTTLECRAB_STORE_MIN_INTERVAL=<secs> Minimum cleanup interval [default: 5]"
492 );
493 println!(
494 " THROTTLECRAB_STORE_MAX_INTERVAL=<secs> Maximum cleanup interval [default: 300]"
495 );
496 println!(
497 " THROTTLECRAB_STORE_MAX_OPERATIONS=<n> Max operations before cleanup [default: 1000000]"
498 );
499 println!();
500
501 println!("General Configuration:");
502 println!(" THROTTLECRAB_BUFFER_SIZE=<size> Channel buffer size [default: 100000]");
503 println!(
504 " THROTTLECRAB_LOG_LEVEL=<level> Log level: error, warn, info, debug, trace [default: info]"
505 );
506 println!();
507
508 println!("Examples:");
509 println!(" # Enable HTTP transport on port 8080");
510 println!(" export THROTTLECRAB_HTTP=true");
511 println!(" export THROTTLECRAB_HTTP_PORT=8080");
512 println!();
513 println!(" # Use adaptive store with custom settings");
514 println!(" export THROTTLECRAB_STORE=adaptive");
515 println!(" export THROTTLECRAB_STORE_MIN_INTERVAL=10");
516 println!(" export THROTTLECRAB_STORE_MAX_INTERVAL=600");
517 println!();
518 println!(" # Run server (CLI args override env vars)");
519 println!(" throttlecrab-server --http-port 9090 # Will use port 9090, not 8080");
520 }
521}
522
523#[cfg(test)]
524mod tests {
525 use super::*;
526 use std::str::FromStr;
527
528 #[test]
529 fn test_store_type_from_str() {
530 assert_eq!(
531 StoreType::from_str("periodic").unwrap(),
532 StoreType::Periodic
533 );
534 assert_eq!(
535 StoreType::from_str("PERIODIC").unwrap(),
536 StoreType::Periodic
537 );
538 assert_eq!(
539 StoreType::from_str("probabilistic").unwrap(),
540 StoreType::Probabilistic
541 );
542 assert_eq!(
543 StoreType::from_str("adaptive").unwrap(),
544 StoreType::Adaptive
545 );
546 assert!(StoreType::from_str("invalid").is_err());
547 }
548
549 #[test]
550 fn test_config_validation_no_transport() {
551 let config = Config {
552 transports: TransportConfig {
553 http: None,
554 grpc: None,
555 redis: None,
556 },
557 store: StoreConfig {
558 store_type: StoreType::Periodic,
559 capacity: 100_000,
560 cleanup_interval: 300,
561 cleanup_probability: 10_000,
562 min_interval: 5,
563 max_interval: 300,
564 max_operations: 1_000_000,
565 },
566 buffer_size: 100_000,
567 log_level: "info".to_string(),
568 };
569
570 assert!(config.validate().is_err());
571 assert!(!config.has_any_transport());
572 }
573
574 #[test]
575 fn test_config_validation_with_transport() {
576 let config = Config {
577 transports: TransportConfig {
578 http: Some(HttpConfig {
579 host: "127.0.0.1".to_string(),
580 port: 8080,
581 }),
582 grpc: None,
583 redis: None,
584 },
585 store: StoreConfig {
586 store_type: StoreType::Periodic,
587 capacity: 100_000,
588 cleanup_interval: 300,
589 cleanup_probability: 10_000,
590 min_interval: 5,
591 max_interval: 300,
592 max_operations: 1_000_000,
593 },
594 buffer_size: 100_000,
595 log_level: "info".to_string(),
596 };
597
598 assert!(config.validate().is_ok());
599 assert!(config.has_any_transport());
600 }
601
602 #[test]
603 fn test_config_multiple_transports() {
604 let config = Config {
605 transports: TransportConfig {
606 http: Some(HttpConfig {
607 host: "0.0.0.0".to_string(),
608 port: 8080,
609 }),
610 grpc: Some(GrpcConfig {
611 host: "0.0.0.0".to_string(),
612 port: 50051,
613 }),
614 redis: None,
615 },
616 store: StoreConfig {
617 store_type: StoreType::Adaptive,
618 capacity: 200_000,
619 cleanup_interval: 300,
620 cleanup_probability: 10_000,
621 min_interval: 10,
622 max_interval: 600,
623 max_operations: 2_000_000,
624 },
625 buffer_size: 50_000,
626 log_level: "debug".to_string(),
627 };
628
629 assert!(config.validate().is_ok());
630 assert!(config.has_any_transport());
631 }
632}