User Disabler in API by #814 + Consistent Listeners in API by #800

src/proxy/shared_state.rs

@@ -1,5 +1,5 @@
-use std::collections::HashSet;
 use std::collections::hash_map::RandomState;
+use std::collections::{HashMap, HashSet};
 use std::net::{IpAddr, SocketAddr};
 use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
@@ -7,6 +7,7 @@ use std::time::Instant;
 
 use dashmap::DashMap;
 use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
 
 use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState};
 use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
@@ -67,10 +68,35 @@ pub(crate) struct ProxySharedState {
     pub(crate) handshake: HandshakeSharedState,
     pub(crate) middle_relay: MiddleRelaySharedState,
     pub(crate) traffic_limiter: Arc<TrafficLimiter>,
+    disabled_users: DashMap<String, ()>,
+    active_user_sessions: DashMap<(String, u64), CancellationToken>,
     pub(crate) conntrack_pressure_active: AtomicBool,
     pub(crate) conntrack_close_tx: Mutex<Option<mpsc::Sender<ConntrackCloseEvent>>>,
 }
 
+#[must_use = "registered user sessions must be kept alive until relay completion"]
+pub(crate) struct UserSessionRegistration {
+    token: CancellationToken,
+    _guard: UserSessionGuard,
+}
+
+impl UserSessionRegistration {
+    pub(crate) fn token(&self) -> CancellationToken {
+        self.token.clone()
+    }
+}
+
+struct UserSessionGuard {
+    shared: Arc<ProxySharedState>,
+    key: (String, u64),
+}
+
+impl Drop for UserSessionGuard {
+    fn drop(&mut self) {
+        self.shared.active_user_sessions.remove(&self.key);
+    }
+}
+
 impl ProxySharedState {
     pub(crate) fn new() -> Arc<Self> {
         Arc::new(Self {
@@ -101,11 +127,82 @@ impl ProxySharedState {
                 relay_idle_mark_seq: AtomicU64::new(0),
             },
             traffic_limiter: TrafficLimiter::new(),
+            disabled_users: DashMap::new(),
+            active_user_sessions: DashMap::new(),
             conntrack_pressure_active: AtomicBool::new(false),
             conntrack_close_tx: Mutex::new(None),
         })
     }
 
+    pub(crate) fn is_user_enabled(&self, user: &str) -> bool {
+        !self.disabled_users.contains_key(user)
+    }
+
+    pub(crate) fn set_user_enabled(&self, user: &str, enabled: bool) -> bool {
+        if enabled {
+            self.disabled_users.remove(user);
+            false
+        } else {
+            self.disabled_users.insert(user.to_string(), ()).is_none()
+        }
+    }
+
+    pub(crate) fn apply_user_enabled_config(
+        &self,
+        user_enabled: &HashMap<String, bool>,
+    ) -> Vec<String> {
+        let desired_disabled = user_enabled
+            .iter()
+            .filter_map(|(user, enabled)| (!*enabled).then_some(user.clone()))
+            .collect::<HashSet<_>>();
+        let current_disabled = self
+            .disabled_users
+            .iter()
+            .map(|entry| entry.key().clone())
+            .collect::<HashSet<_>>();
+
+        for user in current_disabled.difference(&desired_disabled) {
+            self.disabled_users.remove(user);
+        }
+        let newly_disabled = desired_disabled
+            .difference(&current_disabled)
+            .cloned()
+            .collect::<Vec<_>>();
+        for user in desired_disabled {
+            self.disabled_users.insert(user, ());
+        }
+        newly_disabled
+    }
+
+    pub(crate) fn register_user_session(
+        self: &Arc<Self>,
+        user: &str,
+        session_id: u64,
+    ) -> UserSessionRegistration {
+        let token = CancellationToken::new();
+        let key = (user.to_string(), session_id);
+        self.active_user_sessions.insert(key.clone(), token.clone());
+        UserSessionRegistration {
+            token,
+            _guard: UserSessionGuard {
+                shared: Arc::clone(self),
+                key,
+            },
+        }
+    }
+
+    pub(crate) fn cancel_user_sessions(&self, user: &str) -> usize {
+        let tokens = self
+            .active_user_sessions
+            .iter()
+            .filter_map(|entry| (entry.key().0 == user).then(|| entry.value().clone()))
+            .collect::<Vec<_>>();
+        for token in &tokens {
+            token.cancel();
+        }
+        tokens.len()
+    }
+
     pub(crate) fn set_conntrack_close_sender(&self, tx: mpsc::Sender<ConntrackCloseEvent>) {
         match self.conntrack_close_tx.lock() {
             Ok(mut guard) => {
@@ -166,3 +263,48 @@ impl ProxySharedState {
         self.conntrack_pressure_active.load(Ordering::Relaxed)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn user_enabled_config_sync_tracks_disabled_overrides() {
+        let shared = ProxySharedState::new();
+        assert!(shared.is_user_enabled("alice"));
+
+        let mut user_enabled = HashMap::new();
+        user_enabled.insert("alice".to_string(), false);
+        user_enabled.insert("bob".to_string(), true);
+
+        let mut newly_disabled = shared.apply_user_enabled_config(&user_enabled);
+        newly_disabled.sort();
+        assert_eq!(newly_disabled, vec!["alice".to_string()]);
+        assert!(!shared.is_user_enabled("alice"));
+        assert!(shared.is_user_enabled("bob"));
+
+        assert!(shared.apply_user_enabled_config(&user_enabled).is_empty());
+
+        user_enabled.clear();
+        assert!(shared.apply_user_enabled_config(&user_enabled).is_empty());
+        assert!(shared.is_user_enabled("alice"));
+    }
+
+    #[test]
+    fn cancel_user_sessions_cancels_only_registered_matching_user() {
+        let shared = ProxySharedState::new();
+        let alice_1 = shared.register_user_session("alice", 1);
+        let alice_2 = shared.register_user_session("alice", 2);
+        let bob = shared.register_user_session("bob", 1);
+        let alice_1_token = alice_1.token();
+        let alice_2_token = alice_2.token();
+        let bob_token = bob.token();
+
+        drop(alice_1);
+
+        assert_eq!(shared.cancel_user_sessions("alice"), 1);
+        assert!(!alice_1_token.is_cancelled());
+        assert!(alice_2_token.is_cancelled());
+        assert!(!bob_token.is_cancelled());
+    }
+}