checkpoint: integrate authenticated transports

This commit is contained in:
h3cx committed 2026-09-15 10:57:11 +02:00
1 parent 313ab00765
commit 586ff8f223
42 files changed
+6277 -771

No files matched your search

+322 -13
View File
@@ -1,13 +1,4 @@
#[path = "../src/config.rs"]
mod config;
#[path = "../src/datagram.rs"]
mod datagram;
#[path = "../src/protocol/mod.rs"]
mod protocol;
#[path = "../src/stats.rs"]
mod stats;
#[path = "../src/transport/mod.rs"]
mod transport;
use wgcloak::{config, protocol, transport};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -16,7 +7,7 @@ use bytes::Bytes;
use config::Config;
use tokio::{net::UdpSocket, time::timeout};
use transport::{
direct_udp::{DirectUdpServer, DirectUdpSession},
direct_udp::{DirectUdpServer, DirectUdpServerHandle, DirectUdpSession},
TransportError, TransportSession,
};
@@ -51,6 +42,19 @@ fn multi_configs(
(server, client("one"), client("two"))
}
async fn wait_for_active_sessions(server: &DirectUdpServerHandle, expected: usize) {
timeout(Duration::from_secs(2), async {
loop {
if server.active_sessions().await == expected {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap_or_else(|_| panic!("active session count did not become {expected}"));
}
#[tokio::test]
async fn direct_udp_round_trip_and_auth_rejection() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
@@ -66,6 +70,11 @@ async fn direct_udp_round_trip_and_auth_rejection() {
backend.send_to(&buf[..len], peer).await.unwrap();
});
let session = DirectUdpSession::connect(&client_config).await.unwrap();
assert_eq!(
session.authenticated_keepalive_interval(),
Some(Duration::from_secs(15))
);
session.send_keepalive().await.unwrap();
session
.send_datagram(Bytes::from_static(b"opaque"))
.await
@@ -99,19 +108,42 @@ async fn direct_udp_enforces_payload_limit_and_expires_idle_session() {
let (mut server_config, client_config) = configs(45782, backend_addr.port());
server_config.session.idle_timeout_secs = 1;
let server = DirectUdpServer::bind(server_config).await.unwrap();
let stats = server.stats();
let observer = server.handle();
let stats = observer.stats();
let server_task = tokio::spawn(server.run());
let session = DirectUdpSession::connect(&client_config).await.unwrap();
wait_for_active_sessions(&observer, 1).await;
assert!(matches!(
session.send_datagram(Bytes::from(vec![0u8; 1201])).await,
Err(TransportError::Oversized { maximum: 1200 })
));
tokio::time::sleep(Duration::from_secs(3)).await;
assert_eq!(stats.snapshot().expired_sessions, 1);
wait_for_active_sessions(&observer, 0).await;
session.close().await.unwrap();
server_task.abort();
}
#[tokio::test]
async fn direct_udp_observer_tracks_authenticated_session_close() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let backend_addr = backend.local_addr().unwrap();
let (server_config, client_config) = configs(45787, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let observer = server.handle();
let stats = observer.stats();
let server_task = tokio::spawn(server.run());
let session = DirectUdpSession::connect(&client_config).await.unwrap();
wait_for_active_sessions(&observer, 1).await;
assert_eq!(stats.snapshot().expired_sessions, 0);
session.close().await.unwrap();
wait_for_active_sessions(&observer, 0).await;
assert_eq!(stats.snapshot().expired_sessions, 0);
server_task.abort();
}
#[tokio::test]
async fn concurrent_clients_use_independent_backend_sockets() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
@@ -178,6 +210,14 @@ async fn malformed_and_unauthenticated_packets_do_not_kill_server() {
let raw = UdpSocket::bind("127.0.0.1:0").await.unwrap();
raw.send_to(&[1, 2, 3], "127.0.0.1:45784").await.unwrap();
raw.send_to(&[0u8; 64], "127.0.0.1:45784").await.unwrap();
let oversized = protocol::DataPacket::new(1, 1, vec![0u8; 1201], [0; 16])
.encode()
.unwrap();
raw.send_to(&oversized, "127.0.0.1:45784").await.unwrap();
let unexpected = protocol::RejectPacket::new(protocol::RejectReason::AuthenticationFailed)
.encode()
.unwrap();
raw.send_to(&unexpected, "127.0.0.1:45784").await.unwrap();
let mut wrong = client_config.clone();
wrong.client.secret.secret = Some("B".repeat(64));
assert!(matches!(
@@ -201,7 +241,8 @@ async fn malformed_and_unauthenticated_packets_do_not_kill_server() {
.unwrap(),
Bytes::from_static(b"still-alive")
);
assert!(stats.snapshot().datagrams_dropped >= 3);
assert!(stats.snapshot().datagrams_dropped >= 5);
assert!(stats.snapshot().oversized_datagrams >= 1);
session.close().await.unwrap();
backend_task.await.unwrap();
server_task.abort();
@@ -275,6 +316,274 @@ async fn authenticated_source_port_change_updates_return_peer() {
server_task.abort();
}
#[tokio::test]
async fn identical_hello_retransmission_recovers_original_accept_and_session() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let backend_addr = backend.local_addr().unwrap();
let (server_config, _) = configs(45788, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let observer = server.handle();
let server_task = tokio::spawn(server.run());
let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let other_source = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let secret = base64::engine::general_purpose::STANDARD
.decode(SECRET)
.unwrap();
let mut nonce = [0u8; 16];
nonce[0] = 44;
let hello = protocol::HelloPacket::authenticated(
&secret,
Bytes::from_static(b"test"),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
nonce,
)
.unwrap();
let hello_bytes = hello.encode().unwrap();
client
.send_to(&hello_bytes, "127.0.0.1:45788")
.await
.unwrap();
let mut response = [0u8; 2048];
let (len, _) = timeout(Duration::from_secs(2), client.recv_from(&mut response))
.await
.unwrap()
.unwrap();
let first_accept = match protocol::parse_packet(&response[..len]).unwrap() {
protocol::Packet::Accept(accept) => accept,
other => panic!("unexpected handshake response: {other:?}"),
};
client
.send_to(&hello_bytes, "127.0.0.1:45788")
.await
.unwrap();
let (len, _) = timeout(Duration::from_secs(2), client.recv_from(&mut response))
.await
.unwrap()
.unwrap();
let retry_accept = match protocol::parse_packet(&response[..len]).unwrap() {
protocol::Packet::Accept(accept) => accept,
other => panic!("identical HELLO should recover ACCEPT, got {other:?}"),
};
assert_eq!(retry_accept, first_accept);
wait_for_active_sessions(&observer, 1).await;
other_source
.send_to(&hello_bytes, "127.0.0.1:45788")
.await
.unwrap();
let (len, _) = timeout(
Duration::from_secs(2),
other_source.recv_from(&mut response),
)
.await
.unwrap()
.unwrap();
assert!(matches!(
protocol::parse_packet(&response[..len]),
Ok(protocol::Packet::Reject(protocol::RejectPacket {
reason: protocol::RejectReason::ReplayDetected,
..
}))
));
let incompatible = protocol::HelloPacket::authenticated(
&secret,
Bytes::from_static(b"test"),
hello.timestamp + 1,
nonce,
)
.unwrap();
client
.send_to(&incompatible.encode().unwrap(), "127.0.0.1:45788")
.await
.unwrap();
let (len, _) = timeout(Duration::from_secs(2), client.recv_from(&mut response))
.await
.unwrap()
.unwrap();
assert!(matches!(
protocol::parse_packet(&response[..len]),
Ok(protocol::Packet::Reject(protocol::RejectPacket {
reason: protocol::RejectReason::ReplayDetected,
..
}))
));
let session_key = protocol::derive_session_key(
&secret,
&nonce,
&first_accept.server_nonce,
first_accept.session_id,
);
let (peer_tx, mut peer_rx) = tokio::sync::mpsc::unbounded_channel();
let backend_task = tokio::spawn(async move {
let mut buffer = [0u8; 2048];
for _ in 0..2 {
let (len, peer) = backend.recv_from(&mut buffer).await.unwrap();
peer_tx.send(peer).unwrap();
backend.send_to(&buffer[..len], peer).await.unwrap();
}
});
for (sequence, payload) in [(1, b"one".as_slice()), (2, b"two".as_slice())] {
let packet = protocol::DataPacket::authenticated(
&session_key,
first_accept.session_id,
sequence,
Bytes::copy_from_slice(payload),
);
client
.send_to(&packet.encode().unwrap(), "127.0.0.1:45788")
.await
.unwrap();
let (len, _) = timeout(Duration::from_secs(2), client.recv_from(&mut response))
.await
.unwrap()
.unwrap();
assert!(matches!(
protocol::parse_packet(&response[..len]),
Ok(protocol::Packet::Data(_))
));
}
backend_task.await.unwrap();
assert_eq!(peer_rx.recv().await, peer_rx.recv().await);
assert_eq!(observer.active_sessions().await, 1);
let close = protocol::SessionPacket::authenticated(
protocol::PacketType::Close,
&session_key,
first_accept.session_id,
3,
);
client
.send_to(&close.encode().unwrap(), "127.0.0.1:45788")
.await
.unwrap();
wait_for_active_sessions(&observer, 0).await;
client
.send_to(&hello_bytes, "127.0.0.1:45788")
.await
.unwrap();
let (len, _) = timeout(Duration::from_secs(2), client.recv_from(&mut response))
.await
.unwrap()
.unwrap();
assert!(matches!(
protocol::parse_packet(&response[..len]),
Ok(protocol::Packet::Reject(protocol::RejectPacket {
reason: protocol::RejectReason::ReplayDetected,
..
}))
));
server_task.abort();
}
#[tokio::test]
async fn flags_are_authenticated_and_invalid_close_cannot_remove_session() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let backend_addr = backend.local_addr().unwrap();
let (server_config, _) = configs(45789, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let observer = server.handle();
let server_task = tokio::spawn(server.run());
let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let secret = base64::engine::general_purpose::STANDARD
.decode(SECRET)
.unwrap();
let mut nonce = [0u8; 16];
nonce[0] = 45;
let hello = protocol::HelloPacket::authenticated(
&secret,
Bytes::from_static(b"test"),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
nonce,
)
.unwrap();
client
.send_to(&hello.encode().unwrap(), "127.0.0.1:45789")
.await
.unwrap();
let mut response = [0u8; 2048];
let (len, _) = timeout(Duration::from_secs(2), client.recv_from(&mut response))
.await
.unwrap()
.unwrap();
let accept = match protocol::parse_packet(&response[..len]).unwrap() {
protocol::Packet::Accept(accept) => accept,
other => panic!("unexpected handshake response: {other:?}"),
};
let key =
protocol::derive_session_key(&secret, &nonce, &accept.server_nonce, accept.session_id);
let data = protocol::DataPacket::authenticated(&key, accept.session_id, 1, b"valid".as_slice());
let mut mutated_flags = data.encode().unwrap();
mutated_flags[2] = 1;
client
.send_to(&mutated_flags, "127.0.0.1:45789")
.await
.unwrap();
assert!(timeout(
Duration::from_millis(100),
backend.recv_from(&mut [0u8; 32])
)
.await
.is_err());
let invalid_close = protocol::SessionPacket::authenticated(
protocol::PacketType::Close,
&key,
accept.session_id,
1,
);
let mut mutated_close_flags = invalid_close.encode().unwrap();
mutated_close_flags[2] = 1;
client
.send_to(&mutated_close_flags, "127.0.0.1:45789")
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(observer.active_sessions().await, 1);
let valid_data = protocol::DataPacket::authenticated(
&key,
accept.session_id,
1,
Bytes::from_static(b"still-open"),
);
client
.send_to(&valid_data.encode().unwrap(), "127.0.0.1:45789")
.await
.unwrap();
let mut backend_buffer = [0u8; 64];
let (len, _) = timeout(
Duration::from_secs(2),
backend.recv_from(&mut backend_buffer),
)
.await
.unwrap()
.unwrap();
assert_eq!(&backend_buffer[..len], b"still-open");
let valid_close = protocol::SessionPacket::authenticated(
protocol::PacketType::Close,
&key,
accept.session_id,
2,
);
client
.send_to(&valid_close.encode().unwrap(), "127.0.0.1:45789")
.await
.unwrap();
wait_for_active_sessions(&observer, 0).await;
server_task.abort();
}
#[tokio::test]
async fn unavailable_backend_does_not_terminate_listener() {
let (server_config, client_config) = configs(45786, 9);