WEB Manager Session + Generation Fence tests

src/proxy/traffic_limiter/tests.rs

@@ -281,3 +281,47 @@ fn active_lease_observes_policy_removal() {
 
     assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 1);
 }
+
+#[test]
+fn active_lease_observes_policy_addition() {
+    let limiter = TrafficLimiter::new();
+    let lease = limiter
+        .acquire_lease("alice", "203.0.113.7".parse().unwrap())
+        .unwrap();
+    assert_eq!(lease.try_consume(RateDirection::Up, 2).granted, 2);
+
+    let mut user_limits = HashMap::new();
+    user_limits.insert("alice".to_string(), rate(1, 0));
+    limiter.apply_policy(user_limits, HashMap::new());
+
+    assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 1);
+    assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 0);
+}
+
+#[test]
+fn reservation_refund_stays_with_retired_binding() {
+    let limiter = TrafficLimiter::new();
+    let mut user_limits = HashMap::new();
+    user_limits.insert("alice".to_string(), rate(400_000, 400_000));
+    limiter.apply_policy(user_limits, HashMap::new());
+    let lease = limiter
+        .acquire_lease("alice", "203.0.113.7".parse().unwrap())
+        .unwrap();
+
+    let reservation = lease.try_reserve(RateDirection::Down, 800);
+    let old_bucket = Arc::clone(
+        reservation
+            ._binding
+            .user_bucket
+            .as_ref()
+            .expect("the original policy must bind a user bucket"),
+    );
+    let epoch = reservation.user.as_ref().unwrap().epoch;
+    limiter.apply_policy(HashMap::new(), HashMap::new());
+    assert_eq!(lease.try_consume(RateDirection::Down, 1).granted, 1);
+    assert!(lease.binding.load().user_bucket.is_none());
+
+    drop(reservation);
+
+    assert_eq!(old_bucket.down.used_at(epoch), Some(0));
+}

src/proxy/user_admission.rs

@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicU8, Ordering};
+use std::sync::atomic::AtomicU8;
 
-use parking_lot::{Mutex, MutexGuard};
+use parking_lot::Mutex;
 use tokio_util::sync::CancellationToken;
 
 use crate::crypto::sha256;
@@ -12,6 +12,10 @@ const REGISTRATION_PENDING: u8 = 0;
 const REGISTRATION_ACTIVE: u8 = 1;
 const REGISTRATION_DROPPED: u8 = 2;
 
+// Authenticated-owner publication and RAII deregistration.
+mod registration;
+pub(crate) use registration::{UserAdmissionPublication, UserSessionRegistration};
+
 /// Stable secret identity used to fence authentication across runtime generations.
 pub(crate) type UserCredentialId = [u8; 16];
 
@@ -501,96 +505,6 @@ impl UserAdmissionAuthority {
     }
 }
 
-/// Authority lock retained until the caller publishes its owned object.
-pub(crate) struct UserAdmissionPublication<'a> {
-    state: MutexGuard<'a, UserAdmissionState>,
-    authority: Arc<UserAdmissionAuthority>,
-    user: String,
-    registration_id: u64,
-    incarnation: UserIncarnation,
-    token: CancellationToken,
-    active: Arc<AtomicU8>,
-    registration_taken: bool,
-}
-
-impl UserAdmissionPublication<'_> {
-    /// Moves the registered owner out while retaining the authority lock.
-    pub(crate) fn take_registration(&mut self) -> Option<UserSessionRegistration> {
-        if self.registration_taken {
-            return None;
-        }
-        self.registration_taken = true;
-        Some(UserSessionRegistration {
-            authority: Arc::clone(&self.authority),
-            user: self.user.clone(),
-            registration_id: self.registration_id,
-            incarnation: self.incarnation,
-            token: self.token.clone(),
-            active: Arc::clone(&self.active),
-        })
-    }
-
-    /// Commits the owner record after the caller publishes its lifecycle object.
-    pub(crate) fn commit(mut self) {
-        if !self.registration_taken {
-            return;
-        }
-        if self.active.compare_exchange(
-            REGISTRATION_PENDING, REGISTRATION_ACTIVE, Ordering::AcqRel, Ordering::Acquire,
-        ).is_err() {
-            return;
-        }
-        self.state
-            .owners_by_user
-            .entry(self.user.clone())
-            .or_default()
-            .insert(
-                self.registration_id,
-                RegisteredOwner {
-                    token: self.token.clone(),
-                    incarnation: self.incarnation,
-                },
-            );
-    }
-}
-
-/// RAII ownership registered against one user incarnation.
-#[must_use = "registered user ownership must be retained until lifecycle completion"]
-pub(crate) struct UserSessionRegistration {
-    authority: Arc<UserAdmissionAuthority>,
-    user: String,
-    registration_id: u64,
-    incarnation: UserIncarnation,
-    token: CancellationToken,
-    active: Arc<AtomicU8>,
-}
-
-impl UserSessionRegistration {
-    /// Returns the cancellation signal for revocation or credential replacement.
-    pub(crate) fn token(&self) -> CancellationToken {
-        self.token.clone()
-    }
-
-    /// Returns the immutable user incarnation owned by this registration.
-    pub(crate) fn incarnation(&self) -> UserIncarnation {
-        self.incarnation
-    }
-
-    /// Returns whether revocation has cancelled this ownership.
-    pub(crate) fn is_cancelled(&self) -> bool {
-        self.token.is_cancelled()
-    }
-}
-
-impl Drop for UserSessionRegistration {
-    fn drop(&mut self) {
-        if self.active.swap(REGISTRATION_DROPPED, Ordering::AcqRel) == REGISTRATION_ACTIVE {
-            self.authority
-                .unregister(&self.user, self.registration_id, self.incarnation);
-        }
-    }
-}
-
 /// Derives the stable credential identity from one decoded MTProxy secret.
 pub(crate) fn credential_id(secret: &[u8; 16]) -> UserCredentialId {
     let digest = sha256(secret);

src/proxy/user_admission/registration.rs

@@ -0,0 +1,112 @@
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU8, Ordering};
+
+use parking_lot::MutexGuard;
+use tokio_util::sync::CancellationToken;
+
+use super::*;
+
+/// Authority lock retained until the caller publishes its owned object.
+pub(crate) struct UserAdmissionPublication<'a> {
+    /// Locked authority state that linearizes publication with policy mutation.
+    pub(super) state: MutexGuard<'a, UserAdmissionState>,
+    /// Process authority used by the eventual registration guard.
+    pub(super) authority: Arc<UserAdmissionAuthority>,
+    /// Username owning the published lifecycle object.
+    pub(super) user: String,
+    /// Unique registration identity within the process authority.
+    pub(super) registration_id: u64,
+    /// User incarnation authenticated by this publication.
+    pub(super) incarnation: UserIncarnation,
+    /// Revocation signal shared with the lifecycle owner.
+    pub(super) token: CancellationToken,
+    /// Atomic publication state shared with the registration guard.
+    pub(super) active: Arc<AtomicU8>,
+    /// Whether ownership has already moved into a registration guard.
+    pub(super) registration_taken: bool,
+}
+
+impl UserAdmissionPublication<'_> {
+    /// Moves the registered owner out while retaining the authority lock.
+    pub(crate) fn take_registration(&mut self) -> Option<UserSessionRegistration> {
+        if self.registration_taken {
+            return None;
+        }
+        self.registration_taken = true;
+        Some(UserSessionRegistration {
+            authority: Arc::clone(&self.authority),
+            user: self.user.clone(),
+            registration_id: self.registration_id,
+            incarnation: self.incarnation,
+            token: self.token.clone(),
+            active: Arc::clone(&self.active),
+        })
+    }
+
+    /// Commits the owner record after the caller publishes its lifecycle object.
+    pub(crate) fn commit(mut self) {
+        if !self.registration_taken {
+            return;
+        }
+        if self
+            .active
+            .compare_exchange(
+                REGISTRATION_PENDING,
+                REGISTRATION_ACTIVE,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            )
+            .is_err()
+        {
+            return;
+        }
+        self.state
+            .owners_by_user
+            .entry(self.user.clone())
+            .or_default()
+            .insert(
+                self.registration_id,
+                RegisteredOwner {
+                    token: self.token.clone(),
+                    incarnation: self.incarnation,
+                },
+            );
+    }
+}
+
+/// RAII ownership registered against one user incarnation.
+#[must_use = "registered user ownership must be retained until lifecycle completion"]
+pub(crate) struct UserSessionRegistration {
+    authority: Arc<UserAdmissionAuthority>,
+    user: String,
+    registration_id: u64,
+    incarnation: UserIncarnation,
+    token: CancellationToken,
+    active: Arc<AtomicU8>,
+}
+
+impl UserSessionRegistration {
+    /// Returns the cancellation signal for revocation or credential replacement.
+    pub(crate) fn token(&self) -> CancellationToken {
+        self.token.clone()
+    }
+
+    /// Returns the immutable user incarnation owned by this registration.
+    pub(crate) fn incarnation(&self) -> UserIncarnation {
+        self.incarnation
+    }
+
+    /// Returns whether revocation has cancelled this ownership.
+    pub(crate) fn is_cancelled(&self) -> bool {
+        self.token.is_cancelled()
+    }
+}
+
+impl Drop for UserSessionRegistration {
+    fn drop(&mut self) {
+        if self.active.swap(REGISTRATION_DROPPED, Ordering::AcqRel) == REGISTRATION_ACTIVE {
+            self.authority
+                .unregister(&self.user, self.registration_id, self.incarnation);
+        }
+    }
+}

src/proxy/user_admission/tests.rs

@@ -60,6 +60,19 @@ fn stale_candidate_cannot_overwrite_newer_mutation() {
             .is_none()
     );
     assert!(!authority.is_user_enabled("alice"));
+
+    let disabled = HashMap::from([("alice".to_string(), false)]);
+    assert!(
+        authority
+            .apply_config_from_source(1, &users(secret), &disabled)
+            .is_none()
+    );
+    assert!(
+        authority
+            .apply_config_from_source(2, &users(secret), &disabled)
+            .is_some()
+    );
+    assert!(!authority.is_user_enabled("alice"));
 }
 
 #[test]