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

Slot Budget + Config Store Atomic Writer fixes + Trusted Command

Browse Source
This commit is contained in:
Alexey
2026-09-16 22:06:59 +03:00
parent 55f3d19ee0
commit 9a683d8b3d
33 changed files with 1680 additions and 687 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/proxy/handshake/auth_candidates.rs
+94 -18
View File
@@ -76,27 +76,48 @@ pub(super) fn sticky_hint_record_success_in(
user_id: u32,
sni: Option<&str>,
) {
if shared.handshake.sticky_user_by_ip.len() > STICKY_HINT_MAX_ENTRIES {
shared.handshake.sticky_user_by_ip.clear();
}
shared.handshake.sticky_user_by_ip.insert(peer_ip, user_id);
if shared.handshake.sticky_user_by_ip_prefix.len() > STICKY_HINT_MAX_ENTRIES {
shared.handshake.sticky_user_by_ip_prefix.clear();
}
shared
.handshake
.sticky_user_by_ip_prefix
.insert(ip_prefix_hint_key(peer_ip), user_id);
bounded_sticky_hint_upsert(
&shared.handshake.sticky_user_by_ip,
&shared.handshake.sticky_user_by_ip_slots,
peer_ip,
user_id,
);
bounded_sticky_hint_upsert(
&shared.handshake.sticky_user_by_ip_prefix,
&shared.handshake.sticky_user_by_ip_prefix_slots,
ip_prefix_hint_key(peer_ip),
user_id,
);
if let Some(sni) = sni {
if shared.handshake.sticky_user_by_sni_hash.len() > STICKY_HINT_MAX_ENTRIES {
shared.handshake.sticky_user_by_sni_hash.clear();
bounded_sticky_hint_upsert(
&shared.handshake.sticky_user_by_sni_hash,
&shared.handshake.sticky_user_by_sni_hash_slots,
sni_hint_hash(sni),
user_id,
);
}
}
fn bounded_sticky_hint_upsert<K>(
entries: &DashMap<K, u32>,
slots: &crate::slot_budget::SlotBudget,
key: K,
user_id: u32,
) where
K: Eq + Hash,
{
match entries.entry(key) {
Entry::Occupied(mut entry) => {
entry.insert(user_id);
}
Entry::Vacant(entry) => {
let Some(slot) = slots.try_acquire() else {
return;
};
entry.insert(user_id);
slot.commit();
}
shared
.handshake
.sticky_user_by_sni_hash
.insert(sni_hint_hash(sni), user_id);
}
}
@@ -343,6 +364,61 @@ mod web_mode_tests {
}
}
#[cfg(test)]
mod bounded_registry_tests {
use std::sync::Arc;
use super::*;
#[test]
fn parallel_sticky_hints_never_exceed_their_hard_caps() {
const ATTEMPTS: usize = 10_000;
let shared = ProxySharedState::new();
std::thread::scope(|scope| {
for worker in 0..16 {
let shared = Arc::clone(&shared);
scope.spawn(move || {
for index in (worker..ATTEMPTS).step_by(16) {
let octets = (index as u32).to_be_bytes();
let peer_ip = IpAddr::V4(std::net::Ipv4Addr::new(
octets[1], octets[2], octets[3], worker as u8,
));
sticky_hint_record_success_in(
shared.as_ref(),
peer_ip,
index as u32,
Some(&format!("host-{index}.example")),
);
}
});
}
});
assert_eq!(shared.handshake.sticky_user_by_ip.len(), STICKY_HINT_MAX_ENTRIES);
assert_eq!(
shared.handshake.sticky_user_by_ip_prefix.len(),
STICKY_HINT_MAX_ENTRIES
);
assert_eq!(
shared.handshake.sticky_user_by_sni_hash.len(),
STICKY_HINT_MAX_ENTRIES
);
assert_eq!(
shared.handshake.sticky_user_by_ip_slots.used(),
shared.handshake.sticky_user_by_ip.len()
);
assert_eq!(
shared.handshake.sticky_user_by_ip_prefix_slots.used(),
shared.handshake.sticky_user_by_ip_prefix.len()
);
assert_eq!(
shared.handshake.sticky_user_by_sni_hash_slots.used(),
shared.handshake.sticky_user_by_sni_hash.len()
);
}
}
pub(super) fn decode_user_secrets_in(
shared: &ProxySharedState,
config: &ProxyConfig,
src/proxy/handshake/auth_probe.rs
+71 -127
View File
@@ -98,9 +98,14 @@ pub(super) fn auth_probe_is_throttled_in(
};
if auth_probe_state_expired(&entry, now) {
drop(entry);
state.remove_if(&peer_ip, |_, current| {
auth_probe_state_expired(current, now)
});
if state
.remove_if(&peer_ip, |_, current| {
auth_probe_state_expired(current, now)
})
.is_some()
{
shared.handshake.auth_probe_slots.release();
}
return false;
}
now < entry.blocked_until
@@ -118,9 +123,14 @@ pub(super) fn auth_probe_saturation_grace_exhausted_in(
};
if auth_probe_state_expired(&entry, now) {
drop(entry);
state.remove_if(&peer_ip, |_, current| {
auth_probe_state_expired(current, now)
});
if state
.remove_if(&peer_ip, |_, current| {
auth_probe_state_expired(current, now)
})
.is_some()
{
shared.handshake.auth_probe_slots.release();
}
return false;
}
@@ -216,7 +226,13 @@ pub(super) fn auth_probe_record_failure_in(
) {
let peer_ip = normalize_auth_probe_ip(peer_ip);
let state = &shared.handshake.auth_probe;
auth_probe_record_failure_with_state_in(shared, state, peer_ip, now);
auth_probe_record_failure_with_state_and_budget_in(
shared,
state,
Some(&shared.handshake.auth_probe_slots),
peer_ip,
now,
);
}
pub(super) fn auth_probe_record_failure_with_state_in(
@@ -224,6 +240,16 @@ pub(super) fn auth_probe_record_failure_with_state_in(
state: &DashMap<IpAddr, AuthProbeState>,
peer_ip: IpAddr,
now: Instant,
) {
auth_probe_record_failure_with_state_and_budget_in(shared, state, None, peer_ip, now);
}
fn auth_probe_record_failure_with_state_and_budget_in(
shared: &ProxySharedState,
state: &DashMap<IpAddr, AuthProbeState>,
slots: Option<&crate::slot_budget::SlotBudget>,
peer_ip: IpAddr,
now: Instant,
) {
let make_new_state = || AuthProbeState {
fail_streak: 1,
@@ -279,6 +305,9 @@ pub(super) fn auth_probe_record_failure_with_state_in(
})
.is_some()
{
if let Some(slots) = slots {
slots.release();
}
break;
}
continue;
@@ -347,9 +376,15 @@ pub(super) fn auth_probe_record_failure_with_state_in(
}
for stale_key in stale_keys {
state.remove_if(&stale_key, |_, current| {
auth_probe_state_expired(current, now)
});
if state
.remove_if(&stale_key, |_, current| {
auth_probe_state_expired(current, now)
})
.is_some()
&& let Some(slots) = slots
{
slots.release();
}
}
if state.len() < AUTH_PROBE_TRACK_MAX_ENTRIES {
@@ -360,19 +395,38 @@ pub(super) fn auth_probe_record_failure_with_state_in(
auth_probe_note_saturation_in(shared, now);
return;
};
state.remove_if(&evict_key, |_, current| {
current.fail_streak == evict_fail_streak && current.last_seen == evict_last_seen
});
if state
.remove_if(&evict_key, |_, current| {
current.fail_streak == evict_fail_streak
&& current.last_seen == evict_last_seen
})
.is_some()
&& let Some(slots) = slots
{
slots.release();
}
auth_probe_note_saturation_in(shared, now);
}
}
let slot = if let Some(slots) = slots {
let Some(slot) = slots.try_acquire() else {
auth_probe_note_saturation_in(shared, now);
return;
};
Some(slot)
} else {
None
};
match state.entry(peer_ip) {
Entry::Occupied(mut entry) => {
update_existing(entry.get_mut());
}
Entry::Vacant(entry) => {
entry.insert(make_new_state());
if let Some(slot) = slot {
slot.commit();
}
}
}
}
@@ -380,125 +434,15 @@ pub(super) fn auth_probe_record_failure_with_state_in(
pub(super) fn auth_probe_record_success_in(shared: &ProxySharedState, peer_ip: IpAddr) {
let peer_ip = normalize_auth_probe_ip(peer_ip);
let state = &shared.handshake.auth_probe;
state.remove(&peer_ip);
}
#[cfg(test)]
pub(crate) fn auth_probe_record_failure_for_testing(
shared: &ProxySharedState,
peer_ip: IpAddr,
now: Instant,
) {
auth_probe_record_failure_in(shared, peer_ip, now);
}
#[cfg(test)]
pub(crate) fn auth_probe_fail_streak_for_testing_in_shared(
shared: &ProxySharedState,
peer_ip: IpAddr,
) -> Option<u32> {
let peer_ip = normalize_auth_probe_ip(peer_ip);
shared
.handshake
.auth_probe
.get(&peer_ip)
.map(|entry| entry.fail_streak)
}
#[cfg(test)]
pub(crate) fn clear_auth_probe_state_for_testing_in_shared(shared: &ProxySharedState) {
shared.handshake.auth_probe.clear();
match shared.handshake.auth_probe_saturation.lock() {
Ok(mut saturation) => {
*saturation = None;
}
Err(poisoned) => {
let mut saturation = poisoned.into_inner();
*saturation = None;
shared.handshake.auth_probe_saturation.clear_poison();
}
if state.remove(&peer_ip).is_some() {
shared.handshake.auth_probe_slots.release();
}
}
#[cfg(test)]
pub(crate) fn auth_probe_state_for_testing_in_shared(
shared: &ProxySharedState,
) -> &DashMap<IpAddr, AuthProbeState> {
&shared.handshake.auth_probe
}
mod testing;
#[cfg(test)]
pub(crate) fn auth_probe_saturation_state_for_testing_in_shared(
shared: &ProxySharedState,
) -> &Mutex<Option<AuthProbeSaturationState>> {
&shared.handshake.auth_probe_saturation
}
#[cfg(test)]
pub(crate) fn auth_probe_saturation_state_lock_for_testing_in_shared(
shared: &ProxySharedState,
) -> std::sync::MutexGuard<'_, Option<AuthProbeSaturationState>> {
shared
.handshake
.auth_probe_saturation
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
#[cfg(test)]
pub(crate) fn clear_unknown_sni_warn_state_for_testing_in_shared(shared: &ProxySharedState) {
let mut guard = shared
.handshake
.unknown_sni_warn_next_allowed
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
*guard = None;
}
#[cfg(test)]
pub(crate) fn should_emit_unknown_sni_warn_for_testing_in_shared(
shared: &ProxySharedState,
now: Instant,
) -> bool {
should_emit_unknown_sni_warn_in(shared, now)
}
#[cfg(test)]
pub(crate) fn clear_warned_secrets_for_testing_in_shared(shared: &ProxySharedState) {
if let Ok(mut guard) = shared.handshake.invalid_secret_warned.lock() {
guard.clear();
}
}
#[cfg(test)]
pub(crate) fn warned_secrets_for_testing_in_shared(
shared: &ProxySharedState,
) -> &Mutex<HashSet<(String, String)>> {
&shared.handshake.invalid_secret_warned
}
#[cfg(test)]
pub(crate) fn auth_probe_is_throttled_for_testing_in_shared(
shared: &ProxySharedState,
peer_ip: IpAddr,
) -> bool {
auth_probe_is_throttled_in(shared, peer_ip, Instant::now())
}
#[cfg(test)]
pub(crate) fn auth_probe_saturation_is_throttled_for_testing_in_shared(
shared: &ProxySharedState,
) -> bool {
auth_probe_saturation_is_throttled_in(shared, Instant::now())
}
#[cfg(test)]
pub(crate) fn auth_probe_saturation_is_throttled_at_for_testing_in_shared(
shared: &ProxySharedState,
now: Instant,
) -> bool {
auth_probe_saturation_is_throttled_in(shared, now)
}
pub(crate) use testing::*;
#[inline]
pub(super) fn find_matching_tls_domain<'a>(config: &'a ProxyConfig, sni: &str) -> Option<&'a str> {
src/proxy/handshake/auth_probe/testing.rs
+137
View File
@@ -0,0 +1,137 @@
use super::*;
pub(crate) fn auth_probe_record_failure_for_testing(
shared: &ProxySharedState,
peer_ip: IpAddr,
now: Instant,
) {
auth_probe_record_failure_in(shared, peer_ip, now);
}
pub(crate) fn auth_probe_fail_streak_for_testing_in_shared(
shared: &ProxySharedState,
peer_ip: IpAddr,
) -> Option<u32> {
let peer_ip = normalize_auth_probe_ip(peer_ip);
shared
.handshake
.auth_probe
.get(&peer_ip)
.map(|entry| entry.fail_streak)
}
pub(crate) fn clear_auth_probe_state_for_testing_in_shared(shared: &ProxySharedState) {
shared.handshake.auth_probe.clear();
shared.handshake.auth_probe_slots.reset_for_testing();
match shared.handshake.auth_probe_saturation.lock() {
Ok(mut saturation) => {
*saturation = None;
}
Err(poisoned) => {
let mut saturation = poisoned.into_inner();
*saturation = None;
shared.handshake.auth_probe_saturation.clear_poison();
}
}
}
pub(crate) fn auth_probe_state_for_testing_in_shared(
shared: &ProxySharedState,
) -> &DashMap<IpAddr, AuthProbeState> {
&shared.handshake.auth_probe
}
pub(crate) fn auth_probe_slots_for_testing_in_shared(shared: &ProxySharedState) -> usize {
shared.handshake.auth_probe_slots.used()
}
pub(crate) fn auth_probe_saturation_state_for_testing_in_shared(
shared: &ProxySharedState,
) -> &Mutex<Option<AuthProbeSaturationState>> {
&shared.handshake.auth_probe_saturation
}
pub(crate) fn auth_probe_saturation_state_lock_for_testing_in_shared(
shared: &ProxySharedState,
) -> std::sync::MutexGuard<'_, Option<AuthProbeSaturationState>> {
shared
.handshake
.auth_probe_saturation
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
pub(crate) fn clear_unknown_sni_warn_state_for_testing_in_shared(shared: &ProxySharedState) {
let mut guard = shared
.handshake
.unknown_sni_warn_next_allowed
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
*guard = None;
}
pub(crate) fn should_emit_unknown_sni_warn_for_testing_in_shared(
shared: &ProxySharedState,
now: Instant,
) -> bool {
should_emit_unknown_sni_warn_in(shared, now)
}
pub(crate) fn clear_warned_secrets_for_testing_in_shared(shared: &ProxySharedState) {
if let Ok(mut guard) = shared.handshake.invalid_secret_warned.lock() {
guard.clear();
}
}
pub(crate) fn warned_secrets_for_testing_in_shared(
shared: &ProxySharedState,
) -> &Mutex<HashSet<(String, String)>> {
&shared.handshake.invalid_secret_warned
}
pub(crate) fn auth_probe_is_throttled_for_testing_in_shared(
shared: &ProxySharedState,
peer_ip: IpAddr,
) -> bool {
auth_probe_is_throttled_in(shared, peer_ip, Instant::now())
}
pub(crate) fn auth_probe_saturation_is_throttled_for_testing_in_shared(
shared: &ProxySharedState,
) -> bool {
auth_probe_saturation_is_throttled_in(shared, Instant::now())
}
pub(crate) fn auth_probe_saturation_is_throttled_at_for_testing_in_shared(
shared: &ProxySharedState,
now: Instant,
) -> bool {
auth_probe_saturation_is_throttled_in(shared, now)
}
#[test]
fn parallel_distinct_failures_respect_exact_auth_probe_capacity() {
const ATTEMPTS: usize = 10_000;
let shared = ProxySharedState::new();
std::thread::scope(|scope| {
for worker in 0..16 {
let shared = Arc::clone(&shared);
scope.spawn(move || {
for index in (worker..ATTEMPTS).step_by(16) {
let octets = (index as u32).to_be_bytes();
let peer_ip = IpAddr::V4(std::net::Ipv4Addr::new(
octets[1], octets[2], octets[3], worker as u8,
));
auth_probe_record_failure_in(shared.as_ref(), peer_ip, Instant::now());
}
});
}
});
assert_eq!(shared.handshake.auth_probe.len(), AUTH_PROBE_TRACK_MAX_ENTRIES);
assert_eq!(
auth_probe_slots_for_testing_in_shared(shared.as_ref()),
AUTH_PROBE_TRACK_MAX_ENTRIES
);
}