feat: add contains_handshake and assert_contains_handshake helpers (mozilla#3533)

* feat: add contains_handshake and assert_contains_handshake helpers

Add helpers that scan through coalesced long-header packets in a
datagram to find a Handshake packet. This reduces test fragility when
packet sizing causes a small Initial to be coalesced before the
Handshake.

Follows the same scanning pattern as assert_no_1rtt.

Fixes mozilla#3447

* test(transport): migrate pto_handshake_complete to assert_contains_handshake

The two assertions after the server's second handshake burst were
commented out because the datagram can contain a small coalesced
Initial in front of the Handshake packet, which tripped the
first-packet-only check in assert_handshake.

Use the new assert_contains_handshake helper so the check is
deterministic regardless of whether an Initial is coalesced at the
front of the datagram.

---------

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>

Author: mvanhorn

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

@@ -9,7 +9,9 @@ use std::time::{Duration, Instant};
 use neqo_common::qdebug;
 use neqo_crypto::AuthenticationStatus;
 use test_fixture::{
-    assertions::{assert_handshake, assert_initial, is_handshake, is_initial},
+    assertions::{
+        assert_contains_handshake, assert_handshake, assert_initial, is_handshake, is_initial,
+    },
     now, split_datagram,
 };
 
@@ -229,15 +231,15 @@ fn pto_handshake_complete() {
 
     now += HALF_RTT;
     let pkt = server.process(pkt, now).dgram();
-    // This is probably a Handshake packet, but it might also contain
-    // extra Initial data at the start.
-    // assert_handshake(pkt.as_ref().unwrap());
+    // This datagram contains a Handshake packet, but it might also have
+    // extra Initial data coalesced at the start.
+    assert_contains_handshake(pkt.as_ref().unwrap());
 
     now += HALF_RTT;
     let pkt = client.process(pkt, now).dgram();
-    // ...and, if the Initial was sent in that last one,
-    // this will acknowledge it.
-    // assert_handshake(pkt.as_ref().unwrap());
+    // ...and, if the Initial was sent in that last one, this will acknowledge
+    // it, so a coalesced Initial may sit in front of the Handshake.
+    assert_contains_handshake(pkt.as_ref().unwrap());
 
     let cb = client.process_output(now).callback();
     // The client now has a single RTT estimate (20ms), so

test-fixture/src/assertions.rs

@@ -133,6 +133,48 @@ pub fn assert_handshake(payload: &[u8]) {
     assert!(is_handshake(payload));
 }
 
+/// Check if a datagram contains a Handshake packet, scanning through coalesced
+/// long-header packets (e.g., an Initial followed by a Handshake).
+#[must_use]
+pub fn contains_handshake(payload: &[u8]) -> bool {
+    let mut dec = Decoder::from(payload);
+    while let Some(b1) = dec.decode_uint::<u8>() {
+        if payload.iter().skip(dec.offset()).all(|b| *b == 0) {
+            return false;
+        }
+        if b1 & 0x80 == 0 {
+            return false; // Short header, no more long packets.
+        }
+        let Ok(version) = Version::try_from(dec.decode_uint::<version::Wire>().unwrap()) else {
+            return false;
+        };
+        if is_long_packet_type(b1, 0b1010_0000, version) {
+            return true;
+        }
+        dec.skip_vec(1); // DCID
+        dec.skip_vec(1); // SCID
+        let initial_type = if version == Version::Version2 {
+            0b1001_0000
+        } else {
+            0b1000_0000
+        };
+        if (b1 & PACKET_TYPE_MASK) == initial_type {
+            dec.skip_vvec(); // Initial token.
+        }
+        dec.skip_vvec(); // Skip the payload.
+    }
+    false
+}
+
+/// Assert that a datagram contains a Handshake packet, scanning through
+/// coalesced long-header packets.
+pub fn assert_contains_handshake(payload: &[u8]) {
+    assert!(
+        contains_handshake(payload),
+        "expected datagram to contain a Handshake packet"
+    );
+}
+
 pub fn assert_no_1rtt(payload: &[u8]) {
     let mut dec = Decoder::from(payload);
     while let Some(b1) = dec.decode_uint::<u8>() {