From 1af91f506120a3d68d6c938ddc42ea3a73218507 Mon Sep 17 00:00:00 2001 From: Lars Eggert Date: Wed, 27 Mar 2024 02:37:19 +0200 Subject: [PATCH] test: Let packets be modified with a closure during tests (#1773) * test: Let packets be modified with a closure during tests Broken out of #1678 to reduce the size of that PR. * Support dropping packets via `DatagramModifier` --- neqo-transport/src/connection/tests/mod.rs | 83 +++++++++++++++++++--- 1 file changed, 73 insertions(+), 10 deletions(-) diff --git a/neqo-transport/src/connection/tests/mod.rs b/neqo-transport/src/connection/tests/mod.rs index b6ce08f8d1..6f598fb23e 100644 --- a/neqo-transport/src/connection/tests/mod.rs +++ b/neqo-transport/src/connection/tests/mod.rs @@ -170,12 +170,17 @@ impl crate::connection::test_internal::FrameWriter for PingWriter { } } +trait DatagramModifier: FnMut(Datagram) -> Option {} + +impl DatagramModifier for T where T: FnMut(Datagram) -> Option {} + /// Drive the handshake between the client and server. -fn handshake( +fn handshake_with_modifier( client: &mut Connection, server: &mut Connection, now: Instant, rtt: Duration, + mut modifier: impl DatagramModifier, ) -> Instant { let mut a = client; let mut b = server; @@ -212,7 +217,11 @@ fn handshake( did_ping[a.role()] = true; } assert!(had_input || output.is_some()); - input = output; + if let Some(d) = output { + input = modifier(d); + } else { + input = output; + } qtrace!("handshake: t += {:?}", rtt / 2); now += rtt / 2; mem::swap(&mut a, &mut b); @@ -223,6 +232,15 @@ fn handshake( now } +fn handshake( + client: &mut Connection, + server: &mut Connection, + now: Instant, + rtt: Duration, +) -> Instant { + handshake_with_modifier(client, server, now, rtt, Some) +} + fn connect_fail( client: &mut Connection, server: &mut Connection, @@ -234,11 +252,12 @@ fn connect_fail( assert_error(server, &ConnectionError::Transport(server_error)); } -fn connect_with_rtt( +fn connect_with_rtt_and_modifier( client: &mut Connection, server: &mut Connection, now: Instant, rtt: Duration, + modifier: impl DatagramModifier, ) -> Instant { fn check_rtt(stats: &Stats, rtt: Duration) { assert_eq!(stats.rtt, rtt); @@ -246,7 +265,7 @@ fn connect_with_rtt( let n = stats.frame_rx.ack + usize::from(stats.rtt_init_guess); assert_eq!(stats.rttvar, rttvar_after_n_updates(n, rtt)); } - let now = handshake(client, server, now, rtt); + let now = handshake_with_modifier(client, server, now, rtt, modifier); assert_eq!(*client.state(), State::Confirmed); assert_eq!(*server.state(), State::Confirmed); @@ -255,6 +274,15 @@ fn connect_with_rtt( now } +fn connect_with_rtt( + client: &mut Connection, + server: &mut Connection, + now: Instant, + rtt: Duration, +) -> Instant { + connect_with_rtt_and_modifier(client, server, now, rtt, Some) +} + fn connect(client: &mut Connection, server: &mut Connection) { connect_with_rtt(client, server, now(), Duration::new(0, 0)); } @@ -301,8 +329,13 @@ fn assert_idle(client: &mut Connection, server: &mut Connection, rtt: Duration, } /// Connect with an RTT and then force both peers to be idle. -fn connect_rtt_idle(client: &mut Connection, server: &mut Connection, rtt: Duration) -> Instant { - let now = connect_with_rtt(client, server, now(), rtt); +fn connect_rtt_idle_with_modifier( + client: &mut Connection, + server: &mut Connection, + rtt: Duration, + modifier: impl DatagramModifier, +) -> Instant { + let now = connect_with_rtt_and_modifier(client, server, now(), rtt, modifier); assert_idle(client, server, rtt, now); // Drain events from both as well. _ = client.events().count(); @@ -311,8 +344,20 @@ fn connect_rtt_idle(client: &mut Connection, server: &mut Connection, rtt: Durat now } +fn connect_rtt_idle(client: &mut Connection, server: &mut Connection, rtt: Duration) -> Instant { + connect_rtt_idle_with_modifier(client, server, rtt, Some) +} + +fn connect_force_idle_with_modifier( + client: &mut Connection, + server: &mut Connection, + modifier: impl DatagramModifier, +) { + connect_rtt_idle_with_modifier(client, server, Duration::new(0, 0), modifier); +} + fn connect_force_idle(client: &mut Connection, server: &mut Connection) { - connect_rtt_idle(client, server, Duration::new(0, 0)); + connect_force_idle_with_modifier(client, server, Some); } fn fill_stream(c: &mut Connection, stream: StreamId) { @@ -524,12 +569,14 @@ fn assert_full_cwnd(packets: &[Datagram], cwnd: usize) { } /// Send something on a stream from `sender` to `receiver`, maybe allowing for pacing. +/// Takes a modifier function that can be used to modify the datagram before it is sent. /// Return the resulting datagram and the new time. #[must_use] -fn send_something_paced( +fn send_something_paced_with_modifier( sender: &mut Connection, mut now: Instant, allow_pacing: bool, + mut modifier: impl DatagramModifier, ) -> (Datagram, Instant) { let stream_id = sender.stream_create(StreamType::UniDi).unwrap(); assert!(sender.stream_send(stream_id, DEFAULT_STREAM_DATA).is_ok()); @@ -544,16 +591,32 @@ fn send_something_paced( .dgram() .expect("send_something: should have something to send") } - Output::Datagram(d) => d, + Output::Datagram(d) => modifier(d).unwrap(), Output::None => panic!("send_something: got Output::None"), }; (dgram, now) } +fn send_something_paced( + sender: &mut Connection, + now: Instant, + allow_pacing: bool, +) -> (Datagram, Instant) { + send_something_paced_with_modifier(sender, now, allow_pacing, Some) +} + +fn send_something_with_modifier( + sender: &mut Connection, + now: Instant, + modifier: impl DatagramModifier, +) -> Datagram { + send_something_paced_with_modifier(sender, now, false, modifier).0 +} + /// Send something on a stream from `sender` to `receiver`. /// Return the resulting datagram. fn send_something(sender: &mut Connection, now: Instant) -> Datagram { - send_something_paced(sender, now, false).0 + send_something_with_modifier(sender, now, Some) } /// Send something on a stream from `sender` to `receiver`.