SwiftQUIC: Fix variance in Spin Bit testing

Sources/SwiftNetwork/Protocols/BridgeProtocol.swift

@@ -66,6 +66,8 @@ public struct DatagramDrops: Equatable {
     }
 }
 
+public typealias BridgeObserveFirstByteHandler = ((UInt8) -> Void)?
+
 @_spi(Essentials)
 @available(Network 0.1.0, *)
 public struct BridgeDatagramProtocol: NetworkProtocol {
@@ -75,6 +77,7 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
 
     public struct BridgeOptions: PerProtocolOptions {
         public var linkDelay: NetworkDuration = .zero
+        public var observeFirstByteHandler: BridgeObserveFirstByteHandler = nil
         var datagramDrops: DatagramDrops?
 
         init() {}
@@ -96,6 +99,10 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
             self == other
         }
 
+        static public func == (lhs: BridgeOptions, rhs: BridgeOptions) -> Bool {
+            lhs.linkDelay == rhs.linkDelay && lhs.datagramDrops == rhs.datagramDrops
+        }
+
         var isDefault: Bool {
             self == BridgeOptions()
         }
@@ -142,6 +149,7 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
 
         var linkDelay: NetworkDuration = .zero
         var datagramDrops: DatagramDrops? = nil
+        var observeFirstByteHandler: BridgeObserveFirstByteHandler = nil
 
         private var timerSet = false
         func deliverInboundDataAvailableEvent() {
@@ -183,6 +191,7 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
             {
                 self.linkDelay = bridgeOptions.linkDelay
                 self.datagramDrops = bridgeOptions.datagramDrops
+                self.observeFirstByteHandler = bridgeOptions.observeFirstByteHandler
             }
             #endif
 
@@ -275,6 +284,14 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
                 datagrams = remainingDatagrams
             }
             log.datapath("forwarding \(datagrams.count) datagrams to port: \(remotePort)")
+            if let observeFirstByteHandler {
+                datagrams.iterateMutableFrames { frame in
+                    if let bytes = frame.bytes, bytes.byteCount > 0 {
+                        observeFirstByteHandler(bytes[0])
+                    }
+                    return true
+                }
+            }
             remoteInstance.incomingFrames.add(frames: datagrams)
             remoteInstance.deliverInboundDataAvailableEvent()
         }
@@ -323,6 +340,11 @@ extension ProtocolOptions<BridgeDatagramProtocol> {
         set { perProtocolOptions!.linkDelay = newValue }
     }
 
+    public var observeFirstByteHandler: BridgeObserveFirstByteHandler {
+        get { perProtocolOptions!.observeFirstByteHandler }
+        set { perProtocolOptions!.observeFirstByteHandler = newValue }
+    }
+
     public var datagramDrops: DatagramDrops? {
         get { perProtocolOptions!.datagramDrops }
         set { perProtocolOptions!.datagramDrops = newValue }

Tests/SwiftNetworkTests/QUICTestHarness.swift

@@ -138,7 +138,8 @@ final class QUICTestHarness {
         serverDrops: DatagramDrops? = nil,
         timeout: TimeInterval = 5.0,
         clientOptions: ProtocolOptions<QUICProtocol> = QUICProtocol.options(),
-        serverOptions: ProtocolOptions<QUICProtocol> = QUICProtocol.options()
+        serverOptions: ProtocolOptions<QUICProtocol> = QUICProtocol.options(),
+        bridgeObserveFirstByteHandler: BridgeObserveFirstByteHandler = nil
     ) throws(NetworkError) {
         var clientConnected = false
         var serverConnected = false
@@ -167,6 +168,7 @@ final class QUICTestHarness {
 
             let clientBridge = BridgeDatagramProtocol.instance(context: self.context)
             let clientBridgeOptions = BridgeDatagramProtocol.options()
+            clientBridgeOptions.observeFirstByteHandler = bridgeObserveFirstByteHandler
             clientBridgeOptions.setProtocolInstance(clientBridge)
             clientBridgeOptions.linkDelay = clientLinkDelay
             clientBridgeOptions.datagramDrops = clientDrops
@@ -194,6 +196,7 @@ final class QUICTestHarness {
 
             let serverBridge = BridgeDatagramProtocol.instance(context: self.context)
             let serverBridgeOptions = BridgeDatagramProtocol.options()
+            serverBridgeOptions.observeFirstByteHandler = bridgeObserveFirstByteHandler
             serverBridgeOptions.setProtocolInstance(serverBridge)
             serverBridgeOptions.linkDelay = serverLinkDelay
             serverBridgeOptions.datagramDrops = serverDrops
@@ -914,6 +917,7 @@ final class QUICTestHarness {
         extraServerCIDs: [(QUICConnectionID, QUICStatelessResetToken)] = .init(),
         afterHandshake: ((QUICTestHarness) -> Void)? = nil,  // Block to run after handshake is complete
         afterData: ((QUICTestHarness) -> Void)? = nil,  // Block to run after handshake is complete
+        bridgeObserveFirstByteHandler: BridgeObserveFirstByteHandler = nil
     ) {
         // Start with the handshake
         Logger.test.debug("Test phase: Handshake")
@@ -933,7 +937,8 @@ final class QUICTestHarness {
                 serverDrops: serverDrops,
                 timeout: timeout,
                 clientOptions: clientOptions,
-                serverOptions: serverOptions
+                serverOptions: serverOptions,
+                bridgeObserveFirstByteHandler: bridgeObserveFirstByteHandler
             )
         } catch {
             if expectHandshakeError == nil {

Tests/SwiftNetworkTests/SwiftNetworkQUICSpinBitTests.swift

@@ -60,6 +60,11 @@ final class SwiftNetworkQUICSpinBitTests: NetTestCase {
         // it from the server.
 
         var hasSpinBit = true
+        var observedSpinBitValues: Set<Bool> = []
+        let observeFirstByteHandler: BridgeObserveFirstByteHandler = { firstByte in
+            guard (firstByte & 0xC0) == 0x40 else { return }
+            observedSpinBitValues.insert((firstByte & 0x20) != 0)
+        }
         QUICTestHarness().runQUICTest(
             dataBlock: Array("Hello World!".utf8),
             afterHandshake: { harness in
@@ -89,13 +94,17 @@ final class SwiftNetworkQUICSpinBitTests: NetTestCase {
                 let expectation = XCTestExpectation(description: "Wait to validate spin bit")
                 harness.context.async {
                     defer { expectation.fulfill() }
-
-                    let clientSpinBit = harness.state?.clientInstance.currentPath?.spinValue ?? false
-                    let serverSpinBit = harness.state?.serverInstance.currentPath?.spinValue ?? false
-                    XCTAssertFalse(clientSpinBit, "Client should have the spin bit set to false")
-                    XCTAssertTrue(serverSpinBit, "Server should have the spin bit set to true")
+                    XCTAssertFalse(
+                        observedSpinBitValues.isEmpty,
+                        "BridgeDatagramProtocol should have observed short-header packets with spin bit values"
+                    )
+                    XCTAssertTrue(
+                        observedSpinBitValues.contains(true),
+                        "BridgeDatagramProtocol should have observed at least one packet with spin bit set"
+                    )
                 }
-            }
+            },
+            bridgeObserveFirstByteHandler: observeFirstByteHandler
         )
     }