ME Pool improvements

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

src/config/defaults.rs

@@ -277,6 +277,18 @@ pub(crate) fn default_me_reinit_every_secs() -> u64 {
     15 * 60
 }
 
+pub(crate) fn default_me_reinit_singleflight() -> bool {
+    true
+}
+
+pub(crate) fn default_me_reinit_trigger_channel() -> usize {
+    64
+}
+
+pub(crate) fn default_me_reinit_coalesce_window_ms() -> u64 {
+    200
+}
+
 pub(crate) fn default_me_hardswap_warmup_delay_min_ms() -> u64 {
     1000
 }
@@ -301,6 +313,18 @@ pub(crate) fn default_me_config_apply_cooldown_secs() -> u64 {
     300
 }
 
+pub(crate) fn default_me_snapshot_require_http_2xx() -> bool {
+    true
+}
+
+pub(crate) fn default_me_snapshot_reject_empty_map() -> bool {
+    true
+}
+
+pub(crate) fn default_me_snapshot_min_proxy_for_lines() -> u32 {
+    1
+}
+
 pub(crate) fn default_proxy_secret_stable_snapshots() -> u8 {
     2
 }
@@ -309,6 +333,10 @@ pub(crate) fn default_proxy_secret_rotate_runtime() -> bool {
     true
 }
 
+pub(crate) fn default_me_secret_atomic_snapshot() -> bool {
+    true
+}
+
 pub(crate) fn default_proxy_secret_len_max() -> usize {
     256
 }
@@ -321,10 +349,18 @@ pub(crate) fn default_me_pool_drain_ttl_secs() -> u64 {
     90
 }
 
+pub(crate) fn default_me_bind_stale_ttl_secs() -> u64 {
+    default_me_pool_drain_ttl_secs()
+}
+
 pub(crate) fn default_me_pool_min_fresh_ratio() -> f32 {
     0.8
 }
 
+pub(crate) fn default_me_deterministic_writer_sort() -> bool {
+    true
+}
+
 pub(crate) fn default_hardswap() -> bool {
     true
 }

src/config/load.rs

@@ -305,12 +305,24 @@ impl ProxyConfig {
             ));
         }
 
+        if config.general.me_snapshot_min_proxy_for_lines == 0 {
+            return Err(ProxyError::Config(
+                "general.me_snapshot_min_proxy_for_lines must be > 0".to_string(),
+            ));
+        }
+
         if config.general.proxy_secret_stable_snapshots == 0 {
             return Err(ProxyError::Config(
                 "general.proxy_secret_stable_snapshots must be > 0".to_string(),
             ));
         }
 
+        if config.general.me_reinit_trigger_channel == 0 {
+            return Err(ProxyError::Config(
+                "general.me_reinit_trigger_channel must be > 0".to_string(),
+            ));
+        }
+
         if !(32..=4096).contains(&config.general.proxy_secret_len_max) {
             return Err(ProxyError::Config(
                 "general.proxy_secret_len_max must be within [32, 4096]".to_string(),
@@ -535,9 +547,10 @@ impl ProxyConfig {
         for (user, tag) in &self.access.user_ad_tags {
             let zeros = "00000000000000000000000000000000";
             if !is_valid_ad_tag(tag) {
-                return Err(ProxyError::Config(
-                    "general.ad_tag must be exactly 32 hex characters".to_string(),
-                ));
+                return Err(ProxyError::Config(format!(
+                    "access.user_ad_tags['{}'] must be exactly 32 hex characters",
+                    user
+                )));
             }
             if tag == zeros {
                 warn!(user = %user, "user ad_tag is all zeros; register a valid proxy tag via @MTProxybot to enable sponsored channel");
@@ -1100,6 +1113,27 @@ mod tests {
         let _ = std::fs::remove_file(path);
     }
 
+    #[test]
+    fn invalid_user_ad_tag_reports_access_user_ad_tags_key() {
+        let toml = r#"
+            [censorship]
+            tls_domain = "example.com"
+
+            [access.users]
+            alice = "00000000000000000000000000000000"
+
+            [access.user_ad_tags]
+            alice = "not_hex"
+        "#;
+        let dir = std::env::temp_dir();
+        let path = dir.join("telemt_invalid_user_ad_tag_message_test.toml");
+        std::fs::write(&path, toml).unwrap();
+        let cfg = ProxyConfig::load(&path).unwrap();
+        let err = cfg.validate().unwrap_err().to_string();
+        assert!(err.contains("access.user_ad_tags['alice'] must be exactly 32 hex characters"));
+        let _ = std::fs::remove_file(path);
+    }
+
     #[test]
     fn invalid_dns_override_is_rejected() {
         let toml = r#"

src/config/types.rs

@@ -130,6 +130,34 @@ impl MeSocksKdfPolicy {
     }
 }
 
+/// Stale ME writer bind policy during drain window.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
+#[serde(rename_all = "lowercase")]
+pub enum MeBindStaleMode {
+    Never,
+    #[default]
+    Ttl,
+    Always,
+}
+
+impl MeBindStaleMode {
+    pub fn as_u8(self) -> u8 {
+        match self {
+            MeBindStaleMode::Never => 0,
+            MeBindStaleMode::Ttl => 1,
+            MeBindStaleMode::Always => 2,
+        }
+    }
+
+    pub fn from_u8(raw: u8) -> Self {
+        match raw {
+            0 => MeBindStaleMode::Never,
+            2 => MeBindStaleMode::Always,
+            _ => MeBindStaleMode::Ttl,
+        }
+    }
+}
+
 /// Telemetry controls for hot-path counters and ME diagnostics.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct TelemetryConfig {
@@ -454,6 +482,18 @@ pub struct GeneralConfig {
     #[serde(default = "default_me_config_apply_cooldown_secs")]
     pub me_config_apply_cooldown_secs: u64,
 
+    /// Ensure getProxyConfig snapshots are applied only for 2xx HTTP responses.
+    #[serde(default = "default_me_snapshot_require_http_2xx")]
+    pub me_snapshot_require_http_2xx: bool,
+
+    /// Reject empty getProxyConfig snapshots instead of marking them applied.
+    #[serde(default = "default_me_snapshot_reject_empty_map")]
+    pub me_snapshot_reject_empty_map: bool,
+
+    /// Minimum parsed `proxy_for` rows required to accept a snapshot.
+    #[serde(default = "default_me_snapshot_min_proxy_for_lines")]
+    pub me_snapshot_min_proxy_for_lines: u32,
+
     /// Number of identical getProxySecret snapshots required before runtime secret rotation.
     #[serde(default = "default_proxy_secret_stable_snapshots")]
     pub proxy_secret_stable_snapshots: u8,
@@ -462,6 +502,10 @@ pub struct GeneralConfig {
     #[serde(default = "default_proxy_secret_rotate_runtime")]
     pub proxy_secret_rotate_runtime: bool,
 
+    /// Keep key-selector and secret bytes from one snapshot during ME handshake.
+    #[serde(default = "default_me_secret_atomic_snapshot")]
+    pub me_secret_atomic_snapshot: bool,
+
     /// Maximum allowed proxy-secret length in bytes for startup and runtime refresh.
     #[serde(default = "default_proxy_secret_len_max")]
     pub proxy_secret_len_max: usize,
@@ -471,6 +515,14 @@ pub struct GeneralConfig {
     #[serde(default = "default_me_pool_drain_ttl_secs")]
     pub me_pool_drain_ttl_secs: u64,
 
+    /// Policy for new binds on stale draining writers.
+    #[serde(default)]
+    pub me_bind_stale_mode: MeBindStaleMode,
+
+    /// TTL for stale bind allowance when `me_bind_stale_mode = \"ttl\"`.
+    #[serde(default = "default_me_bind_stale_ttl_secs")]
+    pub me_bind_stale_ttl_secs: u64,
+
     /// Minimum desired-DC coverage ratio required before draining stale writers.
     /// Range: 0.0..=1.0.
     #[serde(default = "default_me_pool_min_fresh_ratio")]
@@ -491,6 +543,22 @@ pub struct GeneralConfig {
     #[serde(default = "default_proxy_config_reload_secs")]
     pub proxy_config_auto_reload_secs: u64,
 
+    /// Serialize ME reinit cycles across all trigger sources.
+    #[serde(default = "default_me_reinit_singleflight")]
+    pub me_reinit_singleflight: bool,
+
+    /// Trigger queue capacity for reinit scheduler.
+    #[serde(default = "default_me_reinit_trigger_channel")]
+    pub me_reinit_trigger_channel: usize,
+
+    /// Trigger coalescing window before starting a reinit cycle.
+    #[serde(default = "default_me_reinit_coalesce_window_ms")]
+    pub me_reinit_coalesce_window_ms: u64,
+
+    /// Deterministic candidate sort for ME writer binding path.
+    #[serde(default = "default_me_deterministic_writer_sort")]
+    pub me_deterministic_writer_sort: bool,
+
     /// Enable NTP drift check at startup.
     #[serde(default = "default_ntp_check")]
     pub ntp_check: bool,
@@ -565,14 +633,24 @@ impl Default for GeneralConfig {
             me_hardswap_warmup_pass_backoff_base_ms: default_me_hardswap_warmup_pass_backoff_base_ms(),
             me_config_stable_snapshots: default_me_config_stable_snapshots(),
             me_config_apply_cooldown_secs: default_me_config_apply_cooldown_secs(),
+            me_snapshot_require_http_2xx: default_me_snapshot_require_http_2xx(),
+            me_snapshot_reject_empty_map: default_me_snapshot_reject_empty_map(),
+            me_snapshot_min_proxy_for_lines: default_me_snapshot_min_proxy_for_lines(),
             proxy_secret_stable_snapshots: default_proxy_secret_stable_snapshots(),
             proxy_secret_rotate_runtime: default_proxy_secret_rotate_runtime(),
+            me_secret_atomic_snapshot: default_me_secret_atomic_snapshot(),
             proxy_secret_len_max: default_proxy_secret_len_max(),
             me_pool_drain_ttl_secs: default_me_pool_drain_ttl_secs(),
+            me_bind_stale_mode: MeBindStaleMode::default(),
+            me_bind_stale_ttl_secs: default_me_bind_stale_ttl_secs(),
             me_pool_min_fresh_ratio: default_me_pool_min_fresh_ratio(),
             me_reinit_drain_timeout_secs: default_me_reinit_drain_timeout_secs(),
             proxy_secret_auto_reload_secs: default_proxy_secret_reload_secs(),
             proxy_config_auto_reload_secs: default_proxy_config_reload_secs(),
+            me_reinit_singleflight: default_me_reinit_singleflight(),
+            me_reinit_trigger_channel: default_me_reinit_trigger_channel(),
+            me_reinit_coalesce_window_ms: default_me_reinit_coalesce_window_ms(),
+            me_deterministic_writer_sort: default_me_deterministic_writer_sort(),
             ntp_check: default_ntp_check(),
             ntp_servers: default_ntp_servers(),
             auto_degradation_enabled: default_true(),