diff --git a/Sources/SwiftNetwork/Protocols/BridgeProtocol.swift b/Sources/SwiftNetwork/Protocols/BridgeProtocol.swift index 24658f9..cf73d12 100644 --- a/Sources/SwiftNetwork/Protocols/BridgeProtocol.swift +++ b/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 { 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 } diff --git a/Tests/SwiftNetworkTests/QUICTestHarness.swift b/Tests/SwiftNetworkTests/QUICTestHarness.swift index 42aaba7..e895b40 100644 --- a/Tests/SwiftNetworkTests/QUICTestHarness.swift +++ b/Tests/SwiftNetworkTests/QUICTestHarness.swift @@ -138,7 +138,8 @@ final class QUICTestHarness { serverDrops: DatagramDrops? = nil, timeout: TimeInterval = 5.0, clientOptions: ProtocolOptions = QUICProtocol.options(), - serverOptions: ProtocolOptions = QUICProtocol.options() + serverOptions: ProtocolOptions = 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 { diff --git a/Tests/SwiftNetworkTests/SwiftNetworkQUICSpinBitTests.swift b/Tests/SwiftNetworkTests/SwiftNetworkQUICSpinBitTests.swift index 4243975..af91342 100644 --- a/Tests/SwiftNetworkTests/SwiftNetworkQUICSpinBitTests.swift +++ b/Tests/SwiftNetworkTests/SwiftNetworkQUICSpinBitTests.swift @@ -60,6 +60,11 @@ final class SwiftNetworkQUICSpinBitTests: NetTestCase { // it from the server. var hasSpinBit = true + var observedSpinBitValues: Set = [] + 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 ) }