Skip to content

Fix data races in NetworkTransport and MockNetworkConnection - #301

Open
varun06 wants to merge 1 commit into
modelcontextprotocol:mainfrom
varun06:network-transport-races
Open

varun06 wants to merge 1 commit into
modelcontextprotocol:mainfrom
varun06:network-transport-races

Conversation

@varun06

@varun06 varun06 commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

Motivation and Context

  • NetworkTransport.attemptConnection() set connection.stateUpdateHandler on the
    actor's executor but cleared it from AsyncStream's onTermination, which can run on
    any thread. The handler is now cleared with defer on the actor after the for await
    loop exits.
  • The test double MockNetworkConnection was @unchecked Sendable with unprotected
    mutable state accessed from both the transport actor and @MainActor tasks. Its state
    is now lock-protected, and handlers are snapshotted before being invoked.

Related upstream work in the same area: #283, #216, #213.

How Has This Been Tested?

swift test passes. (A TSan run of NetworkTransportTests currently aborts on main in Heartbeat parsing — a misaligned load(as: UInt64.self) — which is unrelated to this change.)

Breaking Changes

None.

Types of changes

  • Bug fix (non-breaking change which fixes an issue)

NetworkTransport.attemptConnection() was writing connection.stateUpdateHandler
from two contexts: the actor's executor (setting the handler) and AsyncStream's
onTermination callback (clearing it to nil), which can fire on any thread.
Fix: remove onTermination and use defer to clear the handler on the actor's
executor after the for-await loop exits, regardless of normal return or throw.

MockNetworkConnection was @unchecked Sendable with entirely unprotected mutable
state. The Task { @mainactor in } wrappers were inconsistent: simulateFailure
wrote mockError directly while receive read it on @mainactor, and stateUpdateHandler
was set from NetworkTransport's actor executor while being read/written on @mainactor.
Fix: add NSLock protecting all mutable state; snapshot handler before calling it
to avoid holding the lock during the callback.
@varun06
varun06 force-pushed the network-transport-races branch from 03776c3 to bef2f07 Compare October 21, 2026 19:17
Comment on lines +74 to +88
let (state, error, data): (NWConnection.State, Swift.Error?, Data?) = lock.withLock {
let d = dataToReceive.isEmpty ? nil : dataToReceive.removeFirst()
return (mockState, mockError, d)
}

let data = self.dataToReceive.removeFirst()
completion(data, nil, self.dataToReceive.isEmpty, nil)
if state == .cancelled {
completion(nil, nil, true, NWError.posix(POSIXErrorCode.ECANCELED))
return
}
if let error {
completion(nil, nil, false, error as? NWError)
return
}
let isLast = lock.withLock { dataToReceive.isEmpty }
completion(data ?? Data(), nil, data == nil ? false : isLast, nil)

@andreyvelich andreyvelich Oct 21, 2026 •

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@varun06 This removes a queued item before the cancelled/error checks, so it gets dropped on those paths. isLast is also read under a separate lock. Do both under one lock:

Suggested change
let (state, error, data): (NWConnection.State, Swift.Error?, Data?) = lock.withLock {
let d = dataToReceive.isEmpty ? nil : dataToReceive.removeFirst()
return (mockState, mockError, d)
}
let data = self.dataToReceive.removeFirst()
completion(data, nil, self.dataToReceive.isEmpty, nil)
if state == .cancelled {
completion(nil, nil, true, NWError.posix(POSIXErrorCode.ECANCELED))
return
}
if let error {
completion(nil, nil, false, error as? NWError)
return
}
let isLast = lock.withLock { dataToReceive.isEmpty }
completion(data ?? Data(), nil, data == nil ? false : isLast, nil)
let (state, error, data, isLast): (NWConnection.State, Swift.Error?, Data?, Bool) =
lock.withLock {
guard mockState != .cancelled, mockError == nil, !dataToReceive.isEmpty else {
return (mockState, mockError, nil, false)
}
let data = dataToReceive.removeFirst()
return (mockState, mockError, data, dataToReceive.isEmpty)
}
if state == .cancelled {
completion(nil, nil, true, NWError.posix(POSIXErrorCode.ECANCELED))
return
}
if let error {
completion(nil, nil, false, error as? NWError)
return
}
completion(data ?? Data(), nil, isLast, nil)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants