test: More Handshake tests

These should pass once the PR series replacing mozilla#1998 has landed.

Broken out of mozilla#1998

neqo-transport/src/connection/tests/handshake.rs

@@ -30,7 +30,7 @@ use super::{
 };
 use crate::{
     connection::{
-        tests::{new_client, new_server},
+        tests::{exchange_ticket, new_client, new_server},
         AddressValidation,
     },
     events::ConnectionEvent,
@@ -1219,6 +1219,44 @@ fn client_initial_retransmits_identical() {
     }
 }
 
+#[test]
+fn client_triggered_zerortt_retransmits_identical() {
+    let mut client = default_client();
+    let mut server = default_server();
+    connect(&mut client, &mut server);
+
+    let token = exchange_ticket(&mut client, &mut server, now());
+    let mut client = default_client();
+    client
+        .enable_resumption(now(), token)
+        .expect("should set token");
+    let mut server = resumed_server(&client);
+
+    // Write 0-RTT before generating any packets.
+    // This should result in a datagram that coalesces Initial and 0-RTT.
+    let client_stream_id = client.stream_create(StreamType::UniDi).unwrap();
+    client.stream_send(client_stream_id, &[1, 2, 3]).unwrap();
+    let client_0rtt = client.process(None, now());
+    assert!(client_0rtt.as_dgram_ref().is_some());
+    let stats1 = client.stats().frame_tx;
+
+    assertions::assert_coalesced_0rtt(&client_0rtt.as_dgram_ref().unwrap()[..]);
+
+    let s1 = server.process(client_0rtt.as_dgram_ref(), now());
+    assert!(s1.as_dgram_ref().is_some()); // Should produce ServerHello etc...
+
+    // Drop the Initial packet from this.
+    let (_, s_hs) = split_datagram(s1.as_dgram_ref().unwrap());
+    assert!(s_hs.is_some());
+
+    // Passing only the server handshake packet to the client should trigger a retransmit.
+    _ = client.process(s_hs.as_ref(), now()).dgram();
+    let stats2 = client.stats().frame_tx;
+    assert_eq!(stats2.all(), stats1.all() * 2);
+    assert_eq!(stats2.crypto, stats1.crypto * 2);
+    assert_eq!(stats2.stream, stats1.stream * 2);
+}
+
 #[test]
 fn server_initial_retransmits_identical() {
     let mut now = now();
@@ -1253,6 +1291,67 @@ fn server_initial_retransmits_identical() {
     }
 }
 
+#[test]
+fn server_triggered_initial_retransmits_identical() {
+    let mut now = now();
+    let mut client = default_client();
+    let mut server = default_server();
+
+    let ci = client.process(None, now).dgram();
+    now += DEFAULT_RTT / 2;
+    let si1 = server.process(ci.as_ref(), now);
+    let stats1 = server.stats().frame_tx;
+
+    // Drop si and wait for client to retransmit.
+    let pto = client.process_output(now).callback();
+    now += pto;
+    let ci = client.process_output(now).dgram();
+
+    // Feed the RTX'ed ci into the server before its PTO fires. The server
+    // should process it and then retransmit its Initial packet including
+    // any coalesced Handshake data.
+    let si2 = server.process(ci.as_ref(), now);
+    let stats2 = server.stats().frame_tx;
+    assert_eq!(si1.dgram().unwrap().len(), si2.dgram().unwrap().len());
+    assert_eq!(stats2.all(), stats1.all() * 2);
+    assert_eq!(stats2.crypto, stats1.crypto * 2);
+    assert_eq!(stats2.ack, stats1.ack * 2);
+}
+
+#[test]
+fn client_handshake_retransmits_identical() {
+    let mut now = now();
+    let mut client = default_client();
+    let mut ci = client.process(None, now).dgram();
+    let mut server = default_server();
+    let mut si = server.process(ci.take().as_ref(), now).dgram();
+
+    now += DEFAULT_RTT;
+
+    _ = client.process(si.take().as_ref(), now).callback();
+    maybe_authenticate(&mut client);
+
+    // Force the client to retransmit its coalesced Handshake/Short packet a number of times and
+    // make sure the retranmissions are identical to the original. Also, verify the PTO
+    // durations.
+    for i in 1..=3 {
+        _ = client.process(None, now).dgram().unwrap();
+        let pto = client.process(None, now).callback();
+        assert_eq!(pto, DEFAULT_RTT * 3 * (1 << (i - 1)));
+        now += pto;
+
+        assert_eq!(
+            client.stats().frame_tx,
+            FrameStats {
+                crypto: i + 1,
+                ack: i + 1,
+                new_connection_id: i * 7,
+                ..Default::default()
+            }
+        );
+    }
+}
+
 #[test]
 fn grease_quic_bit_transport_parameter() {
     fn get_remote_tp(conn: &Connection) -> bool {

neqo-transport/src/connection/tests/zerortt.rs

@@ -6,17 +6,25 @@
 
 use std::{cell::RefCell, rc::Rc, time::Duration};
 
-use neqo_common::{event::Provider, qdebug};
+use neqo_common::{event::Provider, qdebug, Datagram, Decoder, Role};
 use neqo_crypto::{AllowZeroRtt, AntiReplay};
-use test_fixture::{assertions, now};
+use test_fixture::{
+    assertions,
+    header_protection::{
+        apply_header_protection, decode_initial_header, initial_aead_and_hp,
+        remove_header_protection,
+    },
+    now, split_datagram,
+};
 
 use super::{
     super::Connection, connect, default_client, default_server, exchange_ticket, new_server,
     resumed_server, CountingConnectionIdGenerator,
 };
 use crate::{
-    events::ConnectionEvent, ConnectionParameters, Error, StreamType, Version,
-    MIN_INITIAL_PACKET_SIZE,
+    connection::tests::{new_client, DEFAULT_RTT},
+    events::ConnectionEvent,
+    ConnectionParameters, Error, StreamType, Version, MIN_INITIAL_PACKET_SIZE,
 };
 
 #[test]
@@ -320,3 +328,156 @@ fn zero_rtt_loss_accepted() {
         );
     }
 }
+
+#[test]
+#[allow(clippy::too_many_lines)]
+fn zerortt_reorder_frames() {
+    const ACK_FRAME: &[u8] = &[3, 0, 0, 0, 0, 1, 0, 0];
+    const ACK_FRAME_2: &[u8] = &[3, 1, 0, 0, 1, 2, 0, 0];
+
+    let mut client = new_client(
+        ConnectionParameters::default()
+            .versions(Version::Version1, vec![Version::Version1])
+            .grease(false),
+    );
+    let mut server = new_server(
+        ConnectionParameters::default()
+            .versions(Version::Version1, vec![Version::Version1])
+            .grease(false),
+    );
+    let mut now = now();
+    connect(&mut client, &mut server);
+
+    let token = exchange_ticket(&mut client, &mut server, now);
+    let mut client = new_client(
+        ConnectionParameters::default()
+            .versions(Version::Version1, vec![Version::Version1])
+            .grease(false),
+    );
+    client
+        .enable_resumption(now, token)
+        .expect("should set token");
+    let mut server = resumed_server(&client);
+
+    // Write 0-RTT before generating any packets.
+    // This should result in a datagram that coalesces Initial and 0-RTT.
+    let client_stream_id = client.stream_create(StreamType::UniDi).unwrap();
+    client.stream_send(client_stream_id, &[1, 2, 3]).unwrap();
+    let client_0rtt = client.process(None, now);
+    assert!(client_0rtt.as_dgram_ref().is_some());
+    assertions::assert_coalesced_0rtt(&client_0rtt.as_dgram_ref().unwrap()[..]);
+
+    let (_, client_dcid, _, _) =
+        decode_initial_header(client_0rtt.as_dgram_ref().unwrap(), Role::Client).unwrap();
+    let client_dcid = client_dcid.to_owned();
+
+    now += DEFAULT_RTT / 2;
+    let server_hs = server.process(client_0rtt.as_dgram_ref(), now);
+    assert!(server_hs.as_dgram_ref().is_some()); // Should produce ServerHello etc...
+
+    let server_stream_id = server
+        .events()
+        .find_map(|evt| match evt {
+            ConnectionEvent::NewStream { stream_id } => Some(stream_id),
+            _ => None,
+        })
+        .expect("should have received a new stream event");
+    assert_eq!(client_stream_id, server_stream_id.as_u64());
+
+    // Now, only deliver the ACK from the server's Intial packet.
+    let (server_initial, _server_hs) = split_datagram(server_hs.as_dgram_ref().unwrap());
+    let (protected_header, _, _, payload) =
+        decode_initial_header(&server_initial, Role::Server).unwrap();
+
+    // Now decrypt the packet.
+    let (aead, hp) = initial_aead_and_hp(&client_dcid, Role::Server);
+    let (header, pn) = remove_header_protection(&hp, protected_header, payload);
+    assert_eq!(pn, 0);
+    let pn_len = header.len() - protected_header.len();
+    let mut buf = vec![0; payload.len()];
+    let mut plaintext = aead
+        .decrypt(pn, &header, &payload[pn_len..], &mut buf)
+        .unwrap()
+        .to_owned();
+
+    // Now we need to find the frames.  Make some really strong assumptions.
+    let mut dec = Decoder::new(&plaintext[..]);
+    assert_eq!(dec.decode(ACK_FRAME.len()), Some(ACK_FRAME));
+    let end = dec.offset();
+
+    // Remove the CRYPTO frame.
+    plaintext[end..].fill(0);
+
+    // And rebuild a packet.
+    let mut packet = header.clone();
+    packet.resize(MIN_INITIAL_PACKET_SIZE, 0);
+    aead.encrypt(pn, &header, &plaintext, &mut packet[header.len()..])
+        .unwrap();
+    apply_header_protection(&hp, &mut packet, protected_header.len()..header.len());
+    let modified = Datagram::new(
+        server_initial.source(),
+        server_initial.destination(),
+        server_initial.tos(),
+        packet,
+    );
+
+    // Deliver the ACK and make the client RTX.
+    now += DEFAULT_RTT / 2;
+    now += client.process(Some(&modified), now).callback();
+    let client_out = client.process(None, now);
+
+    // The server should get the retransmission.
+    now += DEFAULT_RTT / 2;
+    let server_initial = server.process(client_out.as_dgram_ref(), now);
+    let (server_initial, _) = split_datagram(server_initial.as_dgram_ref().unwrap());
+
+    // Reorder the ACK and CRYPTO frames in the server's Initial packet.
+    let (protected_header, _, _, payload) =
+        decode_initial_header(&server_initial, Role::Server).unwrap();
+
+    // Now decrypt the packet.
+    let (aead, hp) = initial_aead_and_hp(&client_dcid, Role::Server);
+    let (header, pn) = remove_header_protection(&hp, protected_header, payload);
+    assert_eq!(pn, 1);
+    let pn_len = header.len() - protected_header.len();
+    let mut buf = vec![0; payload.len()];
+    let mut plaintext = aead
+        .decrypt(pn, &header, &payload[pn_len..], &mut buf)
+        .unwrap()
+        .to_owned();
+
+    // Now we need to find the frames.  Make some really strong assumptions.
+    let mut dec = Decoder::new(&plaintext[..]);
+    assert_eq!(dec.decode(ACK_FRAME_2.len()), Some(ACK_FRAME_2));
+    assert_eq!(dec.decode_varint(), Some(0x06)); // CRYPTO
+    assert_eq!(dec.decode_varint(), Some(0x00)); // offset
+    dec.skip_vvec(); // Skip over the payload.
+    let end = dec.offset();
+
+    // Move the ACK frame after the CRYPTO frame.
+    plaintext[..end].rotate_left(ACK_FRAME_2.len());
+
+    // And rebuild a packet.
+    let mut packet = header.clone();
+    packet.resize(MIN_INITIAL_PACKET_SIZE, 0);
+    aead.encrypt(pn, &header, &plaintext, &mut packet[header.len()..])
+        .unwrap();
+    apply_header_protection(&hp, &mut packet, protected_header.len()..header.len());
+    let modified = Datagram::new(
+        server_initial.source(),
+        server_initial.destination(),
+        server_initial.tos(),
+        packet,
+    );
+
+    // Deliver the server's Initial (ACK + CRYPTO) after a delay long enough to trigger the
+    // application space PTO.
+    now += DEFAULT_RTT * 5;
+    let probe = client.process(Some(&modified), now).dgram().unwrap();
+    assertions::assert_initial(&probe[..], true);
+
+    now += client.process(None, now).callback();
+    let probe = client.process(None, now).dgram().unwrap();
+    assertions::assert_initial(&probe[..], true);
+    assertions::assert_coalesced_0rtt(&probe[..]);
+}