agnosticdev/SpinBitVariance: Added first byte observer to protocol options

Author: agnosticdev

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,7 +149,7 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
 
         var linkDelay: NetworkDuration = .zero
         var datagramDrops: DatagramDrops? = nil
-        var observeFirstByte: ((UInt8) -> Void)? = nil
+        var observeFirstByteHandler: BridgeObserveFirstByteHandler = nil
 
         private var timerSet = false
         func deliverInboundDataAvailableEvent() {
@@ -184,6 +191,7 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
             {
                 self.linkDelay = bridgeOptions.linkDelay
                 self.datagramDrops = bridgeOptions.datagramDrops
+                self.observeFirstByteHandler = bridgeOptions.observeFirstByteHandler
             }
             #endif
 
@@ -276,10 +284,10 @@ public struct BridgeDatagramProtocol: NetworkProtocol {
                 datagrams = remainingDatagrams
             }
             log.datapath("forwarding \(datagrams.count) datagrams to port: \(remotePort)")
-            if let observeFirstByte {
+            if let observeFirstByteHandler {
                 datagrams.iterateMutableFrames { frame in
                     if let bytes = frame.bytes, bytes.byteCount > 0 {
-                        observeFirstByte(bytes[0])
+                        observeFirstByteHandler(bytes[0])
                     }
                     return true
                 }
@@ -332,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

@@ -76,9 +76,6 @@ final class QUICTestHarness {
 
         let clientInstance: QUICProtocol.Instance
         let serverInstance: QUICProtocol.Instance
-
-        let clientBridge: BridgeDatagramProtocol.Instance
-        let serverBridge: BridgeDatagramProtocol.Instance
     }
     var state: QUICHarnessState? = nil
 
@@ -141,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
@@ -168,9 +166,9 @@ final class QUICTestHarness {
             clientOptions.setProtocolInstance(clientReference)
             clientParameters.defaultStack.transport = .quic(clientOptions)
 
-            let clientBridgeInstance = BridgeDatagramProtocol.Instance(context: self.context)
-            let clientBridge = clientBridgeInstance.reference
+            let clientBridge = BridgeDatagramProtocol.instance(context: self.context)
             let clientBridgeOptions = BridgeDatagramProtocol.options()
+            clientBridgeOptions.observeFirstByteHandler = bridgeObserveFirstByteHandler
             clientBridgeOptions.setProtocolInstance(clientBridge)
             clientBridgeOptions.linkDelay = clientLinkDelay
             clientBridgeOptions.datagramDrops = clientDrops
@@ -196,9 +194,9 @@ final class QUICTestHarness {
             serverOptions.setProtocolInstance(serverReference)
             serverParameters.defaultStack.transport = .quic(serverOptions)
 
-            let serverBridgeInstance = BridgeDatagramProtocol.Instance(context: self.context)
-            let serverBridge = serverBridgeInstance.reference
+            let serverBridge = BridgeDatagramProtocol.instance(context: self.context)
             let serverBridgeOptions = BridgeDatagramProtocol.options()
+            serverBridgeOptions.observeFirstByteHandler = bridgeObserveFirstByteHandler
             serverBridgeOptions.setProtocolInstance(serverBridge)
             serverBridgeOptions.linkDelay = serverLinkDelay
             serverBridgeOptions.datagramDrops = serverDrops
@@ -297,9 +295,7 @@ final class QUICTestHarness {
                 clientReference: clientReference,
                 serverReference: serverReference,
                 clientInstance: clientInstance,
-                serverInstance: serverInstance,
-                clientBridge: clientBridgeInstance,
-                serverBridge: serverBridgeInstance
+                serverInstance: serverInstance
             )
 
             clientHarness.waitForError { error in
@@ -550,18 +546,6 @@ final class QUICTestHarness {
         wait(for: [stopCompleteExpectation], timeout: 5.0)
     }
 
-    // Installs a callback on the client and server bridge that fires passing the raw first byte to the
-    // caller for a short header packet.  Used to detect a spin bit passing between the client and server.
-    func observeFirstByteForOutboundShortHeaderPacket(handler: @escaping (_ firstByte: UInt8) -> Void) {
-        guard let state else { return }
-        let observe: ((UInt8) -> Void) = { firstByte in
-            guard (firstByte & 0xC0) == 0x40 else { return }
-            handler(firstByte)
-        }
-        state.clientBridge.observeFirstByte = observe
-        state.serverBridge.observeFirstByte = observe
-    }
-
     func echoDataOnStream(
         dataGenerator: TestDataGenerator,
         streamIndex: Int,
@@ -933,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")
@@ -952,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

@@ -61,6 +61,10 @@ final class SwiftNetworkQUICSpinBitTests: NetTestCase {
 
         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
@@ -79,11 +83,6 @@ final class SwiftNetworkQUICSpinBitTests: NetTestCase {
                         hasSpinBit = false
                         return
                     }
-                    // If we have made it this far we should at least be able to detect at least one spin bit value going back and forth
-                    // Detect at least 1 outbound datagram sending a spin bit value
-                    harness.observeFirstByteForOutboundShortHeaderPacket { firstByte in
-                        observedSpinBitValues.insert((firstByte & 0x20) != 0)
-                    }
                 }
                 self.wait(for: [expectation], timeout: 5.0)
             },
@@ -104,7 +103,8 @@ final class SwiftNetworkQUICSpinBitTests: NetTestCase {
                         "BridgeDatagramProtocol should have observed at least one packet with spin bit set"
                     )
                 }
-            }
+            },
+            bridgeObserveFirstByteHandler: observeFirstByteHandler
         )
     }