mirror of
https://github.com/telemt/telemt.git
synced 2026-10-10 19:35:58 +03:00
Proxy Shared User Drafts
This commit is contained in:
@@ -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,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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
@@ -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() {
|
||||
|
||||
@@ -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(®istration_id)
|
||||
.is_some_and(|owner| owner.incarnation == incarnation)
|
||||
{
|
||||
owners.remove(®istration_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"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user