Proxy Shared User Drafts

This commit is contained in:
Alexey
2026-09-14 20:19:48 +03:00
parent b37f1ebdeb
commit 0ac236955a
59 changed files with 1485 additions and 401 deletions
+74 -9
View File
@@ -14,6 +14,7 @@ use crate::proxy::handshake::HandshakeSuccess;
use crate::proxy::middle_relay::{handle_via_middle_proxy, handle_via_middle_proxy_with_conntrack};
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::{ConntrackClosePolicy, ProxySharedState};
use crate::proxy::user_admission::UserIncarnation;
use crate::stats::Stats;
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
use crate::transport::UpstreamManager;
@@ -59,13 +60,21 @@ where
W: AsyncWrite + Unpin + Send + 'static,
{
let user = success.user.clone();
if !deps.shared.is_user_enabled(&user) {
let Some(credential_id) = deps.config.runtime_user_credential_id(&user) else {
warn!(user = %user, "Authenticated user is absent from the runtime credential snapshot");
return Err(ProxyError::UserDisabled { user });
};
let Some(user_incarnation) = deps
.shared
.authenticated_user_incarnation(&user, credential_id)
else {
warn!(user = %user, "Disabled user rejected");
return Err(ProxyError::UserDisabled { user });
}
};
let user_reservation = acquire_user_connection_reservation(
let user_reservation = acquire_user_connection_reservation_for_incarnation(
&user,
user_incarnation,
&deps.config,
Arc::clone(&deps.stats),
peer_addr,
@@ -79,11 +88,20 @@ where
let route_snapshot = deps.route_runtime.snapshot();
let session_id = deps.rng.u64();
let Some(user_session) = deps.shared.register_user_session(&user, session_id) else {
let Some(user_session) = deps
.shared
.register_authenticated_user_session(&user, credential_id)
else {
user_reservation.release_deferred();
warn!(user = %user, "Disabled user rejected during final admission");
return Err(ProxyError::UserDisabled { user });
};
if user_session.incarnation() != user_incarnation {
drop(user_session);
user_reservation.release_deferred();
warn!(user = %user, "User incarnation changed during admission");
return Err(ProxyError::UserDisabled { user });
}
let session_cancel = user_session.token();
let selected_me_pool = if deps.config.general.use_middle_proxy
&& matches!(route_snapshot.mode, RelayRouteMode::Middle)
@@ -216,6 +234,7 @@ pub(crate) struct UserConnectionReservation {
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
incarnation: UserIncarnation,
tracks_ip: bool,
active: bool,
}
@@ -228,12 +247,25 @@ impl UserConnectionReservation {
user: String,
ip: IpAddr,
tracks_ip: bool,
) -> Self {
Self::new_for_incarnation(stats, ip_tracker, user, ip, 0, tracks_ip)
}
/// Creates a reservation fenced to one authenticated user incarnation.
pub(crate) fn new_for_incarnation(
stats: Arc<Stats>,
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
incarnation: UserIncarnation,
tracks_ip: bool,
) -> Self {
Self {
stats,
ip_tracker,
user,
ip,
incarnation,
tracks_ip,
active: true,
}
@@ -246,7 +278,9 @@ impl UserConnectionReservation {
}
self.active = false;
if self.tracks_ip {
self.ip_tracker.remove_ip(&self.user, self.ip).await;
self.ip_tracker
.remove_ip_for_incarnation(&self.user, self.incarnation, self.ip)
.await;
}
self.stats.decrement_user_curr_connects(&self.user);
}
@@ -259,7 +293,11 @@ impl UserConnectionReservation {
self.active = false;
self.stats.decrement_user_curr_connects(&self.user);
if self.tracks_ip {
self.ip_tracker.enqueue_cleanup(self.user.clone(), self.ip);
self.ip_tracker.enqueue_cleanup_for_incarnation(
self.user.clone(),
self.incarnation,
self.ip,
);
}
}
}
@@ -273,7 +311,11 @@ impl Drop for UserConnectionReservation {
self.stats.increment_session_drop_fallback_total();
self.stats.decrement_user_curr_connects(&self.user);
if self.tracks_ip {
self.ip_tracker.enqueue_cleanup(self.user.clone(), self.ip);
self.ip_tracker.enqueue_cleanup_for_incarnation(
self.user.clone(),
self.incarnation,
self.ip,
);
}
}
}
@@ -285,6 +327,25 @@ pub(crate) async fn acquire_user_connection_reservation(
stats: Arc<Stats>,
peer_addr: SocketAddr,
ip_tracker: Arc<UserIpTracker>,
) -> Result<UserConnectionReservation> {
acquire_user_connection_reservation_for_incarnation(
user,
0,
config,
stats,
peer_addr,
ip_tracker,
)
.await
}
async fn acquire_user_connection_reservation_for_incarnation(
user: &str,
incarnation: UserIncarnation,
config: &ProxyConfig,
stats: Arc<Stats>,
peer_addr: SocketAddr,
ip_tracker: Arc<UserIpTracker>,
) -> Result<UserConnectionReservation> {
if let Some(expiration) = config.access.user_expirations.get(user)
&& chrono::Utc::now() > *expiration
@@ -316,7 +377,10 @@ pub(crate) async fn acquire_user_connection_reservation(
});
}
if let Err(reason) = ip_tracker.check_and_add(user, peer_addr.ip()).await {
if let Err(reason) = ip_tracker
.check_and_add_for_incarnation(user, incarnation, peer_addr.ip())
.await
{
stats.decrement_user_curr_connects(user);
warn!(
user = %user,
@@ -329,11 +393,12 @@ pub(crate) async fn acquire_user_connection_reservation(
});
}
Ok(UserConnectionReservation::new(
Ok(UserConnectionReservation::new_for_incarnation(
stats,
ip_tracker,
user.to_string(),
peer_addr.ip(),
incarnation,
true,
))
}
+1
View File
@@ -73,6 +73,7 @@ pub mod route_mode;
pub mod session_eviction;
pub mod shared_state;
pub mod traffic_limiter;
pub(crate) mod user_admission;
pub use client::ClientHandler;
#[allow(unused_imports)]
+126 -139
View File
@@ -6,14 +6,16 @@ use std::sync::{Arc, Mutex};
use std::time::Instant;
use dashmap::DashMap;
use parking_lot::Mutex as ParkingMutex;
use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
use tokio_util::sync::CancellationToken;
use crate::proxy::direct_buffer_budget::{DirectBufferBudget, fallback_direct_buffer_hard_limit};
use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState};
use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
use crate::proxy::traffic_limiter::TrafficLimiter;
use crate::proxy::user_admission::{
UserAdmissionAuthority, UserAdmissionPublication, UserCredentialId, UserIncarnation,
UserMutationResult, UserSessionRegistration,
};
const HANDSHAKE_RECENT_USER_RING_LEN: usize = 64;
const MASKING_FALLBACK_MAX_CONCURRENT: usize = 512;
@@ -76,57 +78,17 @@ pub(crate) struct MiddleRelaySharedState {
pub(crate) relay_idle_mark_seq: AtomicU64,
}
#[derive(Default)]
struct UserAdmissionState {
disabled_users: HashSet<String>,
sessions_by_user: HashMap<String, HashMap<u64, CancellationToken>>,
}
pub(crate) struct ProxySharedState {
pub(crate) handshake: HandshakeSharedState,
pub(crate) middle_relay: MiddleRelaySharedState,
pub(crate) traffic_limiter: Arc<TrafficLimiter>,
pub(crate) direct_buffer_budget: Arc<DirectBufferBudget>,
user_admission: ParkingMutex<UserAdmissionState>,
user_admission: Arc<UserAdmissionAuthority>,
pub(crate) conntrack_pressure_active: AtomicBool,
pub(crate) conntrack_close_tx: Mutex<Option<mpsc::Sender<ConntrackCloseEvent>>>,
masking_fallback_permits: Arc<Semaphore>,
}
#[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) {
let mut admission = self.shared.user_admission.lock();
let remove_user = admission
.sessions_by_user
.get_mut(&self.key.0)
.map(|sessions| {
sessions.remove(&self.key.1);
sessions.is_empty()
})
.unwrap_or(false);
if remove_user {
admission.sessions_by_user.remove(&self.key.0);
}
}
}
impl ProxySharedState {
pub(crate) fn new() -> Arc<Self> {
Self::new_with_direct_buffer_budget(DirectBufferBudget::new(
@@ -137,6 +99,17 @@ impl ProxySharedState {
/// Creates process state with the startup-resolved Direct buffer envelope.
pub(crate) fn new_with_direct_buffer_budget(
direct_buffer_budget: Arc<DirectBufferBudget>,
) -> Arc<Self> {
Self::new_with_direct_buffer_budget_and_user_admission(
direct_buffer_budget,
UserAdmissionAuthority::new(),
)
}
/// Creates generation state around one process-owned user authority.
pub(crate) fn new_with_direct_buffer_budget_and_user_admission(
direct_buffer_budget: Arc<DirectBufferBudget>,
user_admission: Arc<UserAdmissionAuthority>,
) -> Arc<Self> {
Arc::new(Self {
handshake: HandshakeSharedState {
@@ -167,7 +140,7 @@ impl ProxySharedState {
},
traffic_limiter: TrafficLimiter::new(),
direct_buffer_budget,
user_admission: ParkingMutex::new(UserAdmissionState::default()),
user_admission,
conntrack_pressure_active: AtomicBool::new(false),
conntrack_close_tx: Mutex::new(None),
masking_fallback_permits: Arc::new(Semaphore::new(MASKING_FALLBACK_MAX_CONCURRENT)),
@@ -183,106 +156,91 @@ impl ProxySharedState {
}
pub(crate) fn is_user_enabled(&self, user: &str) -> bool {
!self.user_admission.lock().disabled_users.contains(user)
self.user_admission.is_user_enabled(user)
}
pub(crate) fn set_user_enabled(&self, user: &str, enabled: bool) -> (bool, usize) {
let (newly_disabled, tokens) = {
let mut admission = self.user_admission.lock();
if enabled {
admission.disabled_users.remove(user);
(false, Vec::new())
} else {
let newly_disabled = admission.disabled_users.insert(user.to_string());
let tokens = admission
.sessions_by_user
.get(user)
.map(|sessions| sessions.values().cloned().collect())
.unwrap_or_default();
(newly_disabled, tokens)
}
};
for token in &tokens {
token.cancel();
}
(newly_disabled, tokens.len())
/// Returns the process authority shared by every runtime generation.
pub(crate) fn user_admission(&self) -> Arc<UserAdmissionAuthority> {
Arc::clone(&self.user_admission)
}
pub(crate) fn apply_user_enabled_config(
/// Reconciles the complete user authentication policy from configuration.
pub(crate) fn apply_user_config(
&self,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Vec<(String, usize)> {
let desired_disabled = user_enabled
.iter()
.filter_map(|(user, enabled)| (!*enabled).then_some(user.clone()))
.collect::<HashSet<_>>();
let cancellations = {
let mut admission = self.user_admission.lock();
let newly_disabled = desired_disabled
.difference(&admission.disabled_users)
.cloned()
.collect::<Vec<_>>();
admission.disabled_users = desired_disabled;
newly_disabled
.into_iter()
.map(|user| {
let tokens = admission
.sessions_by_user
.get(&user)
.map(|sessions| sessions.values().cloned().collect())
.unwrap_or_default();
(user, tokens)
})
.collect::<Vec<(String, Vec<CancellationToken>)>>()
};
cancellations
.into_iter()
.map(|(user, tokens)| {
for token in &tokens {
token.cancel();
}
(user, tokens.len())
})
.collect()
self.user_admission.apply_config(users, user_enabled)
}
/// Applies a candidate user policy only when its captured epoch is current.
pub(crate) fn apply_user_config_if_epoch(
&self,
expected_epoch: u64,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
self.user_admission
.apply_config_if_epoch(expected_epoch, users, user_enabled)
}
/// Applies one persisted user mutation before asynchronous config reload.
pub(crate) fn stage_user(
&self,
user: &str,
secret: &str,
enabled: bool,
) -> Option<UserMutationResult> {
self.user_admission.stage_user(user, secret, enabled)
}
/// Installs a deletion tombstone and cancels every current owner.
pub(crate) fn delete_user(&self, user: &str) -> UserMutationResult {
self.user_admission.delete_user(user)
}
/// Returns the current incarnation for an exact authenticated credential.
pub(crate) fn authenticated_user_incarnation(
&self,
user: &str,
credential_id: UserCredentialId,
) -> Option<UserIncarnation> {
self.user_admission
.authenticated_incarnation(user, credential_id)
}
/// Starts an atomic publication boundary for an authenticated owner.
pub(crate) fn claim_authenticated_user(
self: &Arc<Self>,
user: &str,
credential_id: UserCredentialId,
) -> Option<UserAdmissionPublication<'_>> {
self.user_admission
.claim_authenticated(user, credential_id)
}
pub(crate) fn register_user_session(
self: &Arc<Self>,
user: &str,
session_id: u64,
_session_id: u64,
) -> Option<UserSessionRegistration> {
let token = CancellationToken::new();
let key = (user.to_string(), session_id);
let mut admission = self.user_admission.lock();
if admission.disabled_users.contains(user) {
return None;
}
admission
.sessions_by_user
.entry(key.0.clone())
.or_default()
.insert(session_id, token.clone());
Some(UserSessionRegistration {
token,
_guard: UserSessionGuard {
shared: Arc::clone(self),
key,
},
})
self.user_admission.register_legacy(user)
}
/// Registers a relay session against the exact credential that authenticated it.
pub(crate) fn register_authenticated_user_session(
self: &Arc<Self>,
user: &str,
credential_id: UserCredentialId,
) -> Option<UserSessionRegistration> {
let mut publication = self.claim_authenticated_user(user, credential_id)?;
let registration = publication.take_registration()?;
publication.commit();
Some(registration)
}
pub(crate) fn cancel_user_sessions(&self, user: &str) -> usize {
let tokens: Vec<CancellationToken> = self
.user_admission
.lock()
.sessions_by_user
.get(user)
.map(|sessions| sessions.values().cloned().collect())
.unwrap_or_default();
for token in &tokens {
token.cancel();
}
tokens.len()
self.user_admission.cancel_user_owners(user)
}
pub(crate) fn set_conntrack_close_sender(&self, tx: mpsc::Sender<ConntrackCloseEvent>) {
@@ -350,31 +308,53 @@ impl ProxySharedState {
mod tests {
use super::*;
const ALICE_SECRET: &str = "00112233445566778899aabbccddeeff";
fn configured_shared() -> Arc<ProxySharedState> {
let shared = ProxySharedState::new();
let users = HashMap::from([
("alice".to_string(), ALICE_SECRET.to_string()),
(
"bob".to_string(),
"ffeeddccbbaa99887766554433221100".to_string(),
),
]);
shared.apply_user_config(&users, &HashMap::new());
shared
}
#[test]
fn user_enabled_config_sync_tracks_disabled_overrides() {
let shared = ProxySharedState::new();
let shared = configured_shared();
assert!(shared.is_user_enabled("alice"));
let users = HashMap::from([
("alice".to_string(), ALICE_SECRET.to_string()),
(
"bob".to_string(),
"ffeeddccbbaa99887766554433221100".to_string(),
),
]);
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);
let mut newly_disabled = shared.apply_user_config(&users, &user_enabled);
newly_disabled.sort();
assert_eq!(newly_disabled, vec![("alice".to_string(), 0)]);
assert!(!shared.is_user_enabled("alice"));
assert!(shared.is_user_enabled("bob"));
assert!(shared.apply_user_enabled_config(&user_enabled).is_empty());
assert!(shared.apply_user_config(&users, &user_enabled).is_empty());
user_enabled.clear();
assert!(shared.apply_user_enabled_config(&user_enabled).is_empty());
assert!(shared.apply_user_config(&users, &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 shared = configured_shared();
let alice_1 = shared.register_user_session("alice", 1).unwrap();
let alice_2 = shared.register_user_session("alice", 2).unwrap();
let bob = shared.register_user_session("bob", 1).unwrap();
@@ -392,9 +372,11 @@ mod tests {
#[test]
fn disabled_user_cannot_register_after_the_cancellation_snapshot() {
let shared = ProxySharedState::new();
let shared = configured_shared();
assert_eq!(shared.set_user_enabled("alice", false), (true, 0));
let result = shared.stage_user("alice", ALICE_SECRET, false).unwrap();
assert!(result.newly_disabled);
assert_eq!(result.cancelled, 0);
assert_eq!(shared.cancel_user_sessions("alice"), 0);
let late = shared.register_user_session("alice", 1);
@@ -406,15 +388,19 @@ mod tests {
#[test]
fn disabling_user_cancels_existing_sessions_before_return() {
let shared = ProxySharedState::new();
let shared = configured_shared();
let registration = shared.register_user_session("alice", 1).unwrap();
let token = registration.token();
assert_eq!(shared.set_user_enabled("alice", false), (true, 1));
let result = shared.stage_user("alice", ALICE_SECRET, false).unwrap();
assert!(result.newly_disabled);
assert_eq!(result.cancelled, 1);
assert!(token.is_cancelled());
assert!(shared.register_user_session("alice", 2).is_none());
assert_eq!(shared.set_user_enabled("alice", true), (false, 0));
let result = shared.stage_user("alice", ALICE_SECRET, true).unwrap();
assert!(!result.newly_disabled);
assert_eq!(result.cancelled, 0);
assert!(shared.register_user_session("alice", 3).is_some());
}
@@ -439,7 +425,8 @@ mod tests {
for session_id in 0..ITERATIONS as u64 {
let user = format!("user-{session_id}");
barrier.wait();
shared.set_user_enabled(&user, false);
let secret = "00112233445566778899aabbccddeeff";
shared.stage_user(&user, secret, false).unwrap();
}
for registration in register.join().unwrap().into_iter().flatten() {
+598
View File
@@ -0,0 +1,598 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use parking_lot::{Mutex, MutexGuard};
use tokio_util::sync::CancellationToken;
use crate::crypto::sha256;
/// Stable secret identity used to fence authentication across runtime generations.
pub(crate) type UserCredentialId = [u8; 16];
/// Monotonic identity of one configured username lifetime.
pub(crate) type UserIncarnation = u64;
#[derive(Clone, Copy, PartialEq, Eq)]
struct EffectiveUser {
credential_id: UserCredentialId,
enabled: bool,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum UserOverride {
Present(EffectiveUser),
Deleted,
}
struct UserRecord {
configured: Option<EffectiveUser>,
mutation_override: Option<UserOverride>,
incarnation: UserIncarnation,
}
impl UserRecord {
fn effective(&self) -> Option<EffectiveUser> {
match self.mutation_override {
Some(UserOverride::Present(user)) => Some(user),
Some(UserOverride::Deleted) => None,
None => self.configured,
}
}
}
struct RegisteredOwner {
token: CancellationToken,
incarnation: UserIncarnation,
}
#[derive(Default)]
struct UserAdmissionState {
initialized: bool,
epoch: u64,
next_incarnation: UserIncarnation,
next_registration_id: u64,
users: HashMap<String, UserRecord>,
owners_by_user: HashMap<String, HashMap<u64, RegisteredOwner>>,
}
impl UserAdmissionState {
fn allocate_incarnation(&mut self) -> UserIncarnation {
self.next_incarnation = self.next_incarnation.checked_add(1).unwrap_or(u64::MAX);
self.next_incarnation
}
fn allocate_registration_id(&mut self) -> Option<u64> {
let next = self.next_registration_id.checked_add(1)?;
self.next_registration_id = next;
Some(next)
}
fn bump_epoch(&mut self) {
self.epoch = self.epoch.checked_add(1).unwrap_or(u64::MAX);
}
fn owner_tokens(&self, user: &str) -> Vec<CancellationToken> {
self.owners_by_user
.get(user)
.map(|owners| owners.values().map(|owner| owner.token.clone()).collect())
.unwrap_or_default()
}
}
/// Result of one durable user mutation applied to the process admission authority.
pub(crate) struct UserMutationResult {
/// Incarnation invalidated or created by the mutation.
pub(crate) incarnation: UserIncarnation,
/// Number of live owners cancelled by the mutation.
pub(crate) cancelled: usize,
/// Whether the effective enabled state changed from enabled to disabled.
pub(crate) newly_disabled: bool,
}
/// Process-owned user authentication and live-owner authority.
pub(crate) struct UserAdmissionAuthority {
state: Mutex<UserAdmissionState>,
}
impl UserAdmissionAuthority {
/// Creates an uninitialized authority for isolated tests and startup wiring.
pub(crate) fn new() -> Arc<Self> {
Arc::new(Self {
state: Mutex::new(UserAdmissionState::default()),
})
}
/// Returns the mutation epoch used to reject stale candidate configuration.
pub(crate) fn epoch(&self) -> u64 {
self.state.lock().epoch
}
/// Reconciles the complete configured user set into the process authority.
pub(crate) fn apply_config(
&self,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Vec<(String, usize)> {
self.apply_config_locked(None, users, user_enabled)
.unwrap_or_default()
}
/// Applies a candidate configuration only if no newer authority mutation occurred.
pub(crate) fn apply_config_if_epoch(
&self,
expected_epoch: u64,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
self.apply_config_locked(Some(expected_epoch), users, user_enabled)
}
fn apply_config_locked(
&self,
expected_epoch: Option<u64>,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
let configured = users
.iter()
.filter_map(|(user, secret)| {
credential_id_from_hex(secret).map(|credential_id| {
(
user.clone(),
EffectiveUser {
credential_id,
enabled: user_enabled.get(user).copied().unwrap_or(true),
},
)
})
})
.collect::<HashMap<_, _>>();
let cancellations = {
let mut state = self.state.lock();
if expected_epoch.is_some_and(|epoch| state.epoch != epoch) {
return None;
}
let mut changed = !state.initialized;
state.initialized = true;
let existing_users = state.users.keys().cloned().collect::<Vec<_>>();
let mut cancellations = Vec::new();
for user in existing_users {
let desired = configured.get(&user).copied();
let old_effective = state.users.get(&user).and_then(UserRecord::effective);
let override_matches = state.users.get(&user).is_some_and(|record| {
matches!(
(record.mutation_override, desired),
(Some(UserOverride::Present(current)), Some(next)) if current == next
) || matches!(
(record.mutation_override, desired),
(Some(UserOverride::Deleted), None)
)
});
if let Some(record) = state.users.get_mut(&user) {
if record.configured != desired || override_matches {
changed = true;
}
record.configured = desired;
if override_matches {
record.mutation_override = None;
}
}
let new_effective = state.users.get(&user).and_then(UserRecord::effective);
if old_effective != new_effective {
let identity_changed = old_effective.map(|entry| entry.credential_id)
!= new_effective.map(|entry| entry.credential_id);
if identity_changed {
let incarnation = state.allocate_incarnation();
if let Some(record) = state.users.get_mut(&user) {
record.incarnation = incarnation;
}
}
if identity_changed
|| old_effective.is_some_and(|entry| entry.enabled)
&& new_effective.is_none_or(|entry| !entry.enabled)
{
let tokens = state.owner_tokens(&user);
cancellations.push((user, tokens));
}
}
}
for (user, desired) in configured {
if state.users.contains_key(&user) {
continue;
}
changed = true;
let incarnation = state.allocate_incarnation();
state.users.insert(
user,
UserRecord {
configured: Some(desired),
mutation_override: None,
incarnation,
},
);
}
if changed {
state.bump_epoch();
}
cancellations
};
Some(cancel_owners(cancellations))
}
/// Applies one persisted user value ahead of asynchronous runtime reload.
pub(crate) fn stage_user(
&self,
user: &str,
secret: &str,
enabled: bool,
) -> Option<UserMutationResult> {
let credential_id = credential_id_from_hex(secret)?;
let desired = EffectiveUser {
credential_id,
enabled,
};
let (incarnation, newly_disabled, tokens) = {
let mut state = self.state.lock();
let previous = state.users.get(user).and_then(UserRecord::effective);
let identity_changed = previous.map(|entry| entry.credential_id) != Some(credential_id);
let incarnation = if identity_changed {
state.allocate_incarnation()
} else {
state
.users
.get(user)
.map(|record| record.incarnation)
.unwrap_or_else(|| state.allocate_incarnation())
};
let record = state.users.entry(user.to_string()).or_insert(UserRecord {
configured: None,
mutation_override: None,
incarnation,
});
record.mutation_override = Some(UserOverride::Present(desired));
record.incarnation = incarnation;
state.initialized = true;
state.bump_epoch();
let newly_disabled = previous.is_some_and(|entry| entry.enabled) && !enabled;
let tokens = if identity_changed || !enabled {
state.owner_tokens(user)
} else {
Vec::new()
};
(incarnation, newly_disabled, tokens)
};
let cancelled = tokens.len();
for token in tokens {
token.cancel();
}
Some(UserMutationResult {
incarnation,
cancelled,
newly_disabled,
})
}
/// Installs a deletion tombstone and cancels every owner of the old incarnation.
pub(crate) fn delete_user(&self, user: &str) -> UserMutationResult {
let (incarnation, newly_disabled, tokens) = {
let mut state = self.state.lock();
let previous = state.users.get(user).and_then(UserRecord::effective);
let incarnation = state.allocate_incarnation();
let record = state.users.entry(user.to_string()).or_insert(UserRecord {
configured: None,
mutation_override: None,
incarnation,
});
record.mutation_override = Some(UserOverride::Deleted);
record.incarnation = incarnation;
state.initialized = true;
state.bump_epoch();
(
incarnation,
previous.is_some_and(|entry| entry.enabled),
state.owner_tokens(user),
)
};
let cancelled = tokens.len();
for token in tokens {
token.cancel();
}
UserMutationResult {
incarnation,
cancelled,
newly_disabled,
}
}
/// Returns whether the effective process policy currently enables a user.
pub(crate) fn is_user_enabled(&self, user: &str) -> bool {
let state = self.state.lock();
if !state.initialized {
return true;
}
state
.users
.get(user)
.and_then(UserRecord::effective)
.is_some_and(|entry| entry.enabled)
}
/// Returns the authenticated incarnation for an exact current credential.
pub(crate) fn authenticated_incarnation(
&self,
user: &str,
credential_id: UserCredentialId,
) -> Option<UserIncarnation> {
let state = self.state.lock();
if !state.initialized {
return Some(0);
}
let record = state.users.get(user)?;
let effective = record.effective()?;
(effective.enabled && effective.credential_id == credential_id).then_some(record.incarnation)
}
/// Starts a short publication critical section for one authenticated owner.
pub(crate) fn claim_authenticated(
self: &Arc<Self>,
user: &str,
credential_id: UserCredentialId,
) -> Option<UserAdmissionPublication<'_>> {
let mut state = self.state.lock();
let incarnation = if state.initialized {
let record = state.users.get(user)?;
let effective = record.effective()?;
if !effective.enabled || effective.credential_id != credential_id {
return None;
}
record.incarnation
} else {
0
};
let registration_id = state.allocate_registration_id()?;
let token = CancellationToken::new();
let active = Arc::new(AtomicBool::new(false));
Some(UserAdmissionPublication {
state,
authority: Arc::clone(self),
user: user.to_string(),
registration_id,
incarnation,
token,
active,
registration_taken: false,
})
}
/// Registers a legacy owner when no credential snapshot is available.
pub(crate) fn register_legacy(
self: &Arc<Self>,
user: &str,
) -> Option<UserSessionRegistration> {
let credential_id = {
let state = self.state.lock();
if !state.initialized {
[0; 16]
} else {
state.users.get(user)?.effective()?.credential_id
}
};
let mut publication = self.claim_authenticated(user, credential_id)?;
let registration = publication.take_registration()?;
publication.commit();
Some(registration)
}
/// Cancels all current owners without changing admission policy.
pub(crate) fn cancel_user_owners(&self, user: &str) -> usize {
let tokens = self.state.lock().owner_tokens(user);
let count = tokens.len();
for token in tokens {
token.cancel();
}
count
}
fn unregister(&self, user: &str, registration_id: u64, incarnation: UserIncarnation) {
let mut state = self.state.lock();
let remove_user = state
.owners_by_user
.get_mut(user)
.map(|owners| {
if owners
.get(&registration_id)
.is_some_and(|owner| owner.incarnation == incarnation)
{
owners.remove(&registration_id);
}
owners.is_empty()
})
.unwrap_or(false);
if remove_user {
state.owners_by_user.remove(user);
}
}
}
/// 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<AtomicBool>,
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;
}
self.state
.owners_by_user
.entry(self.user.clone())
.or_default()
.insert(
self.registration_id,
RegisteredOwner {
token: self.token.clone(),
incarnation: self.incarnation,
},
);
self.active.store(true, Ordering::Release);
}
}
/// 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<AtomicBool>,
}
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(false, Ordering::AcqRel) {
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);
let mut id = [0; 16];
id.copy_from_slice(&digest[..16]);
id
}
/// Decodes one configured secret and derives its credential identity.
pub(crate) fn credential_id_from_hex(secret: &str) -> Option<UserCredentialId> {
let decoded = hex::decode(secret).ok()?;
let secret: [u8; 16] = decoded.try_into().ok()?;
Some(credential_id(&secret))
}
fn cancel_owners(
cancellations: Vec<(String, Vec<CancellationToken>)>,
) -> Vec<(String, usize)> {
cancellations
.into_iter()
.map(|(user, tokens)| {
let count = tokens.len();
for token in tokens {
token.cancel();
}
(user, count)
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn users(secret: &str) -> HashMap<String, String> {
HashMap::from([("alice".to_string(), secret.to_string())])
}
#[test]
fn shared_authority_rejects_registration_through_an_old_generation() {
let authority = UserAdmissionAuthority::new();
let secret = "00112233445566778899aabbccddeeff";
authority.apply_config(&users(secret), &HashMap::new());
let credential = credential_id_from_hex(secret).unwrap();
assert!(authority.claim_authenticated("alice", credential).is_some());
authority.stage_user("alice", secret, false).unwrap();
assert!(authority.claim_authenticated("alice", credential).is_none());
}
#[test]
fn stale_credential_cannot_cross_delete_and_recreate() {
let authority = UserAdmissionAuthority::new();
let old_secret = "00112233445566778899aabbccddeeff";
let new_secret = "ffeeddccbbaa99887766554433221100";
authority.apply_config(&users(old_secret), &HashMap::new());
let old_credential = credential_id_from_hex(old_secret).unwrap();
let old_incarnation = authority
.authenticated_incarnation("alice", old_credential)
.unwrap();
authority.delete_user("alice");
let recreated = authority.stage_user("alice", new_secret, true).unwrap();
assert!(recreated.incarnation > old_incarnation);
assert!(
authority
.authenticated_incarnation("alice", old_credential)
.is_none()
);
}
#[test]
fn stale_candidate_cannot_overwrite_newer_mutation() {
let authority = UserAdmissionAuthority::new();
let secret = "00112233445566778899aabbccddeeff";
authority.apply_config(&users(secret), &HashMap::new());
let candidate_epoch = authority.epoch();
authority.stage_user("alice", secret, false).unwrap();
assert!(
authority
.apply_config_if_epoch(candidate_epoch, &users(secret), &HashMap::new())
.is_none()
);
assert!(!authority.is_user_enabled("alice"));
}
}