From 1bcf2a22c0c936b787e093e1f396e6c38ce8f797 Mon Sep 17 00:00:00 2001 From: Phoenix Kahlo Date: Sun, 24 Nov 2024 16:53:03 -0600 Subject: [PATCH] test(proto): Add and enhance token tests When we first added tests::util::IncomingConnectionBehavior, we opted to use an enum instead of a callback because it seemed cleaner. However, the number of variants have grown, and adding integration tests for validation tokens from NEW_TOKEN frames threatens to make this logic even more complicated. Moreover, there is another advantage to callbacks we have not been exploiting: a stateful FnMut can assert that incoming connection handling within a test follows a certain expected sequence of Incoming properties. As such, this commit replaces TestEndpoint.incoming_connection_behavior with a handle_incoming callback, modifies some existing tests to exploit this functionality to test more things than they were previously, and adds new integration tests for server and client usage of tokens from NEW_TOKEN frames. --- quinn-proto/src/tests/mod.rs | 278 +++++++++++++++++++++++++++++++++- quinn-proto/src/tests/util.rs | 122 +++++++++++---- 2 files changed, 365 insertions(+), 35 deletions(-) diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs index ac254d555c..10bc927233 100644 --- a/quinn-proto/src/tests/mod.rs +++ b/quinn-proto/src/tests/mod.rs @@ -2,7 +2,7 @@ use std::{ convert::TryInto, mem, net::{Ipv4Addr, Ipv6Addr, SocketAddr}, - sync::Arc, + sync::{Arc, Mutex}, }; use assert_matches::assert_matches; @@ -186,7 +186,7 @@ fn draft_version_compat() { fn stateless_retry() { let _guard = subscribe(); let mut pair = Pair::default(); - pair.server.incoming_connection_behavior = IncomingConnectionBehavior::Validate; + pair.server.handle_incoming = Box::new(validate_incoming); let (client_ch, _server_ch) = pair.connect(); pair.client .connections @@ -200,6 +200,259 @@ fn stateless_retry() { assert_eq!(pair.server.known_cids(), 0); } +#[test] +fn use_token() { + let _guard = subscribe(); + let mut pair = Pair::default(); + let client_config = client_config(); + let (client_ch, _server_ch) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + pair.server.handle_incoming = Box::new(|incoming| { + assert!(incoming.remote_address_validated()); + assert!(incoming.may_retry()); + IncomingConnectionBehavior::Accept + }); + let (client_ch_2, _server_ch_2) = pair.connect_with(client_config); + pair.client + .connections + .get_mut(&client_ch_2) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); +} + +#[test] +fn retry_then_use_token() { + let _guard = subscribe(); + let mut pair = Pair::default(); + let client_config = client_config(); + pair.server.handle_incoming = Box::new(validate_incoming); + let (client_ch, _server_ch) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + pair.server.handle_incoming = Box::new(|incoming| { + assert!(incoming.remote_address_validated()); + assert!(incoming.may_retry()); + IncomingConnectionBehavior::Accept + }); + let (client_ch_2, _server_ch_2) = pair.connect_with(client_config); + pair.client + .connections + .get_mut(&client_ch_2) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); +} + +#[test] +fn use_token_then_retry() { + let _guard = subscribe(); + let mut pair = Pair::default(); + let client_config = client_config(); + let (client_ch, _server_ch) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + pair.server.handle_incoming = Box::new({ + let mut i = 0; + move |incoming| { + if i == 0 { + assert!(incoming.remote_address_validated()); + assert!(incoming.may_retry()); + i += 1; + IncomingConnectionBehavior::Retry + } else if i == 1 { + assert!(incoming.remote_address_validated()); + assert!(!incoming.may_retry()); + i += 1; + IncomingConnectionBehavior::Accept + } else { + panic!("too many handle_incoming iterations") + } + } + }); + let (client_ch_2, _server_ch_2) = pair.connect_with(client_config); + pair.client + .connections + .get_mut(&client_ch_2) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); +} + +#[test] +fn use_same_token_twice() { + #[derive(Default)] + struct EvilTokenStore(Mutex); + + impl TokenStore for EvilTokenStore { + fn insert(&self, _server_name: &str, token: Bytes) { + let mut lock = self.0.lock().unwrap(); + if lock.is_empty() { + *lock = token; + } + } + + fn take(&self, _server_name: &str) -> Option { + let lock = self.0.lock().unwrap(); + if lock.is_empty() { + None + } else { + Some(lock.clone()) + } + } + } + + let _guard = subscribe(); + let mut pair = Pair::default(); + let mut client_config = client_config(); + client_config.token_store(Arc::new(EvilTokenStore::default())); + let (client_ch, _server_ch) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + pair.server.handle_incoming = Box::new(|incoming| { + assert!(incoming.remote_address_validated()); + assert!(incoming.may_retry()); + IncomingConnectionBehavior::Accept + }); + let (client_ch_2, _server_ch_2) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch_2) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + pair.server.handle_incoming = Box::new(|incoming| { + assert!(!incoming.remote_address_validated()); + assert!(incoming.may_retry()); + IncomingConnectionBehavior::Accept + }); + let (client_ch_3, _server_ch_3) = pair.connect_with(client_config); + pair.client + .connections + .get_mut(&client_ch_3) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); +} + +#[test] +fn use_token_expired() { + let _guard = subscribe(); + let mocked_time = Arc::new(FakeTimeSource::new()); + let lifetime = Duration::from_secs(10000); + let mut server_config = server_config(); + server_config + .time_source(Arc::clone(&mocked_time) as _) + .validation_token_lifetime(lifetime); + let mut pair = Pair::new(Default::default(), server_config); + let client_config = client_config(); + let (client_ch, _server_ch) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + pair.server.handle_incoming = Box::new(|incoming| { + assert!(incoming.remote_address_validated()); + assert!(incoming.may_retry()); + IncomingConnectionBehavior::Accept + }); + let (client_ch_2, _server_ch_2) = pair.connect_with(client_config.clone()); + pair.client + .connections + .get_mut(&client_ch_2) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); + + mocked_time.advance(lifetime + Duration::from_secs(1)); + + pair.server.handle_incoming = Box::new(|incoming| { + assert!(!incoming.remote_address_validated()); + assert!(incoming.may_retry()); + IncomingConnectionBehavior::Accept + }); + let (client_ch_3, _server_ch_3) = pair.connect_with(client_config); + pair.client + .connections + .get_mut(&client_ch_3) + .unwrap() + .close(pair.time, VarInt(42), Bytes::new()); + pair.drive(); + assert_eq!(pair.client.known_connections(), 0); + assert_eq!(pair.client.known_cids(), 0); + assert_eq!(pair.server.known_connections(), 0); + assert_eq!(pair.server.known_cids(), 0); +} + #[test] fn server_stateless_reset() { let _guard = subscribe(); @@ -553,7 +806,7 @@ fn high_latency_handshake() { fn zero_rtt_happypath() { let _guard = subscribe(); let mut pair = Pair::default(); - pair.server.incoming_connection_behavior = IncomingConnectionBehavior::Validate; + pair.server.handle_incoming = Box::new(validate_incoming); let config = client_config(); // Establish normal connection @@ -722,7 +975,7 @@ fn test_zero_rtt_incoming_limit(configure_server: CLIENT_PORTS.lock().unwrap().next().unwrap(), ); info!("resuming session"); - pair.server.incoming_connection_behavior = IncomingConnectionBehavior::Wait; + pair.server.handle_incoming = Box::new(|_| IncomingConnectionBehavior::Wait); let client_ch = pair.begin_connect(config); assert!(pair.client_conn_mut(client_ch).has_0rtt()); let s = pair.client_streams(client_ch).open(Dir::Uni).unwrap(); @@ -2992,7 +3245,7 @@ fn pure_sender_voluntarily_acks() { fn reject_manually() { let _guard = subscribe(); let mut pair = Pair::default(); - pair.server.incoming_connection_behavior = IncomingConnectionBehavior::RejectAll; + pair.server.handle_incoming = Box::new(|_| IncomingConnectionBehavior::Reject); // The server should now reject incoming connections. let client_ch = pair.begin_connect(client_config()); @@ -3012,7 +3265,20 @@ fn reject_manually() { fn validate_then_reject_manually() { let _guard = subscribe(); let mut pair = Pair::default(); - pair.server.incoming_connection_behavior = IncomingConnectionBehavior::ValidateThenReject; + pair.server.handle_incoming = Box::new({ + let mut i = 0; + move |incoming| { + if incoming.remote_address_validated() { + assert_eq!(i, 1); + i += 1; + IncomingConnectionBehavior::Reject + } else { + assert_eq!(i, 0); + i += 1; + IncomingConnectionBehavior::Retry + } + } + }); // The server should now retry and reject incoming connections. let client_ch = pair.begin_connect(client_config()); diff --git a/quinn-proto/src/tests/util.rs b/quinn-proto/src/tests/util.rs index 19c80570d0..2780ea9e3a 100644 --- a/quinn-proto/src/tests/util.rs +++ b/quinn-proto/src/tests/util.rs @@ -1,6 +1,6 @@ use std::{ cmp, - collections::{HashMap, VecDeque}, + collections::{HashMap, HashSet, VecDeque}, env, io::{self, Write}, mem, @@ -22,7 +22,7 @@ use tracing::{info_span, trace}; use super::crypto::rustls::{configured_provider, QuicClientConfig, QuicServerConfig}; use super::*; -use crate::{Duration, Instant}; +use crate::{Duration, Instant, SystemTime}; pub(super) const DEFAULT_MTU: usize = 1452; @@ -81,7 +81,7 @@ impl Pair { epoch: now, time: now, mtu: DEFAULT_MTU, - latency: Duration::ZERO, + latency: Duration::new(0, 0), spins: 0, last_spin: false, congestion_experienced: false, @@ -297,25 +297,32 @@ pub(super) struct TestEndpoint { conn_events: HashMap>, pub(super) captured_packets: Vec>, pub(super) capture_inbound_packets: bool, - pub(super) incoming_connection_behavior: IncomingConnectionBehavior, + pub(super) handle_incoming: Box IncomingConnectionBehavior>, pub(super) waiting_incoming: Vec, } #[derive(Debug, Copy, Clone)] pub(super) enum IncomingConnectionBehavior { - AcceptAll, - RejectAll, - Validate, - ValidateThenReject, + Accept, + Reject, + Retry, Wait, } +pub(super) fn validate_incoming(incoming: &Incoming) -> IncomingConnectionBehavior { + if incoming.remote_address_validated() { + IncomingConnectionBehavior::Accept + } else { + IncomingConnectionBehavior::Retry + } +} + impl TestEndpoint { fn new(endpoint: Endpoint, addr: SocketAddr) -> Self { let socket = if env::var_os("SSLKEYLOGFILE").is_some() { let socket = UdpSocket::bind(addr).expect("failed to bind UDP socket"); socket - .set_read_timeout(Some(Duration::from_millis(10))) + .set_read_timeout(Some(Duration::new(0, 10_000_000))) .unwrap(); Some(socket) } else { @@ -334,7 +341,7 @@ impl TestEndpoint { conn_events: HashMap::default(), captured_packets: Vec::new(), capture_inbound_packets: false, - incoming_connection_behavior: IncomingConnectionBehavior::AcceptAll, + handle_incoming: Box::new(|_| IncomingConnectionBehavior::Accept), waiting_incoming: Vec::new(), } } @@ -364,26 +371,15 @@ impl TestEndpoint { { match event { DatagramEvent::NewConnection(incoming) => { - match self.incoming_connection_behavior { - IncomingConnectionBehavior::AcceptAll => { + match (self.handle_incoming)(&incoming) { + IncomingConnectionBehavior::Accept => { let _ = self.try_accept(incoming, now); } - IncomingConnectionBehavior::RejectAll => { + IncomingConnectionBehavior::Reject => { self.reject(incoming); } - IncomingConnectionBehavior::Validate => { - if incoming.remote_address_validated() { - let _ = self.try_accept(incoming, now); - } else { - self.retry(incoming); - } - } - IncomingConnectionBehavior::ValidateThenReject => { - if incoming.remote_address_validated() { - self.reject(incoming); - } else { - self.retry(incoming); - } + IncomingConnectionBehavior::Retry => { + self.retry(incoming); } IncomingConnectionBehavior::Wait => { self.waiting_incoming.push(incoming); @@ -564,14 +560,22 @@ impl Write for TestWriter { } pub(super) fn server_config() -> ServerConfig { - ServerConfig::with_crypto(Arc::new(server_crypto())) + let mut config = ServerConfig::with_crypto(Arc::new(server_crypto())); + config + .validation_tokens_sent(2) + .validation_token_log(Arc::new(SimpleTokenLog::default())); + config } pub(super) fn server_config_with_cert( cert: CertificateDer<'static>, key: PrivateKeyDer<'static>, ) -> ServerConfig { - ServerConfig::with_crypto(Arc::new(server_crypto_with_cert(cert, key))) + let mut config = ServerConfig::with_crypto(Arc::new(server_crypto_with_cert(cert, key))); + config + .validation_tokens_sent(2) + .validation_token_log(Arc::new(SimpleTokenLog::default())); + config } pub(super) fn server_crypto() -> QuicServerConfig { @@ -609,7 +613,9 @@ fn server_crypto_inner( } pub(super) fn client_config() -> ClientConfig { - ClientConfig::new(Arc::new(client_crypto())) + let mut config = ClientConfig::new(Arc::new(client_crypto())); + config.token_store(Arc::new(SimpleTokenStore::default())); + config } pub(super) fn client_config_with_deterministic_pns() -> ClientConfig { @@ -717,3 +723,61 @@ lazy_static! { pub(crate) static ref CERTIFIED_KEY: rcgen::CertifiedKey = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); } + +#[derive(Default)] +struct SimpleTokenLog(Mutex>); + +impl TokenLog for SimpleTokenLog { + fn check_and_insert( + &self, + rand: u128, + _issued: SystemTime, + _lifetime: Duration, + ) -> Result<(), TokenReuseError> { + if self.0.lock().unwrap().insert(rand) { + Ok(()) + } else { + Err(TokenReuseError) + } + } +} + +#[derive(Default)] +struct SimpleTokenStore(Mutex>>); + +impl TokenStore for SimpleTokenStore { + fn insert(&self, server_name: &str, token: Bytes) { + self.0 + .lock() + .unwrap() + .entry(server_name.into()) + .or_default() + .push_back(token); + } + + fn take(&self, server_name: &str) -> Option { + self.0 + .lock() + .unwrap() + .get_mut(server_name) + .and_then(|queue| queue.pop_front()) + } +} + +pub(super) struct FakeTimeSource(Mutex); + +impl FakeTimeSource { + pub(super) fn new() -> Self { + Self(Mutex::new(SystemTime::now())) + } + + pub(super) fn advance(&self, dur: Duration) { + *self.0.lock().unwrap() += dur; + } +} + +impl TimeSource for FakeTimeSource { + fn now(&self) -> SystemTime { + *self.0.lock().unwrap() + } +}