astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-10-10 11:25:57 +03:00
Code Issues Packages Projects Releases Wiki Activity

Hardswap Invariants in tests + Quota fixes

Browse Source
This commit is contained in:
Alexey
2026-09-20 00:28:52 +03:00
parent 89dacbd17e
commit d706b3f3ba
66 changed files with 1821 additions and 505 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/proxy/authenticated.rs
+33 -3
View File
@@ -15,7 +15,7 @@ use crate::proxy::middle_relay::{handle_via_middle_proxy, handle_via_middle_prox
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::stats::{Stats, UserQuotaHandle};
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
@@ -85,6 +85,7 @@ where
warn!(user = %user, error = %error, "User admission check failed");
error
})?;
let quota_handle = user_reservation.quota_handle();
let route_snapshot = deps.route_runtime.snapshot();
let session_id = deps.rng.u64();
@@ -137,6 +138,7 @@ where
session_id,
session_cancel.clone(),
Arc::clone(&deps.shared),
quota_handle.clone(),
)
.await
} else {
@@ -156,6 +158,7 @@ where
session_cancel.clone(),
Arc::clone(&deps.shared),
ConntrackClosePolicy::Suppress,
quota_handle.clone(),
)
.await
}
@@ -171,6 +174,7 @@ where
local_addr,
session_cancel.clone(),
conntrack_close_policy,
quota_handle.clone(),
)
.await
}
@@ -185,6 +189,7 @@ where
local_addr,
session_cancel,
conntrack_close_policy,
quota_handle,
)
.await
};
@@ -202,6 +207,7 @@ async fn run_direct<R, W>(
local_addr: SocketAddr,
session_cancel: tokio_util::sync::CancellationToken,
conntrack_close_policy: ConntrackClosePolicy,
quota_handle: UserQuotaHandle,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
@@ -223,6 +229,7 @@ where
session_cancel,
Arc::clone(&deps.shared),
conntrack_close_policy,
quota_handle,
)
.await
}
@@ -235,6 +242,7 @@ pub(crate) struct UserConnectionReservation {
user: String,
ip: IpAddr,
incarnation: UserIncarnation,
quota_handle: UserQuotaHandle,
tracks_ip: bool,
active: bool,
}
@@ -248,7 +256,16 @@ impl UserConnectionReservation {
ip: IpAddr,
tracks_ip: bool,
) -> Self {
Self::new_for_incarnation(stats, ip_tracker, user, ip, 0, tracks_ip)
let quota_handle = stats.current_user_quota_handle(&user);
Self::new_for_incarnation(
stats,
ip_tracker,
user,
ip,
0,
quota_handle,
tracks_ip,
)
}
/// Creates a reservation fenced to one authenticated user incarnation.
@@ -258,6 +275,7 @@ impl UserConnectionReservation {
user: String,
ip: IpAddr,
incarnation: UserIncarnation,
quota_handle: UserQuotaHandle,
tracks_ip: bool,
) -> Self {
Self {
@@ -266,11 +284,17 @@ impl UserConnectionReservation {
user,
ip,
incarnation,
quota_handle,
tracks_ip,
active: true,
}
}
/// Returns quota ownership pinned to the authenticated user incarnation.
pub(crate) fn quota_handle(&self) -> UserQuotaHandle {
self.quota_handle.clone()
}
/// Releases both admission counters through the asynchronous cleanup path.
pub(crate) async fn release(mut self) {
if !self.active {
@@ -354,8 +378,13 @@ async fn acquire_user_connection_reservation_for_incarnation(
user: user.to_string(),
});
}
let Some(quota_handle) = stats.quota_handle_for_incarnation(user, incarnation) else {
return Err(ProxyError::UserDisabled {
user: user.to_string(),
});
};
if let Some(quota) = config.access.user_data_quota.get(user)
&& stats.get_user_quota_used(user) >= *quota
&& quota_handle.used() >= *quota
{
return Err(ProxyError::DataQuotaExceeded {
user: user.to_string(),
@@ -399,6 +428,7 @@ async fn acquire_user_connection_reservation_for_incarnation(
user.to_string(),
peer_addr.ip(),
incarnation,
quota_handle,
true,
))
}
src/proxy/client/authenticated.rs
+8 -1
View File
@@ -38,7 +38,14 @@ impl RunningClientHandler {
} else {
config
};
let shared = ProxySharedState::new();
let shared = ProxySharedState::new_with_direct_buffer_budget_and_user_admission(
crate::proxy::direct_buffer_budget::DirectBufferBudget::new(
crate::proxy::direct_buffer_budget::fallback_direct_buffer_hard_limit(),
),
crate::proxy::user_admission::UserAdmissionAuthority::new_with_quota_store(
stats.quota_store(),
),
);
shared.apply_user_config(&config.access.users, &config.access.user_enabled);
Self::handle_authenticated_static_with_shared(
client_reader,
src/proxy/direct_relay.rs
+1
View File
@@ -27,6 +27,7 @@ use crate::proxy::shared_state::{
ProxySharedState,
};
use crate::stats::Stats;
use crate::stats::UserQuotaHandle;
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
use crate::transport::UpstreamManager;
#[cfg(unix)]
src/proxy/direct_relay/relay.rs
+4
View File
@@ -59,6 +59,7 @@ where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let quota_handle = stats.current_user_quota_handle(&success.user);
handle_via_direct_with_shared_and_conntrack(
client_reader,
client_writer,
@@ -75,6 +76,7 @@ where
session_cancel,
shared,
ConntrackClosePolicy::Publish,
quota_handle,
)
.await
}
@@ -96,6 +98,7 @@ pub(crate) async fn handle_via_direct_with_shared_and_conntrack<R, W>(
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
conntrack_close_policy: ConntrackClosePolicy,
quota_handle: UserQuotaHandle,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
@@ -171,6 +174,7 @@ where
config.server.max_connections,
user,
Arc::clone(&stats),
quota_handle,
config.access.user_data_quota.get(user).copied(),
traffic_lease,
relay_activity_timeout,
src/proxy/middle_relay.rs
+4 -1
View File
@@ -31,7 +31,8 @@ use crate::proxy::shared_state::{
};
use crate::proxy::traffic_limiter::{RateDirection, TrafficLease, next_refill_delay};
use crate::stats::{
MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, UserStats,
MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats,
UserQuotaHandle, UserStats,
};
use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer};
use crate::transport::middle_proxy::{ConnLease, MePool, MeResponse, proto_flags_for_tag};
@@ -108,6 +109,7 @@ pub(crate) async fn handle_via_middle_proxy<R, W>(
session_id: u64,
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
quota_handle: UserQuotaHandle,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
@@ -129,6 +131,7 @@ where
session_cancel,
shared,
ConntrackClosePolicy::Publish,
quota_handle,
)
.await
}
src/proxy/middle_relay/d2c.rs
+5 -2
View File
@@ -135,6 +135,7 @@ pub(crate) async fn process_me_writer_response<W>(
where
W: AsyncWrite + Unpin + Send + 'static,
{
let quota_handle = quota_limit.map(|_| stats.current_user_quota_handle(user));
process_me_writer_response_with_traffic_lease(
response,
client_writer,
@@ -144,6 +145,7 @@ where
stats,
user,
quota_user_stats,
quota_handle.as_ref(),
quota_limit,
quota_soft_overshoot_bytes,
None,
@@ -165,6 +167,7 @@ pub(crate) async fn process_me_writer_response_with_traffic_lease<W>(
stats: &Stats,
user: &str,
quota_user_stats: Option<&UserStats>,
quota_handle: Option<&UserQuotaHandle>,
quota_limit: Option<u64>,
quota_soft_overshoot_bytes: u64,
traffic_lease: Option<&Arc<TrafficLease>>,
@@ -185,10 +188,10 @@ where
trace!(conn_id, bytes = data.len(), flags, "ME->C data");
}
let data_len = data.len() as u64;
if let (Some(limit), Some(user_stats)) = (quota_limit, quota_user_stats) {
if let (Some(limit), Some(quota_handle)) = (quota_limit, quota_handle) {
let soft_limit = quota_soft_cap(limit, quota_soft_overshoot_bytes);
match reserve_user_quota_with_yield(
user_stats, data_len, soft_limit, stats, cancel, None,
quota_handle, data_len, soft_limit, stats, cancel, None,
)
.await
{
src/proxy/middle_relay/quota.rs
+3 -3
View File
@@ -12,7 +12,7 @@ pub(super) fn quota_soft_cap(limit: u64, overshoot: u64) -> u64 {
}
pub(super) async fn reserve_user_quota_with_yield(
user_stats: &UserStats,
quota_handle: &UserQuotaHandle,
bytes: u64,
limit: u64,
stats: &Stats,
@@ -23,8 +23,8 @@ pub(super) async fn reserve_user_quota_with_yield(
let mut backoff_rounds = 0usize;
loop {
for _ in 0..QUOTA_RESERVE_SPIN_RETRIES {
match user_stats.quota_try_reserve(bytes, limit) {
Ok(total) => return Ok(total),
match quota_handle.try_reserve(bytes, limit) {
Ok(reservation) => return Ok(reservation.commit()),
Err(QuotaReserveError::LimitExceeded) => {
return Err(MiddleQuotaReserveError::LimitExceeded);
}
src/proxy/middle_relay/session.rs
+12 -4
View File
@@ -61,6 +61,7 @@ pub(crate) async fn handle_via_middle_proxy_with_conntrack<R, W>(
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
conntrack_close_policy: ConntrackClosePolicy,
quota_handle: UserQuotaHandle,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
@@ -73,6 +74,7 @@ where
let quota_limit = config.access.user_data_quota.get(&user).copied();
let quota_user_stats = quota_limit.map(|_| stats.get_or_create_user_stats_handle(&user));
let quota_handle = quota_limit.map(|_| quota_handle);
let peer = success.peer;
let traffic_lease = shared.traffic_limiter.acquire_lease(&user, peer.ip());
let proto_tag = success.proto_tag;
@@ -200,6 +202,7 @@ where
let rng_clone = rng.clone();
let user_clone = user.clone();
let quota_user_stats_me_writer = quota_user_stats.clone();
let quota_handle_me_writer = quota_handle.clone();
let traffic_lease_me_writer = traffic_lease.clone();
let flow_cancel_me_writer = flow_cancel.clone();
let last_downstream_activity_ms_clone = last_downstream_activity_ms.clone();
@@ -212,6 +215,7 @@ where
rng_clone,
user_clone,
quota_user_stats_me_writer,
quota_handle_me_writer,
quota_limit,
traffic_lease_me_writer,
flow_cancel_me_writer,
@@ -340,11 +344,11 @@ where
forensics.bytes_c2me = forensics
.bytes_c2me
.saturating_add(payload.len() as u64);
if let (Some(limit), Some(user_stats)) =
(quota_limit, quota_user_stats.as_deref())
if let (Some(limit), Some(quota_handle)) =
(quota_limit, quota_handle.as_ref())
{
match reserve_user_quota_with_yield(
user_stats,
quota_handle,
payload.len() as u64,
limit,
stats.as_ref(),
@@ -379,7 +383,11 @@ where
break;
}
}
stats.add_user_octets_from_handle(user_stats, payload.len() as u64);
if let Some(user_stats) = quota_user_stats.as_deref() {
stats.add_user_octets_from_handle(user_stats, payload.len() as u64);
} else {
stats.add_user_octets_from(&user, payload.len() as u64);
}
} else {
stats.add_user_octets_from(&user, payload.len() as u64);
}
src/proxy/middle_relay/session/tasks.rs
+5
View File
@@ -53,6 +53,7 @@ pub(super) async fn run_me_writer<W>(
rng_clone: Arc<SecureRandom>,
user_clone: String,
quota_user_stats_me_writer: Option<Arc<UserStats>>,
quota_handle_me_writer: Option<UserQuotaHandle>,
quota_limit: Option<u64>,
traffic_lease_me_writer: Option<Arc<TrafficLease>>,
flow_cancel_me_writer: CancellationToken,
@@ -105,6 +106,7 @@ where
stats_clone.as_ref(),
&user_clone,
quota_user_stats_me_writer.as_deref(),
quota_handle_me_writer.as_ref(),
quota_limit,
d2c_flush_policy.quota_soft_overshoot_bytes,
traffic_lease_me_writer.as_ref(),
@@ -167,6 +169,7 @@ where
stats_clone.as_ref(),
&user_clone,
quota_user_stats_me_writer.as_deref(),
quota_handle_me_writer.as_ref(),
quota_limit,
d2c_flush_policy.quota_soft_overshoot_bytes,
traffic_lease_me_writer.as_ref(),
@@ -233,6 +236,7 @@ where
stats_clone.as_ref(),
&user_clone,
quota_user_stats_me_writer.as_deref(),
quota_handle_me_writer.as_ref(),
quota_limit,
d2c_flush_policy.quota_soft_overshoot_bytes,
traffic_lease_me_writer.as_ref(),
@@ -304,6 +308,7 @@ where
stats_clone.as_ref(),
&user_clone,
quota_user_stats_me_writer.as_deref(),
quota_handle_me_writer.as_ref(),
quota_limit,
d2c_flush_policy.quota_soft_overshoot_bytes,
traffic_lease_me_writer.as_ref(),
src/proxy/relay.rs
+2
View File
@@ -290,6 +290,7 @@ where
// ── Combine split halves into bidirectional streams ──────────────
let client_combined = CombinedStream::new(client_reader, client_writer);
let mut server = CombinedStream::new(server_reader, server_writer);
let quota_handle = stats.current_user_quota_handle(&user_owned);
// Wrap client with stats/activity tracking
let mut client = StatsIo::new_with_traffic_lease(
@@ -297,6 +298,7 @@ where
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
quota_handle,
traffic_lease,
quota_limit,
Arc::clone(&quota_exceeded),
src/proxy/relay/adaptive_copy.rs
+4 -1
View File
@@ -19,7 +19,7 @@ use crate::proxy::direct_buffer_budget::{
DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES, DirectBufferBudget, DirectBufferLease,
};
use crate::proxy::traffic_limiter::TrafficLease;
use crate::stats::Stats;
use crate::stats::{Stats, UserQuotaHandle};
use super::WATCHDOG_INTERVAL;
use super::io::{SharedCounters, StatsIo, is_quota_io_error};
@@ -141,6 +141,7 @@ pub(crate) async fn relay_direct_adaptive<CR, CW, SR, SW>(
max_connections: u32,
user: &str,
stats: Arc<Stats>,
quota_handle: UserQuotaHandle,
quota_limit: Option<u64>,
traffic_lease: Option<Arc<TrafficLease>>,
activity_timeout: Duration,
@@ -200,6 +201,7 @@ where
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
quota_handle.clone(),
traffic_lease.clone(),
quota_limit,
Arc::clone(&quota_exceeded),
@@ -210,6 +212,7 @@ where
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
quota_handle,
traffic_lease,
quota_limit,
Arc::clone(&quota_exceeded),
src/proxy/relay/io.rs
+14 -9
View File
@@ -1,5 +1,5 @@
use crate::proxy::traffic_limiter::{RateDirection, TrafficLease, next_refill_delay};
use crate::stats::{Stats, UserStats};
use crate::stats::{Stats, UserQuotaHandle, UserStats};
use std::io;
use std::pin::Pin;
use std::sync::Arc;
@@ -40,6 +40,7 @@ pub(super) struct StatsIo<S> {
stats: Arc<Stats>,
user: String,
user_stats: Arc<UserStats>,
quota_handle: UserQuotaHandle,
traffic_lease: Option<Arc<TrafficLease>>,
c2s_rate_debt_bytes: u64,
c2s_wait: RateWaitState,
@@ -71,11 +72,13 @@ impl<S> StatsIo<S> {
quota_exceeded: Arc<AtomicBool>,
epoch: Instant,
) -> Self {
let quota_handle = stats.current_user_quota_handle(&user);
Self::new_with_traffic_lease(
inner,
counters,
stats,
user,
quota_handle,
None,
quota_limit,
quota_exceeded,
@@ -88,6 +91,7 @@ impl<S> StatsIo<S> {
counters: Arc<SharedCounters>,
stats: Arc<Stats>,
user: String,
quota_handle: UserQuotaHandle,
traffic_lease: Option<Arc<TrafficLease>>,
quota_limit: Option<u64>,
quota_exceeded: Arc<AtomicBool>,
@@ -102,6 +106,7 @@ impl<S> StatsIo<S> {
stats,
user,
user_stats,
quota_handle,
traffic_lease,
c2s_rate_debt_bytes: 0,
c2s_wait: RateWaitState::default(),
@@ -213,7 +218,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for StatsIo<S> {
let mut quota_reservation = None;
let mut read_limit = buf.remaining();
if let Some(limit) = this.quota_limit {
let used_before = this.user_stats.quota_used();
let used_before = this.quota_handle.used();
let remaining = limit.saturating_sub(used_before);
if remaining == 0 {
this.quota_exceeded.store(true, Ordering::Release);
@@ -230,7 +235,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for StatsIo<S> {
let mut reserve_rounds = 0usize;
while quota_reservation.is_none() {
for _ in 0..QUOTA_RESERVE_SPIN_RETRIES {
match this.user_stats.quota_reserve(desired, limit) {
match this.quota_handle.try_reserve(desired, limit) {
Ok(reservation) => {
quota_reservation = Some(reservation);
break;
@@ -305,7 +310,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for StatsIo<S> {
}
}
if let Some(limit) = this.quota_limit
&& this.user_stats.quota_used() >= limit
&& this.quota_handle.used() >= limit
{
this.quota_exceeded.store(true, Ordering::Release);
}
@@ -401,7 +406,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
if !write_buf.is_empty() {
let mut reserve_rounds = 0usize;
while quota_reservation.is_none() {
let used_before = this.user_stats.quota_used();
let used_before = this.quota_handle.used();
let remaining = limit.saturating_sub(used_before);
if remaining == 0 {
this.quota_exceeded.store(true, Ordering::Release);
@@ -412,7 +417,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
let desired = remaining.min(write_buf.len() as u64);
let mut saw_contention = false;
for _ in 0..QUOTA_RESERVE_SPIN_RETRIES {
match this.user_stats.quota_reserve(desired, limit) {
match this.quota_handle.try_reserve(desired, limit) {
Ok(reservation) => {
quota_reservation = Some(reservation);
write_buf = &write_buf[..desired as usize];
@@ -442,7 +447,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
}
}
} else {
let used_before = this.user_stats.quota_used();
let used_before = this.quota_handle.used();
let remaining = limit.saturating_sub(used_before);
if remaining == 0 {
this.quota_exceeded.store(true, Ordering::Release);
@@ -481,7 +486,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
if let (Some(limit), Some(remaining)) = (this.quota_limit, remaining_before) {
if should_immediate_quota_check(remaining, n_to_charge) {
this.quota_bytes_since_check = 0;
if this.user_stats.quota_used() >= limit {
if this.quota_handle.used() >= limit {
this.quota_exceeded.store(true, Ordering::Release);
}
} else {
@@ -490,7 +495,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
let interval = quota_adaptive_interval_bytes(remaining);
if this.quota_bytes_since_check >= interval {
this.quota_bytes_since_check = 0;
if this.user_stats.quota_used() >= limit {
if this.quota_handle.used() >= limit {
this.quota_exceeded.store(true, Ordering::Release);
}
}
src/proxy/shared_state.rs
+27 -4
View File
@@ -198,15 +198,27 @@ impl ProxySharedState {
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(
/// Transfers user-policy ownership to one runtime generation.
pub(crate) fn activate_user_config_source(
&self,
expected_epoch: u64,
source_generation: u64,
expected_epoch: Option<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)
.activate_config_source(source_generation, expected_epoch, users, user_enabled)
}
/// Applies an update only from the active runtime generation.
pub(crate) fn apply_user_config_from_source(
&self,
source_generation: u64,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
self.user_admission
.apply_config_from_source(source_generation, users, user_enabled)
}
/// Applies one persisted user mutation before asynchronous config reload.
@@ -219,6 +231,17 @@ impl ProxySharedState {
self.user_admission.stage_user(user, secret, enabled)
}
/// Applies one prevalidated persisted credential before asynchronous reload.
pub(crate) fn stage_user_credential(
&self,
user: &str,
credential_id: UserCredentialId,
enabled: bool,
) -> UserMutationResult {
self.user_admission
.stage_user_credential(user, credential_id, 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)
src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs
+1
View File
@@ -240,6 +240,7 @@ async fn me_writer_data_write_obeys_flow_cancellation() {
user,
None,
None,
None,
0,
None,
&cancel,
src/proxy/traffic_limiter.rs
+23 -12
View File
@@ -90,20 +90,20 @@ struct DirectionBucket {
struct UserBucket {
rates: AtomicRatePair,
up: DirectionBucket,
down: DirectionBucket,
up: Arc<DirectionBucket>,
down: Arc<DirectionBucket>,
active_leases: AtomicU64,
}
#[derive(Default)]
struct CidrDirectionBucket {
used: DirectionBucket,
active_users: DirectionBucket,
used: Arc<DirectionBucket>,
active_users: Arc<DirectionBucket>,
}
#[derive(Default)]
struct CidrUserDirectionState {
used: DirectionBucket,
used: Arc<DirectionBucket>,
}
struct CidrUserShare {
@@ -155,17 +155,27 @@ struct ShardedRegistry<T> {
mask: usize,
}
pub struct TrafficLease {
struct TrafficLeaseBinding {
limiter: Arc<TrafficLimiter>,
revision: u64,
user_bucket: Option<Arc<UserBucket>>,
cidr_bucket: Option<Arc<CidrBucket>>,
cidr_user_key: Option<String>,
cidr_user_share: Option<Arc<CidrUserShare>>,
}
pub struct TrafficLease {
limiter: Arc<TrafficLimiter>,
user: String,
client_ip: IpAddr,
binding: ArcSwap<TrafficLeaseBinding>,
refresh: ParkingMutex<()>,
}
pub struct TrafficLimiter {
policy: ArcSwap<PolicySnapshot>,
policy_update: ParkingMutex<()>,
published_revision: AtomicU64,
user_buckets: ShardedRegistry<UserBucket>,
cidr_buckets: ShardedRegistry<CidrBucket>,
user_scope: ScopeMetrics,
@@ -173,17 +183,18 @@ pub struct TrafficLimiter {
last_cleanup_epoch_secs: AtomicU64,
}
struct DirectionDebit<'a> {
bucket: &'a DirectionBucket,
struct DirectionDebit {
bucket: Arc<DirectionBucket>,
epoch: u64,
refundable: u64,
}
/// Refunds uncommitted shaping budget when an I/O attempt is cancelled.
#[must_use = "traffic reservations must be settled after the I/O attempt"]
pub(crate) struct TrafficReservation<'a> {
pub(crate) struct TrafficReservation {
result: TrafficConsumeResult,
user: Option<DirectionDebit<'a>>,
cidr: Option<DirectionDebit<'a>>,
cidr_user: Option<DirectionDebit<'a>>,
_binding: Arc<TrafficLeaseBinding>,
user: Option<DirectionDebit>,
cidr: Option<DirectionDebit>,
cidr_user: Option<DirectionDebit>,
}
src/proxy/traffic_limiter/buckets.rs
+17 -17
View File
@@ -71,11 +71,11 @@ impl DirectionBucket {
}
pub(super) fn try_reserve_at(
&self,
self: &Arc<Self>,
epoch: u64,
cap: u64,
requested: u64,
) -> Option<DirectionDebit<'_>> {
) -> Option<DirectionDebit> {
if requested == 0 || cap == 0 || epoch > PACKED_EPOCH_MAX {
return None;
}
@@ -109,7 +109,7 @@ impl DirectionBucket {
) {
Ok(_) => {
return Some(DirectionDebit {
bucket: self,
bucket: Arc::clone(self),
epoch,
refundable: grant,
});
@@ -144,7 +144,7 @@ impl DirectionBucket {
}
}
impl DirectionDebit<'_> {
impl DirectionDebit {
fn granted(&self) -> u64 {
self.refundable
}
@@ -168,7 +168,7 @@ impl DirectionDebit<'_> {
}
}
impl Drop for DirectionDebit<'_> {
impl Drop for DirectionDebit {
fn drop(&mut self) {
self.bucket.refund_at(self.epoch, self.refundable);
}
@@ -178,8 +178,8 @@ impl UserBucket {
pub(super) fn new(revision: u64, limits: RateLimitBps) -> Self {
Self {
rates: AtomicRatePair::new(revision, limits),
up: DirectionBucket::default(),
down: DirectionBucket::default(),
up: Arc::new(DirectionBucket::default()),
down: Arc::new(DirectionBucket::default()),
active_leases: AtomicU64::new(0),
}
}
@@ -192,7 +192,7 @@ impl UserBucket {
&self,
direction: RateDirection,
requested: u64,
) -> (u64, Option<DirectionDebit<'_>>) {
) -> (u64, Option<DirectionDebit>) {
let cap_bps = self.rates.get(direction);
if cap_bps == 0 {
return (requested, None);
@@ -208,12 +208,12 @@ impl UserBucket {
}
impl CidrDirectionBucket {
pub(super) fn try_reserve<'a>(
&'a self,
user_state: &'a CidrUserDirectionState,
pub(super) fn try_reserve(
&self,
user_state: &CidrUserDirectionState,
cap_epoch: u64,
requested: u64,
) -> (u64, Option<DirectionDebit<'a>>, Option<DirectionDebit<'a>>) {
) -> (u64, Option<DirectionDebit>, Option<DirectionDebit>) {
if requested == 0 || cap_epoch == 0 {
return (0, None, None);
}
@@ -260,7 +260,7 @@ impl CidrDirectionBucket {
}
impl CidrUserDirectionState {
pub(super) fn ensure_active(&self, epoch: u64, active_users: &DirectionBucket) -> bool {
pub(super) fn ensure_active(&self, epoch: u64, active_users: &Arc<DirectionBucket>) -> bool {
if epoch > PACKED_EPOCH_MAX {
return false;
}
@@ -340,12 +340,12 @@ impl CidrBucket {
});
}
pub(super) fn try_reserve_for_user<'a>(
&'a self,
pub(super) fn try_reserve_for_user(
&self,
direction: RateDirection,
share: &'a CidrUserShare,
share: &CidrUserShare,
requested: u64,
) -> (u64, Option<DirectionDebit<'a>>, Option<DirectionDebit<'a>>) {
) -> (u64, Option<DirectionDebit>, Option<DirectionDebit>) {
let cap_bps = self.rates.get(direction);
if cap_bps == 0 {
return (requested, None, None);
src/proxy/traffic_limiter/lease.rs
+38 -5
View File
@@ -1,12 +1,41 @@
use super::*;
impl TrafficLease {
fn current_binding(&self) -> Arc<TrafficLeaseBinding> {
let published_revision = self.limiter.published_revision.load(Ordering::Acquire);
let current = self.binding.load_full();
if current.revision == published_revision {
return current;
}
let refresh = self.refresh.lock();
let published_revision = self.limiter.published_revision.load(Ordering::Acquire);
let current = self.binding.load_full();
if current.revision == published_revision {
return current;
}
let policy_update = self.limiter.policy_update.lock();
let policy = self.limiter.policy.load_full();
if current.revision == policy.revision {
return current;
}
let next = self
.limiter
.build_binding(&self.user, self.client_ip, &policy);
self.binding.store(Arc::clone(&next));
drop(policy_update);
drop(refresh);
self.limiter.maybe_cleanup();
next
}
/// Reserves shaping budget until the associated I/O result is settled.
pub(crate) fn try_reserve(
&self,
direction: RateDirection,
requested: u64,
) -> TrafficReservation<'_> {
) -> TrafficReservation {
let binding = self.current_binding();
if requested == 0 {
return TrafficReservation {
result: TrafficConsumeResult {
@@ -14,6 +43,7 @@ impl TrafficLease {
blocked_user: false,
blocked_cidr: false,
},
_binding: binding,
user: None,
cidr: None,
cidr_user: None,
@@ -22,7 +52,7 @@ impl TrafficLease {
let mut granted = requested;
let mut user_debit = None;
if let Some(user_bucket) = self.user_bucket.as_ref() {
if let Some(user_bucket) = binding.user_bucket.as_ref() {
let (user_granted, debit) = user_bucket.try_reserve(direction, granted);
user_debit = debit;
if user_granted == 0 {
@@ -33,6 +63,7 @@ impl TrafficLease {
blocked_user: true,
blocked_cidr: false,
},
_binding: binding,
user: user_debit,
cidr: None,
cidr_user: None,
@@ -44,7 +75,7 @@ impl TrafficLease {
let mut cidr_debit = None;
let mut cidr_user_debit = None;
if let (Some(cidr_bucket), Some(cidr_user_share)) =
(self.cidr_bucket.as_ref(), self.cidr_user_share.as_ref())
(binding.cidr_bucket.as_ref(), binding.cidr_user_share.as_ref())
{
let (cidr_granted, aggregate_debit, share_debit) =
cidr_bucket.try_reserve_for_user(direction, cidr_user_share, granted);
@@ -63,6 +94,7 @@ impl TrafficLease {
blocked_user: false,
blocked_cidr: true,
},
_binding: binding,
user: user_debit,
cidr: cidr_debit,
cidr_user: cidr_user_debit,
@@ -77,6 +109,7 @@ impl TrafficLease {
blocked_user: false,
blocked_cidr: false,
},
_binding: binding,
user: user_debit,
cidr: cidr_debit,
cidr_user: cidr_user_debit,
@@ -105,7 +138,7 @@ impl TrafficLease {
}
}
impl TrafficReservation<'_> {
impl TrafficReservation {
/// Returns the shaping decision associated with this reservation.
pub(crate) fn result(&self) -> TrafficConsumeResult {
self.result
@@ -126,7 +159,7 @@ impl TrafficReservation<'_> {
}
}
impl Drop for TrafficLease {
impl Drop for TrafficLeaseBinding {
fn drop(&mut self) {
if let Some(bucket) = self.user_bucket.as_ref() {
decrement_atomic_saturating(&bucket.active_leases, 1);
src/proxy/traffic_limiter/limiter.rs
+24 -7
View File
@@ -6,6 +6,7 @@ impl TrafficLimiter {
Arc::new(Self {
policy: ArcSwap::from_pointee(PolicySnapshot::default()),
policy_update: ParkingMutex::new(()),
published_revision: AtomicU64::new(0),
user_buckets: ShardedRegistry::new(REGISTRY_SHARDS),
cidr_buckets: ShardedRegistry::new(REGISTRY_SHARDS),
user_scope: ScopeMetrics::default(),
@@ -92,6 +93,7 @@ impl TrafficLimiter {
cidr_auto_rules_v6,
cidr_rule_keys,
}));
self.published_revision.store(revision, Ordering::Release);
drop(policy_update);
self.maybe_cleanup();
@@ -102,7 +104,26 @@ impl TrafficLimiter {
user: &str,
client_ip: IpAddr,
) -> Option<Arc<TrafficLease>> {
let policy_update = self.policy_update.lock();
let policy = self.policy.load_full();
let binding = self.build_binding(user, client_ip, &policy);
drop(policy_update);
self.maybe_cleanup();
Some(Arc::new(TrafficLease {
limiter: Arc::clone(self),
user: user.to_string(),
client_ip,
binding: ArcSwap::from(binding),
refresh: ParkingMutex::new(()),
}))
}
pub(super) fn build_binding(
self: &Arc<Self>,
user: &str,
client_ip: IpAddr,
policy: &PolicySnapshot,
) -> Arc<TrafficLeaseBinding> {
let mut user_bucket = None;
if let Some(limit) = policy.user_limits.get(user).copied() {
let bucket = self.user_buckets.get_or_insert_with(
@@ -144,18 +165,14 @@ impl TrafficLimiter {
cidr_bucket = Some(bucket);
}
if user_bucket.is_none() && cidr_bucket.is_none() {
return None;
}
self.maybe_cleanup();
Some(Arc::new(TrafficLease {
Arc::new(TrafficLeaseBinding {
limiter: Arc::clone(self),
revision: policy.revision,
user_bucket,
cidr_bucket,
cidr_user_key,
cidr_user_share,
}))
})
}
pub fn metrics_snapshot(&self) -> TrafficLimiterMetricsSnapshot {
src/proxy/traffic_limiter/tests.rs
+31 -6
View File
@@ -77,7 +77,7 @@ fn auto_cidr_bucket_key_canonicalizes_network_address() {
#[test]
fn refund_from_an_old_epoch_does_not_reduce_the_current_epoch() {
let bucket = DirectionBucket::default();
let bucket = Arc::new(DirectionBucket::default());
let old_debit = bucket.try_reserve_at(7, 100, 80).unwrap();
let current_debit = bucket.try_reserve_at(8, 100, 60).unwrap();
@@ -166,7 +166,7 @@ fn stale_policy_revision_cannot_restore_an_old_rate() {
#[test]
fn dropped_debit_refunds_only_its_packed_epoch() {
let bucket = DirectionBucket::default();
let bucket = Arc::new(DirectionBucket::default());
let debit = bucket.try_reserve_at(11, 100, 80).unwrap();
drop(debit);
@@ -229,9 +229,10 @@ fn dropped_traffic_reservation_refunds_user_and_cidr_debits() {
let epoch = reservation.user.as_ref().unwrap().epoch;
drop(reservation);
let user_bucket = lease.user_bucket.as_ref().unwrap();
let cidr_bucket = lease.cidr_bucket.as_ref().unwrap();
let cidr_user = lease.cidr_user_share.as_ref().unwrap();
let binding = lease.binding.load_full();
let user_bucket = binding.user_bucket.as_ref().unwrap();
let cidr_bucket = binding.cidr_bucket.as_ref().unwrap();
let cidr_user = binding.cidr_user_share.as_ref().unwrap();
assert_eq!(user_bucket.down.used_at(epoch), Some(0));
assert_eq!(cidr_bucket.down.used.used_at(epoch), Some(0));
assert_eq!(cidr_user.down.used.used_at(epoch), Some(0));
@@ -252,7 +253,31 @@ fn partial_traffic_settlement_charges_only_committed_bytes() {
reservation.settle_written(300);
assert_eq!(
lease.user_bucket.as_ref().unwrap().down.used_at(epoch),
lease
.binding
.load_full()
.user_bucket
.as_ref()
.unwrap()
.down
.used_at(epoch),
Some(300)
);
}
#[test]
fn active_lease_observes_policy_removal() {
let limiter = TrafficLimiter::new();
let mut user_limits = HashMap::new();
user_limits.insert("alice".to_string(), rate(1, 0));
limiter.apply_policy(user_limits, HashMap::new());
let lease = limiter
.acquire_lease("alice", "203.0.113.7".parse().unwrap())
.unwrap();
assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 1);
assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 0);
limiter.apply_policy(HashMap::new(), HashMap::new());
assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 1);
}
src/proxy/user_admission.rs
+84 -8
View File
@@ -6,6 +6,7 @@ use parking_lot::{Mutex, MutexGuard};
use tokio_util::sync::CancellationToken;
use crate::crypto::sha256;
use crate::stats::QuotaStore;
const REGISTRATION_PENDING: u8 = 0;
const REGISTRATION_ACTIVE: u8 = 1;
@@ -54,6 +55,8 @@ struct RegisteredOwner {
struct UserAdmissionState {
initialized: bool,
epoch: u64,
active_config_source: Option<u64>,
stale_config_source_rejections: u64,
next_incarnation: UserIncarnation,
next_registration_id: u64,
users: HashMap<String, UserRecord>,
@@ -97,13 +100,20 @@ pub(crate) struct UserMutationResult {
/// Process-owned user authentication and live-owner authority.
pub(crate) struct UserAdmissionAuthority {
state: Mutex<UserAdmissionState>,
quota_store: Arc<QuotaStore>,
}
impl UserAdmissionAuthority {
/// Creates an uninitialized authority for isolated tests and startup wiring.
pub(crate) fn new() -> Arc<Self> {
Self::new_with_quota_store(Arc::new(QuotaStore::default()))
}
/// Creates an authority coupled to the process-scoped quota identity store.
pub(crate) fn new_with_quota_store(quota_store: Arc<QuotaStore>) -> Arc<Self> {
Arc::new(Self {
state: Mutex::new(UserAdmissionState::default()),
quota_store,
})
}
@@ -118,23 +128,41 @@ impl UserAdmissionAuthority {
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Vec<(String, usize)> {
self.apply_config_locked(None, users, user_enabled)
self.activate_config_source(0, 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(
/// Transfers configuration ownership to one runtime generation.
pub(crate) fn activate_config_source(
&self,
expected_epoch: u64,
source_generation: u64,
expected_epoch: Option<u64>,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
self.apply_config_locked(Some(expected_epoch), users, user_enabled)
self.apply_config_locked(source_generation, expected_epoch, true, users, user_enabled)
}
/// Reconciles an update only while its runtime generation owns configuration.
pub(crate) fn apply_config_from_source(
&self,
source_generation: u64,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
self.apply_config_locked(source_generation, None, false, users, user_enabled)
}
/// Returns the number of rejected updates from non-owning generations.
pub(crate) fn stale_config_source_rejections(&self) -> u64 {
self.state.lock().stale_config_source_rejections
}
fn apply_config_locked(
&self,
source_generation: u64,
expected_epoch: Option<u64>,
activate_source: bool,
users: &HashMap<String, String>,
user_enabled: &HashMap<String, bool>,
) -> Option<Vec<(String, usize)>> {
@@ -154,6 +182,21 @@ impl UserAdmissionAuthority {
.collect::<HashMap<_, _>>();
let cancellations = {
let mut state = self.state.lock();
if activate_source {
if state
.active_config_source
.is_some_and(|active| source_generation < active)
{
state.stale_config_source_rejections =
state.stale_config_source_rejections.saturating_add(1);
return None;
}
state.active_config_source = Some(source_generation);
} else if state.active_config_source != Some(source_generation) {
state.stale_config_source_rejections =
state.stale_config_source_rejections.saturating_add(1);
return None;
}
if expected_epoch.is_some_and(|epoch| state.epoch != epoch) {
return None;
}
@@ -195,6 +238,19 @@ impl UserAdmissionAuthority {
if let Some(record) = state.users.get_mut(&user) {
record.incarnation = incarnation;
}
match (old_effective, new_effective) {
(None, Some(_)) => {
self.quota_store.activate_fresh(&user, incarnation);
}
(Some(_), Some(_)) => {
self.quota_store
.advance_preserving_usage(&user, incarnation);
}
(Some(_), None) => {
self.quota_store.retire_through(&user, incarnation);
}
(None, None) => {}
}
}
if identity_changed
|| old_effective.is_some_and(|entry| entry.enabled)
@@ -213,13 +269,14 @@ impl UserAdmissionAuthority {
changed = true;
let incarnation = state.allocate_incarnation();
state.users.insert(
user,
user.clone(),
UserRecord {
configured: Some(desired),
mutation_override: None,
incarnation,
},
);
self.quota_store.activate_fresh(&user, incarnation);
}
if changed {
@@ -238,6 +295,16 @@ impl UserAdmissionAuthority {
enabled: bool,
) -> Option<UserMutationResult> {
let credential_id = credential_id_from_hex(secret)?;
Some(self.stage_user_credential(user, credential_id, enabled))
}
/// Applies one already validated credential mutation ahead of runtime reload.
pub(crate) fn stage_user_credential(
&self,
user: &str,
credential_id: UserCredentialId,
enabled: bool,
) -> UserMutationResult {
let desired = EffectiveUser {
credential_id,
enabled,
@@ -262,6 +329,14 @@ impl UserAdmissionAuthority {
});
record.mutation_override = Some(UserOverride::Present(desired));
record.incarnation = incarnation;
if identity_changed {
if previous.is_some() {
self.quota_store
.advance_preserving_usage(user, incarnation);
} else {
self.quota_store.activate_fresh(user, incarnation);
}
}
state.initialized = true;
state.bump_epoch();
let newly_disabled = previous.is_some_and(|entry| entry.enabled) && !enabled;
@@ -276,11 +351,11 @@ impl UserAdmissionAuthority {
for token in tokens {
token.cancel();
}
Some(UserMutationResult {
UserMutationResult {
incarnation,
cancelled,
newly_disabled,
})
}
}
/// Installs a deletion tombstone and cancels every owner of the old incarnation.
@@ -296,6 +371,7 @@ impl UserAdmissionAuthority {
});
record.mutation_override = Some(UserOverride::Deleted);
record.incarnation = incarnation;
self.quota_store.retire_through(user, incarnation);
state.initialized = true;
state.bump_epoch();
(
src/proxy/user_admission/tests.rs
+66 -2
View File
@@ -43,18 +43,49 @@ fn stale_credential_cannot_cross_delete_and_recreate() {
fn stale_candidate_cannot_overwrite_newer_mutation() {
let authority = UserAdmissionAuthority::new();
let secret = "00112233445566778899aabbccddeeff";
authority.apply_config(&users(secret), &HashMap::new());
authority
.activate_config_source(1, None, &users(secret), &HashMap::new())
.unwrap();
let candidate_epoch = authority.epoch();
authority.stage_user("alice", secret, false).unwrap();
assert!(
authority
.apply_config_if_epoch(candidate_epoch, &users(secret), &HashMap::new())
.activate_config_source(
2,
Some(candidate_epoch),
&users(secret),
&HashMap::new(),
)
.is_none()
);
assert!(!authority.is_user_enabled("alice"));
}
#[test]
fn stale_generation_snapshot_cannot_reopen_reconciled_user() {
let authority = UserAdmissionAuthority::new();
let secret = "00112233445566778899aabbccddeeff";
authority
.activate_config_source(1, None, &users(secret), &HashMap::new())
.unwrap();
authority.stage_user("alice", secret, false).unwrap();
let disabled = HashMap::from([("alice".to_string(), false)]);
let epoch = authority.epoch();
authority
.activate_config_source(2, Some(epoch), &users(secret), &disabled)
.unwrap();
assert!(
authority
.apply_config_from_source(1, &users(secret), &HashMap::new())
.is_none()
);
assert!(!authority.is_user_enabled("alice"));
assert_eq!(authority.stale_config_source_rejections(), 1);
}
#[test]
fn registration_dropped_before_publication_cannot_leave_an_owner() {
let authority = UserAdmissionAuthority::new();
@@ -71,3 +102,36 @@ fn registration_dropped_before_publication_cannot_leave_an_owner() {
assert_eq!(authority.cancel_user_owners("alice"), 0);
}
#[test]
fn quota_identity_follows_credential_rotation_and_recreation() {
let quota_store = Arc::new(QuotaStore::default());
let authority = UserAdmissionAuthority::new_with_quota_store(quota_store.clone());
let old_secret = "00112233445566778899aabbccddeeff";
let new_secret = "ffeeddccbbaa99887766554433221100";
authority.apply_config(&users(old_secret), &HashMap::new());
let old_incarnation = authority
.authenticated_incarnation("alice", credential_id_from_hex(old_secret).unwrap())
.unwrap();
let old_quota = quota_store
.handle_exact("alice", old_incarnation)
.unwrap();
old_quota.charge(40);
let rotated = authority.stage_user("alice", new_secret, true).unwrap();
let rotated_quota = quota_store
.handle_exact("alice", rotated.incarnation)
.unwrap();
old_quota.charge(20);
assert_eq!(rotated_quota.used(), 40);
authority.delete_user("alice");
let recreated = authority.stage_user("alice", old_secret, true).unwrap();
assert_eq!(
quota_store
.handle_exact("alice", recreated.incarnation)
.unwrap()
.used(),
0
);
}