h3cx/wgcloak
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
49ec9d8c33c4720bfde0f93d43ab81b182e83705
wgcloak/tests/direct_udp.rs
T

605 lines
21 KiB
Rust
Raw Blame History

use wgcloak::{config, protocol, transport};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use base64::Engine;
use bytes::Bytes;
use config::Config;
use tokio::{net::UdpSocket, time::timeout};
use transport::{
direct_udp::{DirectUdpServer, DirectUdpServerHandle, DirectUdpSession},
TransportError, TransportSession,
};
const SECRET: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
fn configs(port: u16, backend: u16) -> (config::ServerConfig, config::ClientConfig) {
let server = Config::from_str(&format!(
"[server]\nlisten = \"127.0.0.1\"\n[backend]\ntarget = \"127.0.0.1:{backend}\"\n[transports.direct_udp]\nport = {port}\n[[auth.clients]]\nid = \"test\"\nsecret = \"{SECRET}\"\n"
)).unwrap().into_server().unwrap();
let client = Config::from_str(&format!(
"[client]\nserver = \"127.0.0.1\"\nid = \"test\"\nsecret = \"{SECRET}\"\n[transport.direct_udp]\nport = {port}\n[network]\nmax_datagram_size = 1200\n"
)).unwrap().into_client().unwrap();
(server, client)
}
fn multi_configs(
port: u16,
backend: u16,
) -> (
config::ServerConfig,
config::ClientConfig,
config::ClientConfig,
) {
let server = Config::from_str(&format!(
"[server]\nlisten = \"127.0.0.1\"\n[backend]\ntarget = \"127.0.0.1:{backend}\"\n[transports.direct_udp]\nport = {port}\n[[auth.clients]]\nid = \"one\"\nsecret = \"{SECRET}\"\n[[auth.clients]]\nid = \"two\"\nsecret = \"{SECRET}\"\n"
)).unwrap().into_server().unwrap();
let client = |id: &str| {
Config::from_str(&format!(
"[client]\nserver = \"127.0.0.1\"\nid = \"{id}\"\nsecret = \"{SECRET}\"\n[transport.direct_udp]\nport = {port}\n[network]\nmax_datagram_size = 1200\n"
)).unwrap().into_client().unwrap()
};
(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();
let backend_addr = backend.local_addr().unwrap();
let port = 45781;
let (server_config, client_config) = configs(port, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let stats = server.stats();
let server_task = tokio::spawn(server.run());
let backend_task = tokio::spawn(async move {
let mut buf = [0u8; 2048];
let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
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
.unwrap();
assert_eq!(
timeout(Duration::from_secs(2), session.recv_datagram())
.await
.unwrap()
.unwrap(),
Bytes::from_static(b"opaque")
);
let mut wrong = client_config.clone();
wrong.client.secret.secret =
Some("BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB".to_owned());
let result = DirectUdpSession::connect(&wrong).await;
assert!(matches!(
result,
Err(TransportError::AuthenticationRejected(_))
));
session.close().await.unwrap();
backend_task.await.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(stats.snapshot().expired_sessions, 0);
server_task.abort();
}
#[tokio::test]
async fn direct_udp_enforces_payload_limit_and_expires_idle_session() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let backend_addr = backend.local_addr().unwrap();
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 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();
let backend_addr = backend.local_addr().unwrap();
let (server_config, mut client_one, mut client_two) = multi_configs(45783, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let stats = server.stats();
let server_task = tokio::spawn(server.run());
client_one.transport.direct_udp.port = 45783;
client_two.transport.direct_udp.port = 45783;
let backend_task = tokio::spawn(async move {
let mut peers = Vec::new();
let mut buf = [0u8; 2048];
for _ in 0..2 {
let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
peers.push(peer);
backend.send_to(&buf[..len], peer).await.unwrap();
}
peers
});
let first = DirectUdpSession::connect(&client_one).await.unwrap();
let second = DirectUdpSession::connect(&client_two).await.unwrap();
assert!(stats.snapshot().handshake_attempts >= 2);
first
.send_datagram(Bytes::from_static(b"first"))
.await
.unwrap();
second
.send_datagram(Bytes::from_static(b"second"))
.await
.unwrap();
assert_eq!(
timeout(Duration::from_secs(2), first.recv_datagram())
.await
.unwrap()
.unwrap(),
Bytes::from_static(b"first")
);
assert_eq!(
timeout(Duration::from_secs(2), second.recv_datagram())
.await
.unwrap()
.unwrap(),
Bytes::from_static(b"second")
);
let peers = backend_task.await.unwrap();
assert_ne!(
peers[0], peers[1],
"each authenticated client must have its own backend socket"
);
first.close().await.unwrap();
second.close().await.unwrap();
server_task.abort();
}
#[tokio::test]
async fn malformed_and_unauthenticated_packets_do_not_kill_server() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let backend_addr = backend.local_addr().unwrap();
let (server_config, client_config) = configs(45784, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let stats = server.stats();
let server_task = tokio::spawn(server.run());
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!(
DirectUdpSession::connect(&wrong).await,
Err(TransportError::AuthenticationRejected(_))
));
let backend_task = tokio::spawn(async move {
let mut buf = [0u8; 2048];
let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
backend.send_to(&buf[..len], peer).await.unwrap();
});
let session = DirectUdpSession::connect(&client_config).await.unwrap();
session
.send_datagram(Bytes::from_static(b"still-alive"))
.await
.unwrap();
assert_eq!(
timeout(Duration::from_secs(2), session.recv_datagram())
.await
.unwrap()
.unwrap(),
Bytes::from_static(b"still-alive")
);
assert!(stats.snapshot().datagrams_dropped >= 5);
assert!(stats.snapshot().oversized_datagrams >= 1);
session.close().await.unwrap();
backend_task.await.unwrap();
server_task.abort();
}
#[tokio::test]
async fn authenticated_source_port_change_updates_return_peer() {
let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let backend_addr = backend.local_addr().unwrap();
let (server_config, _) = configs(45785, backend_addr.port());
let server = DirectUdpServer::bind(server_config).await.unwrap();
let server_task = tokio::spawn(server.run());
let first = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let second = 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] = 1;
let hello = protocol::HelloPacket::authenticated(
&secret,
Bytes::from_static(b"test"),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
nonce,
)
.unwrap();
first
.send_to(&hello.encode().unwrap(), "127.0.0.1:45785")
.await
.unwrap();
let mut wire = [0u8; 2048];
let (len, _) = timeout(Duration::from_secs(2), first.recv_from(&mut wire))
.await
.unwrap()
.unwrap();
let accept = match protocol::parse_packet(&wire[..len]).unwrap() {
protocol::Packet::Accept(packet) => packet,
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,
Bytes::from_static(b"rebound"),
);
second
.send_to(&data.encode().unwrap(), "127.0.0.1:45785")
.await
.unwrap();
let backend_task = tokio::spawn(async move {
let mut buf = [0u8; 2048];
let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
backend.send_to(&buf[..len], peer).await.unwrap();
});
let (len, _) = timeout(Duration::from_secs(2), second.recv_from(&mut wire))
.await
.unwrap()
.unwrap();
match protocol::parse_packet(&wire[..len]).unwrap() {
protocol::Packet::Data(packet) => {
assert_eq!(packet.payload, Bytes::from_static(b"rebound"))
}
other => panic!("unexpected rebound response: {other:?}"),
}
backend_task.await.unwrap();
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);
let server = DirectUdpServer::bind(server_config).await.unwrap();
let server_task = tokio::spawn(server.run());
let session = DirectUdpSession::connect(&client_config).await.unwrap();
session
.send_datagram(Bytes::from_static(b"backend-may-be-down"))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(
!server_task.is_finished(),
"backend failure must not terminate the public listener"
);
session.close().await.unwrap();
server_task.abort();
}
Reference in new issue View Git Blame Copy Permalink