astelm
/
telemt
mirror of https://github.com/telemt/telemt.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(proxy): refactor auth probe failure handling and add concurrent failure tests

This commit is contained in:
David Osipov 2026-03-17 16:25:29 +04:00
parent 0c6bb3a641
commit 93caab1aec
No known key found for this signature in database
GPG Key ID: 0E55C4A47454E82E
4 changed files with 455 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
src/proxy/handshake.rs
View File

@ -11,6 +11,7 @@ use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::time::{Duration, Instant};
use dashmap::DashMap;
use dashmap::mapref::entry::Entry;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use tracing::{debug, warn, trace};
use zeroize::Zeroize;
@ -118,21 +119,30 @@ fn auth_probe_record_failure_with_state(
peer_ip: IpAddr,
now: Instant,
) {
if let Some(mut entry) = state.get_mut(&peer_ip) {
if auth_probe_state_expired(&entry, now) {
*entry = AuthProbeState {
let make_new_state = || AuthProbeState {
fail_streak: 1,
blocked_until: now + auth_probe_backoff(1),
last_seen: now,
};
return;
}
let update_existing = |entry: &mut AuthProbeState| {
if auth_probe_state_expired(entry, now) {
*entry = make_new_state();
} else {
entry.fail_streak = entry.fail_streak.saturating_add(1);
entry.last_seen = now;
entry.blocked_until = now + auth_probe_backoff(entry.fail_streak);
return;
}
};
match state.entry(peer_ip) {
Entry::Occupied(mut entry) => {
update_existing(entry.get_mut());
return;
}
Entry::Vacant(_) => {}
}
if state.len() >= AUTH_PROBE_TRACK_MAX_ENTRIES {
let mut stale_keys = Vec::new();
let mut eviction_candidates = Vec::new();
@ -155,11 +165,14 @@ fn auth_probe_record_failure_with_state(
}
}
state.insert(peer_ip, AuthProbeState {
fail_streak: 1,
blocked_until: now + auth_probe_backoff(1),
last_seen: now,
});
match state.entry(peer_ip) {
Entry::Occupied(mut entry) => {
update_existing(entry.get_mut());
}
Entry::Vacant(entry) => {
entry.insert(make_new_state());
}
}
}
fn auth_probe_record_success(peer_ip: IpAddr) {

114
src/proxy/handshake_security_tests.rs
View File

@ -4,6 +4,7 @@ use dashmap::DashMap;
use std::net::{IpAddr, Ipv4Addr};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Barrier;
fn make_valid_tls_handshake(secret: &[u8], timestamp: u32) -> Vec<u8> {
let session_id_len: usize = 32;
@ -994,3 +995,116 @@ fn auth_probe_eviction_offset_varies_with_input() {
assert_eq!(a, b, "same input must yield deterministic offset");
assert_ne!(a, c, "different peer IPs should not collapse to one offset");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn auth_probe_concurrent_failures_do_not_lose_fail_streak_updates() {
let _guard = auth_probe_test_lock()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
clear_auth_probe_state_for_testing();
let peer_ip: IpAddr = "198.51.100.90".parse().unwrap();
let tasks = 128usize;
let barrier = Arc::new(Barrier::new(tasks));
let mut handles = Vec::with_capacity(tasks);
for _ in 0..tasks {
let barrier = barrier.clone();
handles.push(tokio::spawn(async move {
barrier.wait().await;
auth_probe_record_failure(peer_ip, Instant::now());
}));
}
for handle in handles {
handle
.await
.expect("concurrent failure recording task must not panic");
}
let streak = auth_probe_fail_streak_for_testing(peer_ip)
.expect("tracked peer must exist after concurrent failure burst");
assert_eq!(
streak as usize,
tasks,
"concurrent failures for one source must account every attempt"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn invalid_probe_noise_from_other_ips_does_not_break_valid_tls_handshake() {
let _guard = auth_probe_test_lock()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
clear_auth_probe_state_for_testing();
let secret = [0x31u8; 16];
let config = Arc::new(test_config_with_secret_hex("31313131313131313131313131313131"));
let replay_checker = Arc::new(ReplayChecker::new(4096, Duration::from_secs(60)));
let rng = Arc::new(SecureRandom::new());
let victim_peer: SocketAddr = "198.51.100.91:44391".parse().unwrap();
let valid = Arc::new(make_valid_tls_handshake(&secret, 0));
let mut invalid = vec![0x42u8; tls::TLS_DIGEST_POS + tls::TLS_DIGEST_LEN + 1 + 32];
invalid[tls::TLS_DIGEST_POS + tls::TLS_DIGEST_LEN] = 32;
let invalid = Arc::new(invalid);
let mut noise_tasks = Vec::new();
for idx in 0..96u16 {
let config = config.clone();
let replay_checker = replay_checker.clone();
let rng = rng.clone();
let invalid = invalid.clone();
noise_tasks.push(tokio::spawn(async move {
let octet = ((idx % 200) + 1) as u8;
let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(203, 0, 113, octet)), 45000 + idx);
let result = handle_tls_handshake(
&invalid,
tokio::io::empty(),
tokio::io::sink(),
peer,
&config,
&replay_checker,
&rng,
None,
)
.await;
assert!(matches!(result, HandshakeResult::BadClient { .. }));
}));
}
let victim_config = config.clone();
let victim_replay_checker = replay_checker.clone();
let victim_rng = rng.clone();
let victim_valid = valid.clone();
let victim_task = tokio::spawn(async move {
handle_tls_handshake(
&victim_valid,
tokio::io::empty(),
tokio::io::sink(),
victim_peer,
&victim_config,
&victim_replay_checker,
&victim_rng,
None,
)
.await
});
for task in noise_tasks {
task.await.expect("noise task must not panic");
}
let victim_result = victim_task
.await
.expect("victim handshake task must not panic");
assert!(
matches!(victim_result, HandshakeResult::Success(_)),
"invalid probe noise from other IPs must not block a valid victim handshake"
);
assert_eq!(
auth_probe_fail_streak_for_testing(victim_peer.ip()),
None,
"successful victim handshake must not retain pre-auth failure streak"
);
}

10
src/proxy/middle_relay.rs
View File

@ -7,7 +7,6 @@ use std::time::{Duration, Instant};
#[cfg(test)]
use std::sync::Mutex;
use bytes::Bytes;
use dashmap::DashMap;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::{mpsc, oneshot, watch};
@ -24,11 +23,11 @@ use crate::proxy::route_mode::{
cutover_stagger_delay,
};
use crate::stats::Stats;
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer};
use crate::transport::middle_proxy::{MePool, MeResponse, proto_flags_for_tag};
enum C2MeCommand {
Data { payload: Bytes, flags: u32 },
Data { payload: PooledBuffer, flags: u32 },
Close,
}
@ -686,7 +685,7 @@ async fn read_client_payload<R>(
forensics: &RelayForensicsState,
frame_counter: &mut u64,
stats: &Stats,
) -> Result<Option<(Bytes, bool)>>
) -> Result<Option<(PooledBuffer, bool)>>
where
R: AsyncRead + Unpin + Send + 'static,
{
@ -807,8 +806,7 @@ where
payload.truncate(secure_payload_len);
}
*frame_counter += 1;
let payload_bytes = Bytes::copy_from_slice(&payload[..]);
return Ok(Some((payload_bytes, quickack)));
return Ok(Some((payload, quickack)));
}
}

311
src/proxy/middle_relay_security_tests.rs
View File

@ -1,7 +1,9 @@
use super::*;
use bytes::Bytes;
use crate::crypto::AesCtr;
use crate::crypto::SecureRandom;
use crate::stats::Stats;
use crate::stream::{BufferPool, CryptoReader};
use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer};
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
@ -9,6 +11,21 @@ use tokio::io::AsyncWriteExt;
use tokio::io::duplex;
use tokio::time::{Duration as TokioDuration, timeout};
fn make_pooled_payload(data: &[u8]) -> PooledBuffer {
let pool = Arc::new(BufferPool::with_config(data.len().max(1), 4));
let mut payload = pool.get();
payload.resize(data.len(), 0);
payload[..data.len()].copy_from_slice(data);
payload
}
fn make_pooled_payload_from(pool: &Arc<BufferPool>, data: &[u8]) -> PooledBuffer {
let mut payload = pool.get();
payload.resize(data.len(), 0);
payload[..data.len()].copy_from_slice(data);
payload
}
#[test]
fn should_yield_sender_only_on_budget_with_backlog() {
assert!(!should_yield_c2me_sender(0, true));
@ -23,7 +40,7 @@ async fn enqueue_c2me_command_uses_try_send_fast_path() {
enqueue_c2me_command(
&tx,
C2MeCommand::Data {
payload: Bytes::from_static(&[1, 2, 3]),
payload: make_pooled_payload(&[1, 2, 3]),
flags: 0,
},
)
@ -47,7 +64,7 @@ async fn enqueue_c2me_command_uses_try_send_fast_path() {
async fn enqueue_c2me_command_falls_back_to_send_when_queue_is_full() {
let (tx, mut rx) = mpsc::channel::<C2MeCommand>(1);
tx.send(C2MeCommand::Data {
payload: Bytes::from_static(&[9]),
payload: make_pooled_payload(&[9]),
flags: 9,
})
.await
@ -58,7 +75,7 @@ async fn enqueue_c2me_command_falls_back_to_send_when_queue_is_full() {
enqueue_c2me_command(
&tx2,
C2MeCommand::Data {
payload: Bytes::from_static(&[7, 7]),
payload: make_pooled_payload(&[7, 7]),
flags: 7,
},
)
@ -84,6 +101,74 @@ async fn enqueue_c2me_command_falls_back_to_send_when_queue_is_full() {
}
}
#[tokio::test]
async fn enqueue_c2me_command_closed_channel_recycles_payload() {
let pool = Arc::new(BufferPool::with_config(64, 4));
let payload = make_pooled_payload_from(&pool, &[1, 2, 3, 4]);
let (tx, rx) = mpsc::channel::<C2MeCommand>(1);
drop(rx);
let result = enqueue_c2me_command(
&tx,
C2MeCommand::Data {
payload,
flags: 0,
},
)
.await;
assert!(result.is_err(), "closed queue must fail enqueue");
drop(result);
assert!(
pool.stats().pooled >= 1,
"payload must return to pool when enqueue fails on closed channel"
);
}
#[tokio::test]
async fn enqueue_c2me_command_full_then_closed_recycles_waiting_payload() {
let pool = Arc::new(BufferPool::with_config(64, 4));
let (tx, rx) = mpsc::channel::<C2MeCommand>(1);
tx.send(C2MeCommand::Data {
payload: make_pooled_payload_from(&pool, &[9]),
flags: 1,
})
.await
.unwrap();
let tx2 = tx.clone();
let pool2 = pool.clone();
let blocked_send = tokio::spawn(async move {
enqueue_c2me_command(
&tx2,
C2MeCommand::Data {
payload: make_pooled_payload_from(&pool2, &[7, 7, 7]),
flags: 2,
},
)
.await
});
tokio::time::sleep(TokioDuration::from_millis(10)).await;
drop(rx);
let result = timeout(TokioDuration::from_secs(1), blocked_send)
.await
.expect("blocked send task must finish")
.expect("blocked send task must not panic");
assert!(
result.is_err(),
"closing receiver while sender is blocked must fail enqueue"
);
drop(result);
assert!(
pool.stats().pooled >= 2,
"both queued and blocked payloads must return to pool after channel close"
);
}
#[test]
fn desync_dedup_cache_is_bounded() {
let _guard = desync_dedup_test_lock()
@ -150,6 +235,12 @@ fn make_crypto_reader(reader: tokio::io::DuplexStream) -> CryptoReader<tokio::io
CryptoReader::new(reader, AesCtr::new(&key, iv))
}
fn make_crypto_writer(writer: tokio::io::DuplexStream) -> CryptoWriter<tokio::io::DuplexStream> {
let key = [0u8; 32];
let iv = 0u128;
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
}
fn encrypt_for_reader(plaintext: &[u8]) -> Vec<u8> {
let key = [0u8; 32];
let iv = 0u128;
@ -476,3 +567,215 @@ async fn read_client_payload_returns_buffer_to_pool_after_emit() {
"emitted payload buffer must be returned to pool to avoid pool drain"
);
}
#[tokio::test]
async fn read_client_payload_keeps_pool_buffer_checked_out_until_frame_drop() {
let _guard = desync_dedup_test_lock()
.lock()
.expect("middle relay test lock must be available");
let pool = Arc::new(BufferPool::with_config(64, 2));
pool.preallocate(1);
assert_eq!(pool.stats().pooled, 1, "one pooled buffer must be available");
let (reader, mut writer) = duplex(1024);
let mut crypto_reader = make_crypto_reader(reader);
let stats = Stats::new();
let forensics = make_forensics_state();
let mut frame_counter = 0;
let payload = [0x41u8, 0x42, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48];
let mut plaintext = Vec::with_capacity(4 + payload.len());
plaintext.extend_from_slice(&(payload.len() as u32).to_le_bytes());
plaintext.extend_from_slice(&payload);
let encrypted = encrypt_for_reader(&plaintext);
writer.write_all(&encrypted).await.unwrap();
let (frame, quickack) = read_client_payload(
&mut crypto_reader,
ProtoTag::Intermediate,
1024,
TokioDuration::from_secs(1),
&pool,
&forensics,
&mut frame_counter,
&stats,
)
.await
.expect("payload read must succeed")
.expect("frame must be present");
assert!(!quickack);
assert_eq!(frame.as_ref(), &payload);
assert_eq!(
pool.stats().pooled,
0,
"buffer must stay checked out while frame payload is alive"
);
drop(frame);
assert!(
pool.stats().pooled >= 1,
"buffer must return to pool only after frame drop"
);
}
#[tokio::test]
async fn enqueue_c2me_close_unblocks_after_queue_drain() {
let (tx, mut rx) = mpsc::channel::<C2MeCommand>(1);
tx.send(C2MeCommand::Data {
payload: make_pooled_payload(&[0x41]),
flags: 0,
})
.await
.unwrap();
let tx2 = tx.clone();
let close_task = tokio::spawn(async move { enqueue_c2me_command(&tx2, C2MeCommand::Close).await });
tokio::time::sleep(TokioDuration::from_millis(10)).await;
let first = timeout(TokioDuration::from_millis(100), rx.recv())
.await
.unwrap()
.expect("first queued item must be present");
assert!(matches!(first, C2MeCommand::Data { .. }));
close_task.await.unwrap().expect("close enqueue must succeed after drain");
let second = timeout(TokioDuration::from_millis(100), rx.recv())
.await
.unwrap()
.expect("close command must follow after queue drain");
assert!(matches!(second, C2MeCommand::Close));
}
#[tokio::test]
async fn enqueue_c2me_close_full_then_receiver_drop_fails_cleanly() {
let (tx, rx) = mpsc::channel::<C2MeCommand>(1);
tx.send(C2MeCommand::Data {
payload: make_pooled_payload(&[0x42]),
flags: 0,
})
.await
.unwrap();
let tx2 = tx.clone();
let close_task = tokio::spawn(async move { enqueue_c2me_command(&tx2, C2MeCommand::Close).await });
tokio::time::sleep(TokioDuration::from_millis(10)).await;
drop(rx);
let result = timeout(TokioDuration::from_secs(1), close_task)
.await
.expect("close task must finish")
.expect("close task must not panic");
assert!(
result.is_err(),
"close enqueue must fail cleanly when receiver is dropped under pressure"
);
}
#[tokio::test]
async fn process_me_writer_response_ack_obeys_flush_policy() {
let (writer_side, _reader_side) = duplex(1024);
let mut writer = make_crypto_writer(writer_side);
let rng = SecureRandom::new();
let mut frame_buf = Vec::new();
let stats = Stats::new();
let bytes_me2c = AtomicU64::new(0);
let immediate = process_me_writer_response(
MeResponse::Ack(0x11223344),
&mut writer,
ProtoTag::Intermediate,
&rng,
&mut frame_buf,
&stats,
"user",
&bytes_me2c,
77,
true,
false,
)
.await
.expect("ack response must be processed");
assert!(matches!(
immediate,
MeWriterResponseOutcome::Continue {
frames: 1,
bytes: 4,
flush_immediately: true,
}
));
let delayed = process_me_writer_response(
MeResponse::Ack(0x55667788),
&mut writer,
ProtoTag::Intermediate,
&rng,
&mut frame_buf,
&stats,
"user",
&bytes_me2c,
77,
false,
false,
)
.await
.expect("ack response must be processed");
assert!(matches!(
delayed,
MeWriterResponseOutcome::Continue {
frames: 1,
bytes: 4,
flush_immediately: false,
}
));
}
#[tokio::test]
async fn process_me_writer_response_data_updates_byte_accounting() {
let (writer_side, _reader_side) = duplex(1024);
let mut writer = make_crypto_writer(writer_side);
let rng = SecureRandom::new();
let mut frame_buf = Vec::new();
let stats = Stats::new();
let bytes_me2c = AtomicU64::new(0);
let payload = vec![1u8, 2, 3, 4, 5, 6, 7, 8, 9];
let outcome = process_me_writer_response(
MeResponse::Data {
flags: 0,
data: Bytes::from(payload.clone()),
},
&mut writer,
ProtoTag::Intermediate,
&rng,
&mut frame_buf,
&stats,
"user",
&bytes_me2c,
88,
false,
false,
)
.await
.expect("data response must be processed");
assert!(matches!(
outcome,
MeWriterResponseOutcome::Continue {
frames: 1,
bytes,
flush_immediately: false,
} if bytes == payload.len()
));
assert_eq!(
bytes_me2c.load(std::sync::atomic::Ordering::Relaxed),
payload.len() as u64,
"ME->C byte accounting must increase by emitted payload size"
);
}