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 native: Option<NativeConfig>,
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 NativeConfig {
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(long, help = "Enable Native transport", env = "THROTTLECRAB_NATIVE")]
222 pub native: bool,
223 #[arg(
224 long,
225 value_name = "HOST",
226 help = "Native host",
227 default_value = "127.0.0.1",
228 env = "THROTTLECRAB_NATIVE_HOST"
229 )]
230 pub native_host: String,
231 #[arg(
232 long,
233 value_name = "PORT",
234 help = "Native port",
235 default_value_t = 8072,
236 env = "THROTTLECRAB_NATIVE_PORT"
237 )]
238 pub native_port: u16,
239
240 #[arg(
242 long,
243 value_name = "TYPE",
244 help = "Store type: periodic, probabilistic, adaptive",
245 default_value = "periodic",
246 env = "THROTTLECRAB_STORE"
247 )]
248 pub store: StoreType,
249 #[arg(
250 long,
251 value_name = "SIZE",
252 help = "Initial store capacity",
253 default_value_t = 100_000,
254 env = "THROTTLECRAB_STORE_CAPACITY"
255 )]
256 pub store_capacity: usize,
257
258 #[arg(
260 long,
261 value_name = "SECS",
262 help = "Cleanup interval for periodic store (seconds)",
263 default_value_t = 300,
264 env = "THROTTLECRAB_STORE_CLEANUP_INTERVAL"
265 )]
266 pub store_cleanup_interval: u64,
267 #[arg(
268 long,
269 value_name = "N",
270 help = "Cleanup probability for probabilistic store (1 in N)",
271 default_value_t = 10_000,
272 env = "THROTTLECRAB_STORE_CLEANUP_PROBABILITY"
273 )]
274 pub store_cleanup_probability: u64,
275 #[arg(
276 long,
277 value_name = "SECS",
278 help = "Minimum cleanup interval for adaptive store (seconds)",
279 default_value_t = 5,
280 env = "THROTTLECRAB_STORE_MIN_INTERVAL"
281 )]
282 pub store_min_interval: u64,
283 #[arg(
284 long,
285 value_name = "SECS",
286 help = "Maximum cleanup interval for adaptive store (seconds)",
287 default_value_t = 300,
288 env = "THROTTLECRAB_STORE_MAX_INTERVAL"
289 )]
290 pub store_max_interval: u64,
291 #[arg(
292 long,
293 value_name = "N",
294 help = "Maximum operations before cleanup for adaptive store",
295 default_value_t = 1_000_000,
296 env = "THROTTLECRAB_STORE_MAX_OPERATIONS"
297 )]
298 pub store_max_operations: usize,
299
300 #[arg(
302 long,
303 value_name = "SIZE",
304 help = "Channel buffer size",
305 default_value_t = 100_000,
306 env = "THROTTLECRAB_BUFFER_SIZE"
307 )]
308 pub buffer_size: usize,
309 #[arg(
310 long,
311 value_name = "LEVEL",
312 help = "Log level: error, warn, info, debug, trace",
313 default_value = "info",
314 env = "THROTTLECRAB_LOG_LEVEL"
315 )]
316 pub log_level: String,
317
318 #[arg(
320 long,
321 help = "List all environment variables and exit",
322 action = clap::ArgAction::SetTrue
323 )]
324 pub list_env_vars: bool,
325}
326
327impl Config {
328 pub fn from_env_and_args() -> Result<Self> {
342 let args = Args::parse();
347
348 if args.list_env_vars {
350 Self::print_env_vars();
351 std::process::exit(0);
352 }
353
354 let mut config = Config {
356 transports: TransportConfig {
357 http: None,
358 grpc: None,
359 native: None,
360 },
361 store: StoreConfig {
362 store_type: args.store,
363 capacity: args.store_capacity,
364 cleanup_interval: args.store_cleanup_interval,
365 cleanup_probability: args.store_cleanup_probability,
366 min_interval: args.store_min_interval,
367 max_interval: args.store_max_interval,
368 max_operations: args.store_max_operations,
369 },
370 buffer_size: args.buffer_size,
371 log_level: args.log_level,
372 };
373
374 if args.http {
376 config.transports.http = Some(HttpConfig {
377 host: args.http_host,
378 port: args.http_port,
379 });
380 }
381
382 if args.grpc {
383 config.transports.grpc = Some(GrpcConfig {
384 host: args.grpc_host,
385 port: args.grpc_port,
386 });
387 }
388
389 if args.native {
390 config.transports.native = Some(NativeConfig {
391 host: args.native_host,
392 port: args.native_port,
393 });
394 }
395
396 config.validate()?;
398
399 Ok(config)
400 }
401
402 pub fn has_any_transport(&self) -> bool {
406 self.transports.http.is_some()
407 || self.transports.grpc.is_some()
408 || self.transports.native.is_some()
409 }
410
411 fn validate(&self) -> Result<()> {
420 if !self.has_any_transport() {
421 return Err(anyhow!(
422 "At least one transport must be specified.\n\n\
423 Available transports:\n \
424 --http Enable HTTP transport\n \
425 --grpc Enable gRPC transport\n \
426 --native Enable Native transport\n\n\
427 Example:\n \
428 throttlecrab-server --http --http-port 7070\n \
429 throttlecrab-server --grpc --native\n\n\
430 For more information, try '--help'"
431 ));
432 }
433
434 Ok(())
438 }
439
440 fn print_env_vars() {
446 println!("ThrottleCrab Environment Variables");
447 println!("==================================");
448 println!();
449 println!("All environment variables use the THROTTLECRAB_ prefix.");
450 println!("CLI arguments take precedence over environment variables.");
451 println!();
452
453 println!("Transport Configuration:");
454 println!(" THROTTLECRAB_HTTP=true|false Enable HTTP transport");
455 println!(" THROTTLECRAB_HTTP_HOST=<host> HTTP host [default: 127.0.0.1]");
456 println!(" THROTTLECRAB_HTTP_PORT=<port> HTTP port [default: 8080]");
457 println!();
458 println!(" THROTTLECRAB_GRPC=true|false Enable gRPC transport");
459 println!(" THROTTLECRAB_GRPC_HOST=<host> gRPC host [default: 127.0.0.1]");
460 println!(" THROTTLECRAB_GRPC_PORT=<port> gRPC port [default: 8070]");
461 println!();
462 println!(" THROTTLECRAB_NATIVE=true|false Enable Native transport");
463 println!(" THROTTLECRAB_NATIVE_HOST=<host> Native host [default: 127.0.0.1]");
464 println!(" THROTTLECRAB_NATIVE_PORT=<port> Native port [default: 8072]");
465 println!();
466
467 println!("Store Configuration:");
468 println!(
469 " THROTTLECRAB_STORE=<type> Store type: periodic, probabilistic, adaptive [default: periodic]"
470 );
471 println!(
472 " THROTTLECRAB_STORE_CAPACITY=<size> Initial store capacity [default: 100000]"
473 );
474 println!();
475 println!(" For periodic store:");
476 println!(
477 " THROTTLECRAB_STORE_CLEANUP_INTERVAL=<secs> Cleanup interval in seconds [default: 300]"
478 );
479 println!();
480 println!(" For probabilistic store:");
481 println!(
482 " THROTTLECRAB_STORE_CLEANUP_PROBABILITY=<n> Cleanup probability (1 in N) [default: 10000]"
483 );
484 println!();
485 println!(" For adaptive store:");
486 println!(
487 " THROTTLECRAB_STORE_MIN_INTERVAL=<secs> Minimum cleanup interval [default: 5]"
488 );
489 println!(
490 " THROTTLECRAB_STORE_MAX_INTERVAL=<secs> Maximum cleanup interval [default: 300]"
491 );
492 println!(
493 " THROTTLECRAB_STORE_MAX_OPERATIONS=<n> Max operations before cleanup [default: 1000000]"
494 );
495 println!();
496
497 println!("General Configuration:");
498 println!(" THROTTLECRAB_BUFFER_SIZE=<size> Channel buffer size [default: 100000]");
499 println!(
500 " THROTTLECRAB_LOG_LEVEL=<level> Log level: error, warn, info, debug, trace [default: info]"
501 );
502 println!();
503
504 println!("Examples:");
505 println!(" # Enable HTTP transport on port 8080");
506 println!(" export THROTTLECRAB_HTTP=true");
507 println!(" export THROTTLECRAB_HTTP_PORT=8080");
508 println!();
509 println!(" # Use adaptive store with custom settings");
510 println!(" export THROTTLECRAB_STORE=adaptive");
511 println!(" export THROTTLECRAB_STORE_MIN_INTERVAL=10");
512 println!(" export THROTTLECRAB_STORE_MAX_INTERVAL=600");
513 println!();
514 println!(" # Run server (CLI args override env vars)");
515 println!(" throttlecrab-server --http-port 9090 # Will use port 9090, not 8080");
516 }
517}
518
519#[cfg(test)]
520mod tests {
521 use super::*;
522 use std::str::FromStr;
523
524 #[test]
525 fn test_store_type_from_str() {
526 assert_eq!(
527 StoreType::from_str("periodic").unwrap(),
528 StoreType::Periodic
529 );
530 assert_eq!(
531 StoreType::from_str("PERIODIC").unwrap(),
532 StoreType::Periodic
533 );
534 assert_eq!(
535 StoreType::from_str("probabilistic").unwrap(),
536 StoreType::Probabilistic
537 );
538 assert_eq!(
539 StoreType::from_str("adaptive").unwrap(),
540 StoreType::Adaptive
541 );
542 assert!(StoreType::from_str("invalid").is_err());
543 }
544
545 #[test]
546 fn test_config_validation_no_transport() {
547 let config = Config {
548 transports: TransportConfig {
549 http: None,
550 grpc: None,
551 native: None,
552 },
553 store: StoreConfig {
554 store_type: StoreType::Periodic,
555 capacity: 100_000,
556 cleanup_interval: 300,
557 cleanup_probability: 10_000,
558 min_interval: 5,
559 max_interval: 300,
560 max_operations: 1_000_000,
561 },
562 buffer_size: 100_000,
563 log_level: "info".to_string(),
564 };
565
566 assert!(config.validate().is_err());
567 assert!(!config.has_any_transport());
568 }
569
570 #[test]
571 fn test_config_validation_with_transport() {
572 let config = Config {
573 transports: TransportConfig {
574 http: Some(HttpConfig {
575 host: "127.0.0.1".to_string(),
576 port: 8080,
577 }),
578 grpc: None,
579 native: None,
580 },
581 store: StoreConfig {
582 store_type: StoreType::Periodic,
583 capacity: 100_000,
584 cleanup_interval: 300,
585 cleanup_probability: 10_000,
586 min_interval: 5,
587 max_interval: 300,
588 max_operations: 1_000_000,
589 },
590 buffer_size: 100_000,
591 log_level: "info".to_string(),
592 };
593
594 assert!(config.validate().is_ok());
595 assert!(config.has_any_transport());
596 }
597
598 #[test]
599 fn test_config_multiple_transports() {
600 let config = Config {
601 transports: TransportConfig {
602 http: Some(HttpConfig {
603 host: "0.0.0.0".to_string(),
604 port: 8080,
605 }),
606 grpc: Some(GrpcConfig {
607 host: "0.0.0.0".to_string(),
608 port: 50051,
609 }),
610 native: None,
611 },
612 store: StoreConfig {
613 store_type: StoreType::Adaptive,
614 capacity: 200_000,
615 cleanup_interval: 300,
616 cleanup_probability: 10_000,
617 min_interval: 10,
618 max_interval: 600,
619 max_operations: 2_000_000,
620 },
621 buffer_size: 50_000,
622 log_level: "debug".to_string(),
623 };
624
625 assert!(config.validate().is_ok());
626 assert!(config.has_any_transport());
627 }
628}