Skip to content

Commit eb7a122

Browse files
MacOMNIMacOMNI
authored
fix issue 226 (#230)
* fix issue 226 * update test * update p2p
1 parent 161eb10 commit eb7a122

File tree

4 files changed

+128
-24
lines changed

4 files changed

+128
-24
lines changed

‎Networking/Sources/MsQuicSwift/QuicConnection.swift

+8-1
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,12 @@ public final class QuicConnection: Sendable {
124124
}
125125
}
126126

127+
fileprivate func close() {
128+
storage.write { storage in
129+
storage = nil
130+
}
131+
}
132+
127133
public func shutdown(errorCode: QuicErrorCode = .success) throws {
128134
logger.debug("closing connection")
129135
try storage.write { storage in
@@ -250,12 +256,13 @@ private class ConnectionHandle {
250256
}
251257

252258
case QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE:
253-
logger.trace("Shutdown complete")
259+
logger.debug("Shutdown complete")
254260
if let connection {
255261
connection.handler.shutdownComplete(connection)
256262
}
257263
if event.pointee.SHUTDOWN_COMPLETE.AppCloseInProgress == 0 {
258264
// avoid closing twice
265+
connection?.close()
259266
api.call { api in
260267
api.pointee.ConnectionClose(ptr)
261268
}

‎Networking/Sources/Networking/Peer.swift

+9-15
Original file line numberDiff line numberDiff line change
@@ -294,20 +294,17 @@ final class PeerImpl<Handler: StreamHandler>: Sendable {
294294
}
295295

296296
func reconnect(to address: NetAddr, role: PeerRole) throws {
297-
let state = reconnectStates.read { reconnectStates in
297+
var state = reconnectStates.read { reconnectStates in
298298
reconnectStates[address] ?? .init()
299299
}
300300

301301
guard state.attempt < maxRetryAttempts else {
302302
logger.warning("reconnecting to \(address) exceeded max attempts")
303303
return
304304
}
305-
305+
state.applyBackoff()
306306
reconnectStates.write { reconnectStates in
307-
if var state = reconnectStates[address] {
308-
state.applyBackoff()
309-
reconnectStates[address] = state
310-
}
307+
reconnectStates[address] = state
311308
}
312309
Task {
313310
try await Task.sleep(for: .seconds(state.delay))
@@ -336,20 +333,17 @@ final class PeerImpl<Handler: StreamHandler>: Sendable {
336333
}
337334

338335
func reopenUpStream(connection: Connection<Handler>, kind: Handler.PresistentHandler.StreamKind) {
339-
let state = reopenStates.read { states in
336+
var state = reopenStates.read { states in
340337
states[connection.id] ?? .init()
341338
}
342339

343340
guard state.attempt < maxRetryAttempts else {
344341
logger.warning("Reopen attempt for stream \(kind) on connection \(connection.id) exceeded max attempts")
345342
return
346343
}
347-
344+
state.applyBackoff()
348345
reopenStates.write { states in
349-
if var state = states[connection.id] {
350-
state.applyBackoff()
351-
states[connection.id] = state
352-
}
346+
states[connection.id] = state
353347
}
354348

355349
Task {
@@ -557,10 +551,10 @@ private struct PeerEventHandler<Handler: StreamHandler>: QuicEventHandler {
557551
false
558552
case let .transport(status, _):
559553
switch QuicStatusCode(rawValue: status.rawValue) {
560-
case .badCert:
561-
false
554+
case .aborted, .outOfMemory, .connectionTimeout, .unreachable, .bufferTooSmall, .connectionRefused:
555+
true
562556
default:
563-
!status.isSucceeded
557+
status.isSucceeded
564558
}
565559
case let .byPeer(code):
566560
// Do not reconnect if the closure was initiated by the peer.

‎Networking/Tests/NetworkingTests/MockPeerEventTests.swift

+13-2
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,11 @@ import Utils
77

88
final class MockPeerEventTests {
99
final class MockPeerEventHandler: QuicEventHandler {
10+
enum MockPeerAction {
11+
case none
12+
case mockHandshakeFailure
13+
}
14+
1015
enum EventType {
1116
case newConnection(listener: QuicListener, connection: QuicConnection, info: ConnectionInfo)
1217
case shouldOpen(connection: QuicConnection, certificate: Data?)
@@ -18,8 +23,11 @@ final class MockPeerEventTests {
1823
}
1924

2025
let events: ThreadSafeContainer<[EventType]> = .init([])
26+
let mockAction: MockPeerAction
2127

22-
init() {}
28+
init(_ action: MockPeerAction = .none) {
29+
mockAction = action
30+
}
2331

2432
func newConnection(
2533
_ listener: QuicListener, connection: QuicConnection, info: ConnectionInfo
@@ -32,6 +40,9 @@ final class MockPeerEventTests {
3240
}
3341

3442
func shouldOpen(_: QuicConnection, certificate: Data?) -> QuicStatus {
43+
if mockAction == .mockHandshakeFailure {
44+
return .code(.handshakeFailure)
45+
}
3546
guard let certificate else {
3647
return .code(.requiredCert)
3748
}
@@ -169,7 +180,7 @@ final class MockPeerEventTests {
169180
func connected() async throws {
170181
let serverHandler = MockPeerEventHandler()
171182
let clientHandler = MockPeerEventHandler()
172-
let privateKey1 = try Ed25519.SecretKey(from: Data32())
183+
let privateKey1 = try Ed25519.SecretKey(from: Data32.random())
173184
let cert = try generateSelfSignedCertificate(privateKey: privateKey1)
174185
let serverConfiguration = try QuicConfiguration(
175186
registration: registration,

‎Networking/Tests/NetworkingTests/PeerTests.swift

+98-6
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,98 @@ struct PeerTests {
143143
typealias EphemeralHandler = MockEphemeralStreamHandler
144144
}
145145

146+
@Test
147+
func mockHandshakeFailure() async throws {
148+
let mockPeerTest = try MockPeerEventTests()
149+
let serverHandler = MockPeerEventTests.MockPeerEventHandler(
150+
MockPeerEventTests.MockPeerEventHandler.MockPeerAction.mockHandshakeFailure
151+
)
152+
let alpns = [
153+
PeerRole.validator: Alpn(genesisHeader: Data32(), builder: false).data,
154+
PeerRole.builder: Alpn(genesisHeader: Data32(), builder: true).data,
155+
]
156+
let allAlpns = Array(alpns.values)
157+
// Server setup with bad certificate
158+
let serverConfiguration = try QuicConfiguration(
159+
registration: mockPeerTest.registration,
160+
pkcs12: mockPeerTest.certData,
161+
alpns: allAlpns,
162+
client: false,
163+
settings: QuicSettings.defaultSettings
164+
)
165+
166+
let listener = try QuicListener(
167+
handler: serverHandler,
168+
registration: mockPeerTest.registration,
169+
configuration: serverConfiguration,
170+
listenAddress: NetAddr(ipAddress: "127.0.0.1", port: 0)!,
171+
alpns: allAlpns
172+
)
173+
174+
let listenAddress = try listener.listenAddress()
175+
let peer1 = try Peer(
176+
options: PeerOptions<MockStreamHandler>(
177+
role: .validator,
178+
listenAddress: NetAddr(ipAddress: "127.0.0.1", port: 0)!,
179+
genesisHeader: Data32(),
180+
secretKey: Ed25519.SecretKey(from: Data32.random()),
181+
presistentStreamHandler: MockPresentStreamHandler(),
182+
ephemeralStreamHandler: MockEphemeralStreamHandler(),
183+
serverSettings: .defaultSettings,
184+
clientSettings: .defaultSettings
185+
)
186+
)
187+
188+
let connection1 = try peer1.connect(to: listenAddress, role: .validator)
189+
try? await Task.sleep(for: .milliseconds(3000))
190+
#expect(connection1.isClosed == true)
191+
}
192+
193+
@Test
194+
func mockShutdownBadCert() async throws {
195+
let mockPeerTest = try MockPeerEventTests()
196+
let serverHandler = MockPeerEventTests.MockPeerEventHandler()
197+
let alpns = [
198+
PeerRole.validator: Alpn(genesisHeader: Data32(), builder: false).data,
199+
PeerRole.builder: Alpn(genesisHeader: Data32(), builder: true).data,
200+
]
201+
let allAlpns = Array(alpns.values)
202+
// Server setup with bad certificate
203+
let serverConfiguration = try QuicConfiguration(
204+
registration: mockPeerTest.registration,
205+
pkcs12: mockPeerTest.badCertData,
206+
alpns: allAlpns,
207+
client: false,
208+
settings: QuicSettings.defaultSettings
209+
)
210+
211+
let listener = try QuicListener(
212+
handler: serverHandler,
213+
registration: mockPeerTest.registration,
214+
configuration: serverConfiguration,
215+
listenAddress: NetAddr(ipAddress: "127.0.0.1", port: 0)!,
216+
alpns: allAlpns
217+
)
218+
219+
let listenAddress = try listener.listenAddress()
220+
let peer1 = try Peer(
221+
options: PeerOptions<MockStreamHandler>(
222+
role: .validator,
223+
listenAddress: NetAddr(ipAddress: "127.0.0.1", port: 0)!,
224+
genesisHeader: Data32(),
225+
secretKey: Ed25519.SecretKey(from: Data32.random()),
226+
presistentStreamHandler: MockPresentStreamHandler(),
227+
ephemeralStreamHandler: MockEphemeralStreamHandler(),
228+
serverSettings: .defaultSettings,
229+
clientSettings: .defaultSettings
230+
)
231+
)
232+
233+
let connection1 = try peer1.connect(to: listenAddress, role: .validator)
234+
try? await Task.sleep(for: .milliseconds(1000))
235+
#expect(connection1.isClosed == true)
236+
}
237+
146238
@Test
147239
func reopenUpStream() async throws {
148240
let handler2 = MockPresentStreamHandler()
@@ -197,7 +289,7 @@ struct PeerTests {
197289
peer1.broadcast(
198290
kind: .uniqueA, message: .init(kind: .uniqueA, data: messageData)
199291
)
200-
try await Task.sleep(for: .milliseconds(1000))
292+
try await Task.sleep(for: .milliseconds(2000))
201293
let lastReceivedData2 = await handler2.lastReceivedData
202294
#expect(lastReceivedData2 == messageData)
203295
}
@@ -290,15 +382,15 @@ struct PeerTests {
290382

291383
let connection1 = try peer1.connect(to: peer2.listenAddress(), role: .validator)
292384
let connection2 = try peer2.connect(to: peer1.listenAddress(), role: .validator)
293-
try? await Task.sleep(for: .milliseconds(50))
385+
try? await Task.sleep(for: .milliseconds(1000))
294386
if !connection1.isClosed {
295387
let data = try await connection1.request(MockRequest(kind: .typeA, data: Data("hello world".utf8)))
296-
try? await Task.sleep(for: .milliseconds(50))
388+
try? await Task.sleep(for: .milliseconds(500))
297389
#expect(data == Data("hello world response".utf8))
298390
}
299391
if !connection2.isClosed {
300392
let data = try await connection2.request(MockRequest(kind: .typeA, data: Data("hello world".utf8)))
301-
try? await Task.sleep(for: .milliseconds(50))
393+
try? await Task.sleep(for: .milliseconds(500))
302394
#expect(data == Data("hello world response".utf8))
303395
}
304396
}
@@ -573,7 +665,7 @@ struct PeerTests {
573665
to: peer2.listenAddress(), role: .validator
574666
)
575667

576-
try? await Task.sleep(for: .milliseconds(50))
668+
try? await Task.sleep(for: .milliseconds(500))
577669

578670
peer1.broadcast(
579671
kind: .uniqueA, message: .init(kind: .uniqueA, data: Data("hello world".utf8))
@@ -583,7 +675,7 @@ struct PeerTests {
583675
kind: .uniqueB, message: .init(kind: .uniqueB, data: Data("I am jam".utf8))
584676
)
585677
// Verify last received data
586-
try? await Task.sleep(for: .milliseconds(200))
678+
try? await Task.sleep(for: .milliseconds(500))
587679
await #expect(handler2.lastReceivedData == Data("hello world".utf8))
588680
await #expect(handler1.lastReceivedData == Data("I am jam".utf8))
589681
}

0 commit comments

Comments
 (0)