Drafting Traffic Control

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/config/hot_reload.rs

@@ -121,6 +121,8 @@ pub struct HotFields {
     pub user_max_tcp_conns_global_each: usize,
     pub user_expirations: std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>,
     pub user_data_quota: std::collections::HashMap<String, u64>,
+    pub user_rate_limits: std::collections::HashMap<String, crate::config::RateLimitBps>,
+    pub cidr_rate_limits: std::collections::HashMap<ipnetwork::IpNetwork, crate::config::RateLimitBps>,
     pub user_max_unique_ips: std::collections::HashMap<String, usize>,
     pub user_max_unique_ips_global_each: usize,
     pub user_max_unique_ips_mode: crate::config::UserMaxUniqueIpsMode,
@@ -245,6 +247,8 @@ impl HotFields {
             user_max_tcp_conns_global_each: cfg.access.user_max_tcp_conns_global_each,
             user_expirations: cfg.access.user_expirations.clone(),
             user_data_quota: cfg.access.user_data_quota.clone(),
+            user_rate_limits: cfg.access.user_rate_limits.clone(),
+            cidr_rate_limits: cfg.access.cidr_rate_limits.clone(),
             user_max_unique_ips: cfg.access.user_max_unique_ips.clone(),
             user_max_unique_ips_global_each: cfg.access.user_max_unique_ips_global_each,
             user_max_unique_ips_mode: cfg.access.user_max_unique_ips_mode,
@@ -545,6 +549,8 @@ fn overlay_hot_fields(old: &ProxyConfig, new: &ProxyConfig) -> ProxyConfig {
     cfg.access.user_max_tcp_conns_global_each = new.access.user_max_tcp_conns_global_each;
     cfg.access.user_expirations = new.access.user_expirations.clone();
     cfg.access.user_data_quota = new.access.user_data_quota.clone();
+    cfg.access.user_rate_limits = new.access.user_rate_limits.clone();
+    cfg.access.cidr_rate_limits = new.access.cidr_rate_limits.clone();
     cfg.access.user_max_unique_ips = new.access.user_max_unique_ips.clone();
     cfg.access.user_max_unique_ips_global_each = new.access.user_max_unique_ips_global_each;
     cfg.access.user_max_unique_ips_mode = new.access.user_max_unique_ips_mode;
@@ -1183,6 +1189,18 @@ fn log_changes(
             new_hot.user_data_quota.len()
         );
     }
+    if old_hot.user_rate_limits != new_hot.user_rate_limits {
+        info!(
+            "config reload: user_rate_limits updated ({} entries)",
+            new_hot.user_rate_limits.len()
+        );
+    }
+    if old_hot.cidr_rate_limits != new_hot.cidr_rate_limits {
+        info!(
+            "config reload: cidr_rate_limits updated ({} entries)",
+            new_hot.cidr_rate_limits.len()
+        );
+    }
     if old_hot.user_max_unique_ips != new_hot.user_max_unique_ips {
         info!(
             "config reload: user_max_unique_ips updated ({} entries)",

src/config/load.rs

@@ -861,6 +861,22 @@ impl ProxyConfig {
             ));
         }
 
+        for (user, limit) in &config.access.user_rate_limits {
+            if limit.up_bps == 0 && limit.down_bps == 0 {
+                return Err(ProxyError::Config(format!(
+                    "access.user_rate_limits.{user} must set at least one non-zero direction"
+                )));
+            }
+        }
+
+        for (cidr, limit) in &config.access.cidr_rate_limits {
+            if limit.up_bps == 0 && limit.down_bps == 0 {
+                return Err(ProxyError::Config(format!(
+                    "access.cidr_rate_limits.{cidr} must set at least one non-zero direction"
+                )));
+            }
+        }
+
         if config.general.me_reinit_every_secs == 0 {
             return Err(ProxyError::Config(
                 "general.me_reinit_every_secs must be > 0".to_string(),

src/config/types.rs

@@ -1826,6 +1826,21 @@ pub struct AccessConfig {
     #[serde(default)]
     pub user_data_quota: HashMap<String, u64>,
 
+    /// Per-user transport rate limits in bits-per-second.
+    ///
+    /// Each entry supports independent upload (`up_bps`) and download
+    /// (`down_bps`) ceilings. A value of `0` in one direction means
+    /// "unlimited" for that direction.
+    #[serde(default)]
+    pub user_rate_limits: HashMap<String, RateLimitBps>,
+
+    /// Per-CIDR aggregate transport rate limits in bits-per-second.
+    ///
+    /// Matching uses longest-prefix-wins semantics. A value of `0` in one
+    /// direction means "unlimited" for that direction.
+    #[serde(default)]
+    pub cidr_rate_limits: HashMap<IpNetwork, RateLimitBps>,
+
     #[serde(default)]
     pub user_max_unique_ips: HashMap<String, usize>,
 
@@ -1859,6 +1874,8 @@ impl Default for AccessConfig {
             user_max_tcp_conns_global_each: default_user_max_tcp_conns_global_each(),
             user_expirations: HashMap::new(),
             user_data_quota: HashMap::new(),
+            user_rate_limits: HashMap::new(),
+            cidr_rate_limits: HashMap::new(),
             user_max_unique_ips: HashMap::new(),
             user_max_unique_ips_global_each: default_user_max_unique_ips_global_each(),
             user_max_unique_ips_mode: UserMaxUniqueIpsMode::default(),
@@ -1870,6 +1887,14 @@ impl Default for AccessConfig {
     }
 }
 
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RateLimitBps {
+    #[serde(default)]
+    pub up_bps: u64,
+    #[serde(default)]
+    pub down_bps: u64,
+}
+
 // ============= Aux Structures =============
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]