From a2fcb05e261b69256b3f73af8a1ea76a67a408cb Mon Sep 17 00:00:00 2001 From: Martin Thomson Date: Tue, 8 Oct 2024 17:24:23 +1100 Subject: [PATCH] wrappers for common usage of test_frame_writer This reduces boilerplate and encapsulates a good pattern of usage. I wasn't able to make `test_frame_writer` private, but this is close enough. --- neqo-transport/src/connection/mod.rs | 13 ++++++++++++ .../src/connection/tests/datagram.rs | 15 +++++++------- .../src/connection/tests/migration.rs | 20 ++++++------------- neqo-transport/src/connection/tests/mod.rs | 13 +++++++++++- .../src/connection/tests/recovery.rs | 7 ++++--- 5 files changed, 43 insertions(+), 25 deletions(-) diff --git a/neqo-transport/src/connection/mod.rs b/neqo-transport/src/connection/mod.rs index eb8ff77881..0184ddce7b 100644 --- a/neqo-transport/src/connection/mod.rs +++ b/neqo-transport/src/connection/mod.rs @@ -1146,6 +1146,19 @@ impl Connection { } } + /// A test-only output function that uses the provided writer to + /// pack something extra into the output. + #[cfg(test)] + pub fn test_write_frames(&mut self, writer: W, now: Instant) -> Output + where + W: test_internal::FrameWriter + 'static, + { + self.test_frame_writer = Some(Box::new(writer)); + let res = self.process_output(now); + self.test_frame_writer = None; + res + } + /// Process input and generate output. #[must_use = "Output of the process function must be handled"] pub fn process(&mut self, dgram: Option<&Datagram>, now: Instant) -> Output { diff --git a/neqo-transport/src/connection/tests/datagram.rs b/neqo-transport/src/connection/tests/datagram.rs index 73ea72d4b4..7fa38ec489 100644 --- a/neqo-transport/src/connection/tests/datagram.rs +++ b/neqo-transport/src/connection/tests/datagram.rs @@ -399,10 +399,11 @@ fn dgram_no_allowed() { let mut client = default_client(); let mut server = default_server(); connect_force_idle(&mut client, &mut server); - server.test_frame_writer = Some(Box::new(InsertDatagram { data: DATA_MTU })); - let out = server.process_output(now()).dgram().unwrap(); - server.test_frame_writer = None; + let out = server + .test_write_frames(InsertDatagram { data: DATA_MTU }, now()) + .dgram() + .unwrap(); client.process_input(&out, now()); assert_error(&client, &CloseReason::Transport(Error::ProtocolViolation)); @@ -415,10 +416,10 @@ fn dgram_too_big() { let mut server = default_server(); connect_force_idle(&mut client, &mut server); - server.test_frame_writer = Some(Box::new(InsertDatagram { data: DATA_MTU })); - let out = server.process_output(now()).dgram().unwrap(); - server.test_frame_writer = None; - + let out = server + .test_write_frames(InsertDatagram { data: DATA_MTU }, now()) + .dgram() + .unwrap(); client.process_input(&out, now()); assert_error(&client, &CloseReason::Transport(Error::ProtocolViolation)); diff --git a/neqo-transport/src/connection/tests/migration.rs b/neqo-transport/src/connection/tests/migration.rs index 1499d55f06..541ab7dfeb 100644 --- a/neqo-transport/src/connection/tests/migration.rs +++ b/neqo-transport/src/connection/tests/migration.rs @@ -25,7 +25,7 @@ use super::{ }; use crate::{ cid::LOCAL_ACTIVE_CID_LIMIT, - connection::tests::send_something_paced, + connection::tests::{send_something_paced, send_with_extra}, frame::FRAME_TYPE_NEW_CONNECTION_ID, packet::PacketBuilder, path::MAX_PATH_PROBES, @@ -810,9 +810,7 @@ fn retire_all() { let original_cid = ConnectionId::from(get_cid(&send_something(&mut client, now()))); - server.test_frame_writer = Some(Box::new(RetireAll { cid_gen })); - let ncid = send_something(&mut server, now()); - server.test_frame_writer = None; + let ncid = send_with_extra(&mut server, RetireAll { cid_gen }, now()); let new_cid_before = client.stats().frame_rx.new_connection_id; let retire_cid_before = client.stats().frame_tx.retire_connection_id; @@ -861,9 +859,7 @@ fn retire_prior_to_migration_failure() { // Have the server receive the probe, but separately have it decide to // retire all of the available connection IDs. - server.test_frame_writer = Some(Box::new(RetireAll { cid_gen })); - let retire_all = send_something(&mut server, now()); - server.test_frame_writer = None; + let retire_all = send_with_extra(&mut server, RetireAll { cid_gen }, now()); let resp = server.process(Some(&probe), now()).dgram().unwrap(); assert_v4_path(&resp, true); @@ -916,9 +912,7 @@ fn retire_prior_to_migration_success() { // Have the server receive the probe, but separately have it decide to // retire all of the available connection IDs. - server.test_frame_writer = Some(Box::new(RetireAll { cid_gen })); - let retire_all = send_something(&mut server, now()); - server.test_frame_writer = None; + let retire_all = send_with_extra(&mut server, RetireAll { cid_gen }, now()); let resp = server.process(Some(&probe), now()).dgram().unwrap(); assert_v4_path(&resp, true); @@ -956,13 +950,11 @@ fn error_on_new_path_with_no_connection_id() { let cid_gen: Rc> = Rc::new(RefCell::new(CountingConnectionIdGenerator::default())); - server.test_frame_writer = Some(Box::new(RetireAll { cid_gen })); - let retire_all = send_something(&mut server, now()); + let retire_all = send_with_extra(&mut server, RetireAll { cid_gen }, now()); client.process_input(&retire_all, now()); - server.test_frame_writer = Some(Box::new(GarbageWriter {})); - let garbage = send_something(&mut server, now()); + let garbage = send_with_extra(&mut server, GarbageWriter {}, now()); let dgram = change_path(&garbage, DEFAULT_ADDR_V4); client.process_input(&dgram, now()); diff --git a/neqo-transport/src/connection/tests/mod.rs b/neqo-transport/src/connection/tests/mod.rs index 923daed628..c2cf4db391 100644 --- a/neqo-transport/src/connection/tests/mod.rs +++ b/neqo-transport/src/connection/tests/mod.rs @@ -17,7 +17,7 @@ use neqo_common::{event::Provider, qdebug, qtrace, Datagram, Decoder, Role}; use neqo_crypto::{random, AllowZeroRtt, AuthenticationStatus, ResumptionToken}; use test_fixture::{fixture_init, new_neqo_qlog, now, DEFAULT_ADDR}; -use super::{CloseReason, Connection, ConnectionId, Output, State}; +use super::{test_internal, CloseReason, Connection, ConnectionId, Output, State}; use crate::{ addr_valid::{AddressValidation, ValidateAddress}, cc::CWND_INITIAL_PKTS, @@ -615,6 +615,17 @@ fn send_something(sender: &mut Connection, now: Instant) -> Datagram { send_something_with_modifier(sender, now, Some) } +/// Send something, but add a little something extra into the output. +fn send_with_extra(sender: &mut Connection, writer: W, now: Instant) -> Datagram +where + W: test_internal::FrameWriter + 'static, +{ + sender.test_frame_writer = Some(Box::new(writer)); + let res = send_something_with_modifier(sender, now, Some); + sender.test_frame_writer = None; + res +} + /// Send something on a stream from `sender` through a modifier to `receiver`. /// Return any ACK that might result. fn send_with_modifier_and_receive( diff --git a/neqo-transport/src/connection/tests/recovery.rs b/neqo-transport/src/connection/tests/recovery.rs index 5b50e8be0a..37845c671f 100644 --- a/neqo-transport/src/connection/tests/recovery.rs +++ b/neqo-transport/src/connection/tests/recovery.rs @@ -835,9 +835,10 @@ fn ack_for_unsent() { let mut server = default_server(); connect_force_idle(&mut client, &mut server); - server.test_frame_writer = Some(Box::new(AckforUnsentWriter {})); - let spoofed = server.process_output(now()).dgram().unwrap(); - server.test_frame_writer = None; + let spoofed = server + .test_write_frames(AckforUnsentWriter {}, now()) + .dgram() + .unwrap(); // Now deliver the packet with the spoofed ACK frame client.process_input(&spoofed, now());