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

@@ -170,12 +170,17 @@ impl crate::connection::test_internal::FrameWriter for PingWriter {
     }
 }
 
+trait DatagramModifier: FnMut(Datagram) -> Option<Datagram> {}
+
+impl<T> DatagramModifier for T where T: FnMut(Datagram) -> Option<Datagram> {}
+
 /// 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,19 +252,20 @@ 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);
         // Validate that rttvar has been computed correctly based on the number of RTT updates.
         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`.