feat(transport): UDP multi-client demux by peer address

UdpServer now serves many concurrent peers on one socket (removes v1's "one peer per accept" limitation). PeerSocket becomes an enum: ConnectedClient (client side, unchanged behavior) vs SharedServer (server side, channel-fed inbox). A master loop reads the shared socket and routes datagrams to the right per-peer inbox by source address; an unknown peer's first TYPE_HS datagram spawns a new handshake task that, on success, hands the established UdpConnection to accept(). Cleanup is lazy via mpsc::Closed — handshake failures and connection drops self- evict from the map. A small Arc<MasterTask> keeps the loop alive for the lifetime of UdpServer OR any spawned UdpConnection, so existing single- client tests (which move UdpServer into an accept task) still pass. ReliableHsAdapter and run_reliable_handshake are unchanged. UdpClient API unchanged. Added 3 tests: two concurrent clients with cross-talk isolation, bad-CA client doesn't block legitimate ones, dropped peer doesn't block others. Workspace: 117 tests green, clippy/fmt clean. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-27 01:27:06 +03:00
parent c95e1a482c
commit 4d1bdba55d
2 changed files with 626 additions and 134 deletions
@@ -0,0 +1,333 @@
+//! Multi-client integration tests for the Aura UDP transport (the v2 master-loop demuxer).
+//!
+//! These prove that a single bound [`UdpServer`] can simultaneously serve **many** peers, that bad
+//! peers do not poison the server, and that established connections survive other peers coming and
+//! going. The single-client and lossy-channel tests live in `udp_loopback.rs`; here we focus on
+//! demuxer correctness.
+//!
+//! * [`udp_multi_client_two_concurrent`] — bind one server, drive two clients (different client CNs)
+//!   to it concurrently, accept twice, and verify both connections are independent (no cross-talk;
+//!   each side learns the correct peer id).
+//! * [`udp_bad_ca_does_not_block_other_clients`] — a third client with a foreign CA fails the
+//!   handshake; the server must keep accepting subsequent legitimate clients on the same port.
+//! * [`udp_dropped_connection_does_not_block_other_clients`] — drop one client's connection mid-flight
+//!   and prove the server keeps serving the other plus accepts a fresh one.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use aura_pki::AuraCa;
+use aura_proto::{ClientConfig, PacketConnection, ServerConfig};
+use aura_transport::{UdpClient, UdpConnection, UdpOpts, UdpServer};
+
+const SERVER_NAME: &str = "localhost";
+
+/// Mint a CA, a server cert, and a set of client certs whose CNs are taken from `client_ids`.
+fn make_configs(client_ids: &[&str]) -> (ServerConfig, Vec<ClientConfig>) {
+    let ca = AuraCa::generate("Aura UDP Multi-Client Test CA").expect("generate CA");
+    let server_cert = ca
+        .issue_server_cert(SERVER_NAME)
+        .expect("issue server cert");
+    let ca_pem = ca.ca_cert_pem();
+    let server_cfg = ServerConfig {
+        ca_cert_pem: ca_pem.clone(),
+        server_cert_pem: server_cert.cert_pem,
+        server_key_pem: server_cert.key_pem,
+    };
+    let client_cfgs: Vec<ClientConfig> = client_ids
+        .iter()
+        .map(|id| {
+            let c = ca.issue_client_cert(id).expect("issue client cert");
+            ClientConfig {
+                ca_cert_pem: ca_pem.clone(),
+                client_cert_pem: c.cert_pem,
+                client_key_pem: c.key_pem,
+                server_name: SERVER_NAME.to_string(),
+            }
+        })
+        .collect();
+    (server_cfg, client_cfgs)
+}
+
+/// Mint a **separate** CA + matching client cert; the resulting `ClientConfig` will trust this CA
+/// for the server (so it will reject the real server) and present a cert the real server will not
+/// verify either. Used to drive a handshake failure that must NOT take down the server.
+fn make_foreign_ca_client(server_name: &str, client_cn: &str) -> ClientConfig {
+    let foreign = AuraCa::generate("Foreign CA").expect("generate foreign CA");
+    let client_cert = foreign
+        .issue_client_cert(client_cn)
+        .expect("issue client cert under foreign CA");
+    ClientConfig {
+        ca_cert_pem: foreign.ca_cert_pem(),
+        client_cert_pem: client_cert.cert_pem,
+        client_key_pem: client_cert.key_pem,
+        server_name: server_name.to_string(),
+    }
+}
+
+/// Round-trip a payload `pkt` from `tx` to `rx` and assert byte equality.
+async fn round_trip(tx: &Arc<dyn PacketConnection>, rx: &Arc<dyn PacketConnection>, pkt: &[u8]) {
+    tx.send_packet(pkt).await.expect("send");
+    let got = tokio::time::timeout(Duration::from_secs(5), rx.recv_packet())
+        .await
+        .expect("recv did not arrive within 5s")
+        .expect("recv");
+    assert_eq!(got, pkt, "payload mismatch over round trip");
+}
+
+#[tokio::test]
+async fn udp_multi_client_two_concurrent() {
+    let (server_cfg, client_cfgs) = make_configs(&["client-a", "client-b"]);
+    let opts = UdpOpts::default();
+
+    let server =
+        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind server");
+    let server_addr = server.local_addr().expect("server addr");
+    let server = Arc::new(server);
+
+    // Spawn two server-side accepts in parallel; they must each pull their own connection from the
+    // master-loop's accept queue.
+    let s_a = server.clone();
+    let accept_a = tokio::spawn(async move { s_a.accept().await });
+    let s_b = server.clone();
+    let accept_b = tokio::spawn(async move { s_b.accept().await });
+
+    // Spawn the two clients concurrently. They share the server's bound port.
+    let cfg_a = client_cfgs[0].clone();
+    let cfg_b = client_cfgs[1].clone();
+    let connect_a = tokio::spawn(async move { UdpClient::connect(server_addr, cfg_a, opts).await });
+    let connect_b = tokio::spawn(async move { UdpClient::connect(server_addr, cfg_b, opts).await });
+
+    // Wait for everything to settle (generous timeout — handshake should be sub-second on loopback).
+    let timeout = Duration::from_secs(15);
+    let server_a: UdpConnection = tokio::time::timeout(timeout, accept_a)
+        .await
+        .expect("accept_a within timeout")
+        .expect("accept_a join")
+        .expect("accept_a result");
+    let server_b: UdpConnection = tokio::time::timeout(timeout, accept_b)
+        .await
+        .expect("accept_b within timeout")
+        .expect("accept_b join")
+        .expect("accept_b result");
+    let client_a: UdpConnection = tokio::time::timeout(timeout, connect_a)
+        .await
+        .expect("connect_a within timeout")
+        .expect("connect_a join")
+        .expect("connect_a result");
+    let client_b: UdpConnection = tokio::time::timeout(timeout, connect_b)
+        .await
+        .expect("connect_b within timeout")
+        .expect("connect_b join")
+        .expect("connect_b result");
+
+    // Each server-side connection has a `peer_id` of either `client-a` or `client-b`; the accept
+    // order is *not* guaranteed (whichever handshake finishes first), so detect which is which
+    // and pair them with the matching client connection.
+    let id_a = server_a.peer_id().map(str::to_owned);
+    let id_b = server_b.peer_id().map(str::to_owned);
+    let mut ids = vec![id_a.clone(), id_b.clone()];
+    ids.sort();
+    assert_eq!(
+        ids,
+        vec![Some("client-a".to_string()), Some("client-b".to_string())],
+        "the two server-side connections must carry client-a and client-b CNs (no duplicates)"
+    );
+
+    let (srv_for_a, srv_for_b) = if id_a.as_deref() == Some("client-a") {
+        (server_a, server_b)
+    } else {
+        (server_b, server_a)
+    };
+
+    // Each side sees its own peer id (client side sees the server name).
+    assert_eq!(client_a.peer_id(), Some(SERVER_NAME));
+    assert_eq!(client_b.peer_id(), Some(SERVER_NAME));
+
+    let client_a: Arc<dyn PacketConnection> = Arc::new(client_a);
+    let client_b: Arc<dyn PacketConnection> = Arc::new(client_b);
+    let server_for_a: Arc<dyn PacketConnection> = Arc::new(srv_for_a);
+    let server_for_b: Arc<dyn PacketConnection> = Arc::new(srv_for_b);
+
+    // No cross-talk: A's payload reaches A's server-side conn (not B's), and vice versa.
+    round_trip(&client_a, &server_for_a, b"hi from a").await;
+    round_trip(&client_b, &server_for_b, b"hi from b").await;
+    round_trip(&server_for_a, &client_a, b"reply to a").await;
+    round_trip(&server_for_b, &client_b, b"reply to b").await;
+
+    // And both directions still work concurrently (no head-of-line blocking via the master loop).
+    let a_send = {
+        let c = client_a.clone();
+        let s = server_for_a.clone();
+        tokio::spawn(async move {
+            c.send_packet(b"a-concurrent").await.unwrap();
+            s.recv_packet().await.unwrap()
+        })
+    };
+    let b_send = {
+        let c = client_b.clone();
+        let s = server_for_b.clone();
+        tokio::spawn(async move {
+            c.send_packet(b"b-concurrent").await.unwrap();
+            s.recv_packet().await.unwrap()
+        })
+    };
+    assert_eq!(a_send.await.unwrap(), b"a-concurrent");
+    assert_eq!(b_send.await.unwrap(), b"b-concurrent");
+}
+
+#[tokio::test]
+async fn udp_bad_ca_does_not_block_other_clients() {
+    let (server_cfg, client_cfgs) = make_configs(&["client-good"]);
+    // Use a tighter handshake timeout so the failing peer fails quickly and the test finishes
+    // even if the rogue client retransmits its ClientHello for a while.
+    let opts = UdpOpts {
+        hs_timeout: Duration::from_secs(3),
+        ..UdpOpts::default()
+    };
+
+    let server =
+        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind server");
+    let server_addr = server.local_addr().expect("server addr");
+    let server = Arc::new(server);
+
+    // A rogue client with a foreign CA: its server-side handshake task will fail. The server must
+    // log + drop and keep accepting OTHER peers.
+    let foreign_cfg = make_foreign_ca_client(SERVER_NAME, "rogue");
+    let rogue =
+        tokio::spawn(async move { UdpClient::connect(server_addr, foreign_cfg, opts).await });
+
+    // Give the rogue task a head start so the server's master loop registers it first.
+    tokio::time::sleep(Duration::from_millis(50)).await;
+
+    // Now the legitimate client connects. The server must still accept it.
+    let cfg = client_cfgs[0].clone();
+    let s = server.clone();
+    let accept_good = tokio::spawn(async move { s.accept().await });
+    let connect_good =
+        tokio::spawn(async move { UdpClient::connect(server_addr, cfg, opts).await });
+
+    let timeout = Duration::from_secs(15);
+    let server_good: UdpConnection = tokio::time::timeout(timeout, accept_good)
+        .await
+        .expect("accept_good within timeout")
+        .expect("accept_good join")
+        .expect("accept_good result");
+    let client_good: UdpConnection = tokio::time::timeout(timeout, connect_good)
+        .await
+        .expect("connect_good within timeout")
+        .expect("connect_good join")
+        .expect("connect_good result");
+
+    assert_eq!(
+        server_good.peer_id(),
+        Some("client-good"),
+        "server must learn the good client's CN despite the rogue peer"
+    );
+
+    let server_good: Arc<dyn PacketConnection> = Arc::new(server_good);
+    let client_good: Arc<dyn PacketConnection> = Arc::new(client_good);
+    round_trip(&client_good, &server_good, b"still serving").await;
+    round_trip(&server_good, &client_good, b"yes we are").await;
+
+    // The rogue connect should eventually fail (foreign CA → server's handshake rejects, the
+    // client's handshake adapter then errors out on the deadline / chain mismatch). We do not
+    // care about the exact error; we only require that it *errors*, not that it succeeds.
+    let rogue_result = tokio::time::timeout(Duration::from_secs(10), rogue)
+        .await
+        .expect("rogue task should terminate")
+        .expect("rogue task join");
+    assert!(
+        rogue_result.is_err(),
+        "rogue client (foreign CA) must NOT succeed in establishing a connection"
+    );
+}
+
+/// Establish ONE client against a running multi-client server, then verify the server-side conn
+/// has the expected CN. The accept happens in its own spawned task to avoid blocking the connect.
+async fn establish_one(
+    server: &Arc<UdpServer>,
+    server_addr: std::net::SocketAddr,
+    cfg: ClientConfig,
+    opts: UdpOpts,
+    expect_cn: &str,
+) -> (UdpConnection, UdpConnection) {
+    let s = server.clone();
+    let acc = tokio::spawn(async move { s.accept().await });
+    let con = tokio::spawn(async move { UdpClient::connect(server_addr, cfg, opts).await });
+    let timeout = Duration::from_secs(15);
+    let srv = tokio::time::timeout(timeout, acc)
+        .await
+        .expect("accept timely")
+        .expect("accept join")
+        .expect("accept result");
+    let cli = tokio::time::timeout(timeout, con)
+        .await
+        .expect("connect timely")
+        .expect("connect join")
+        .expect("connect result");
+    assert_eq!(
+        srv.peer_id(),
+        Some(expect_cn),
+        "server learned wrong CN for this client"
+    );
+    (srv, cli)
+}
+
+#[tokio::test]
+async fn udp_dropped_connection_does_not_block_other_clients() {
+    let (server_cfg, client_cfgs) = make_configs(&["client-1", "client-2", "client-3"]);
+    let opts = UdpOpts::default();
+
+    let server =
+        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind server");
+    let server_addr = server.local_addr().expect("server addr");
+    let server = Arc::new(server);
+
+    // Connect clients sequentially so the (server-side, client-side) pairing is unambiguous.
+    let (srv1, cli1) = establish_one(
+        &server,
+        server_addr,
+        client_cfgs[0].clone(),
+        opts,
+        "client-1",
+    )
+    .await;
+    let (srv2, cli2) = establish_one(
+        &server,
+        server_addr,
+        client_cfgs[1].clone(),
+        opts,
+        "client-2",
+    )
+    .await;
+
+    let srv2: Arc<dyn PacketConnection> = Arc::new(srv2);
+    let cli2: Arc<dyn PacketConnection> = Arc::new(cli2);
+    // Sanity: client-2 works.
+    round_trip(&cli2, &srv2, b"keep-1").await;
+
+    // Drop both ends of client-1's pair: the server-side `UdpConnection` is dropped, its
+    // `PeerSocket` (with the master's per-peer inbox receiver) is dropped, and the master loop's
+    // next datagram from client-1's address — if any — will `Closed` and evict the entry.
+    drop(srv1);
+    drop(cli1);
+    tokio::time::sleep(Duration::from_millis(50)).await;
+
+    // The other client must keep working.
+    round_trip(&cli2, &srv2, b"keep-2").await;
+    round_trip(&srv2, &cli2, b"keep-3").await;
+
+    // A fresh client-3 must also still be accepted.
+    let (srv3, cli3) = establish_one(
+        &server,
+        server_addr,
+        client_cfgs[2].clone(),
+        opts,
+        "client-3",
+    )
+    .await;
+    let srv3: Arc<dyn PacketConnection> = Arc::new(srv3);
+    let cli3: Arc<dyn PacketConnection> = Arc::new(cli3);
+    round_trip(&cli3, &srv3, b"hi-3").await;
+}