Skip to content

Commit eb881c4

Browse files
weissiweissi
committed
fix 784: dont crash on huge in-memory bodies
1 parent bdaa3b1 commit eb881c4

File tree

10 files changed

+311
-87
lines changed

10 files changed

+311
-87
lines changed

‎Sources/AsyncHTTPClient/AsyncAwait/Transaction+StateMachine.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,14 @@ extension Transaction {
3434
case finished(error: Error?)
3535
}
3636

37-
fileprivate enum RequestStreamState {
37+
fileprivate enum RequestStreamState: Sendable {
3838
case requestHeadSent
3939
case producing
4040
case paused(continuation: CheckedContinuation<Void, Error>?)
4141
case finished
4242
}
4343

44-
fileprivate enum ResponseStreamState {
44+
fileprivate enum ResponseStreamState: Sendable {
4545
// Waiting for response head. Valid transitions to: streamingBody.
4646
case waitingForResponseHead
4747
// streaming response body. Valid transitions to: finished.

‎Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift‎

Lines changed: 40 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,11 @@ import NIOHTTP1
1919
import NIOSSL
2020

2121
@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
22-
@usableFromInline final class Transaction: @unchecked Sendable {
22+
@usableFromInline
23+
final class Transaction:
24+
// until NIOLockedValueBox learns `sending` because StateMachine cannot be Sendable
25+
@unchecked Sendable
26+
{
2327
let logger: Logger
2428

2529
let request: HTTPClientRequest.Prepared
@@ -28,8 +32,7 @@ import NIOSSL
2832
let preferredEventLoop: EventLoop
2933
let requestOptions: RequestOptions
3034

31-
private let stateLock = NIOLock()
32-
private var state: StateMachine
35+
private let state: NIOLockedValueBox<StateMachine>
3336

3437
init(
3538
request: HTTPClientRequest.Prepared,
@@ -44,7 +47,7 @@ import NIOSSL
4447
self.logger = logger
4548
self.connectionDeadline = connectionDeadline
4649
self.preferredEventLoop = preferredEventLoop
47-
self.state = StateMachine(responseContinuation)
50+
self.state = NIOLockedValueBox(StateMachine(responseContinuation))
4851
}
4952

5053
func cancel() {
@@ -56,8 +59,8 @@ import NIOSSL
5659
private func writeOnceAndOneTimeOnly(byteBuffer: ByteBuffer) {
5760
// This method is synchronously invoked after sending the request head. For this reason we
5861
// can make a number of assumptions, how the state machine will react.
59-
let writeAction = self.stateLock.withLock {
60-
self.state.writeNextRequestPart()
62+
let writeAction = self.state.withLockedValue { state in
63+
state.writeNextRequestPart()
6164
}
6265

6366
switch writeAction {
@@ -99,30 +102,33 @@ import NIOSSL
99102

100103
struct BreakTheWriteLoopError: Swift.Error {}
101104

105+
// FIXME: Refactor this to not use `self.state.unsafe`.
102106
private func writeRequestBodyPart(_ part: ByteBuffer) async throws {
103-
self.stateLock.lock()
104-
switch self.state.writeNextRequestPart() {
107+
self.state.unsafe.lock()
108+
switch self.state.unsafe.withValueAssumingLockIsAcquired({ state in state.writeNextRequestPart() }) {
105109
case .writeAndContinue(let executor):
106-
self.stateLock.unlock()
110+
self.state.unsafe.unlock()
107111
executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
108112

109113
case .writeAndWait(let executor):
110114
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
111-
self.state.waitForRequestBodyDemand(continuation: continuation)
112-
self.stateLock.unlock()
115+
self.state.unsafe.withValueAssumingLockIsAcquired({ state in
116+
state.waitForRequestBodyDemand(continuation: continuation)
117+
})
118+
self.state.unsafe.unlock()
113119

114120
executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
115121
}
116122

117123
case .fail:
118-
self.stateLock.unlock()
124+
self.state.unsafe.unlock()
119125
throw BreakTheWriteLoopError()
120126
}
121127
}
122128

123129
private func requestBodyStreamFinished() {
124-
let finishAction = self.stateLock.withLock {
125-
self.state.finishRequestBodyStream()
130+
let finishAction = self.state.withLockedValue { state in
131+
state.finishRequestBodyStream()
126132
}
127133

128134
switch finishAction {
@@ -150,8 +156,8 @@ extension Transaction: HTTPSchedulableRequest {
150156
var requiredEventLoop: EventLoop? { nil }
151157

152158
func requestWasQueued(_ scheduler: HTTPRequestScheduler) {
153-
self.stateLock.withLock {
154-
self.state.requestWasQueued(scheduler)
159+
self.state.withLockedValue { state in
160+
state.requestWasQueued(scheduler)
155161
}
156162
}
157163
}
@@ -165,8 +171,8 @@ extension Transaction: HTTPExecutableRequest {
165171
// MARK: Request
166172

167173
func willExecuteRequest(_ executor: HTTPRequestExecutor) {
168-
let action = self.stateLock.withLock {
169-
self.state.willExecuteRequest(executor)
174+
let action = self.state.withLockedValue { state in
175+
state.willExecuteRequest(executor)
170176
}
171177

172178
switch action {
@@ -183,8 +189,8 @@ extension Transaction: HTTPExecutableRequest {
183189
func requestHeadSent() {}
184190

185191
func resumeRequestBodyStream() {
186-
let action = self.stateLock.withLock {
187-
self.state.resumeRequestBodyStream()
192+
let action = self.state.withLockedValue { state in
193+
state.resumeRequestBodyStream()
188194
}
189195

190196
switch action {
@@ -214,16 +220,16 @@ extension Transaction: HTTPExecutableRequest {
214220
}
215221

216222
func pauseRequestBodyStream() {
217-
self.stateLock.withLock {
218-
self.state.pauseRequestBodyStream()
223+
self.state.withLockedValue { state in
224+
state.pauseRequestBodyStream()
219225
}
220226
}
221227

222228
// MARK: Response
223229

224230
func receiveResponseHead(_ head: HTTPResponseHead) {
225-
let action = self.stateLock.withLock {
226-
self.state.receiveResponseHead(head, delegate: self)
231+
let action = self.state.withLockedValue { state in
232+
state.receiveResponseHead(head, delegate: self)
227233
}
228234

229235
switch action {
@@ -243,8 +249,8 @@ extension Transaction: HTTPExecutableRequest {
243249
}
244250

245251
func receiveResponseBodyParts(_ buffer: CircularBuffer<ByteBuffer>) {
246-
let action = self.stateLock.withLock {
247-
self.state.receiveResponseBodyParts(buffer)
252+
let action = self.state.withLockedValue { state in
253+
state.receiveResponseBodyParts(buffer)
248254
}
249255
switch action {
250256
case .none:
@@ -260,8 +266,8 @@ extension Transaction: HTTPExecutableRequest {
260266
}
261267

262268
func succeedRequest(_ buffer: CircularBuffer<ByteBuffer>?) {
263-
let succeedAction = self.stateLock.withLock {
264-
self.state.succeedRequest(buffer)
269+
let succeedAction = self.state.withLockedValue { state in
270+
state.succeedRequest(buffer)
265271
}
266272
switch succeedAction {
267273
case .finishResponseStream(let source, let finalResponse):
@@ -276,8 +282,8 @@ extension Transaction: HTTPExecutableRequest {
276282
}
277283

278284
func fail(_ error: Error) {
279-
let action = self.stateLock.withLock {
280-
self.state.fail(error)
285+
let action = self.state.withLockedValue { state in
286+
state.fail(error)
281287
}
282288
self.performFailAction(action)
283289
}
@@ -304,8 +310,8 @@ extension Transaction: HTTPExecutableRequest {
304310
}
305311

306312
func deadlineExceeded() {
307-
let action = self.stateLock.withLock {
308-
self.state.deadlineExceeded()
313+
let action = self.state.withLockedValue { state in
314+
state.deadlineExceeded()
309315
}
310316
self.performDeadlineExceededAction(action)
311317
}
@@ -329,8 +335,8 @@ extension Transaction: HTTPExecutableRequest {
329335
extension Transaction: NIOAsyncSequenceProducerDelegate {
330336
@usableFromInline
331337
func produceMore() {
332-
let action = self.stateLock.withLock {
333-
self.state.produceMore()
338+
let action = self.state.withLockedValue { state in
339+
state.produceMore()
334340
}
335341
switch action {
336342
case .none:

‎Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2ClientRequestHandler.swift‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,9 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {
432432
}
433433
}
434434

435+
@available(*, unavailable)
436+
extension HTTP2ClientRequestHandler: Sendable {}
437+
435438
extension HTTP2ClientRequestHandler: HTTPRequestExecutor {
436439
func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {
437440
if self.eventLoop.inEventLoop {

‎Sources/AsyncHTTPClient/ConnectionPool/HTTPConnectionPool.swift‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,10 @@ protocol HTTPConnectionPoolDelegate {
2121
func connectionPoolDidShutdown(_ pool: HTTPConnectionPool, unclean: Bool)
2222
}
2323

24-
final class HTTPConnectionPool {
24+
final class HTTPConnectionPool:
25+
// TODO: Refactor to use `NIOLockedValueBox` which will allow this to be checked
26+
@unchecked Sendable
27+
{
2528
private let stateLock = NIOLock()
2629
private var _state: StateMachine
2730
/// The connection idle timeout timers. Protected by the stateLock

‎Sources/AsyncHTTPClient/ConnectionPool/HTTPExecutableRequest.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ import NIOSSL
132132
///
133133
/// Use this handle to cancel the request, while it is waiting for a free connection, to execute the request.
134134
/// This protocol is only intended to be implemented by the `HTTPConnectionPool`.
135-
protocol HTTPRequestScheduler {
135+
protocol HTTPRequestScheduler: Sendable {
136136
/// Informs the task queuer that a request has been cancelled.
137137
func cancelRequest(_: HTTPSchedulableRequest)
138138
}

‎Sources/AsyncHTTPClient/HTTPClient.swift‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -222,21 +222,20 @@ public class HTTPClient {
222222
"""
223223
)
224224
}
225-
let errorStorageLock = NIOLock()
226-
let errorStorage: UnsafeMutableTransferBox<Error?> = .init(nil)
225+
let errorStorage: NIOLockedValueBox<Error?> = NIOLockedValueBox(nil)
227226
let continuation = DispatchWorkItem {}
228227
self.shutdown(requiresCleanClose: requiresCleanClose, queue: DispatchQueue(label: "async-http-client.shutdown"))
229228
{ error in
230229
if let error = error {
231-
errorStorageLock.withLock {
232-
errorStorage.wrappedValue = error
230+
errorStorage.withLockedValue { errorStorage in
231+
errorStorage = error
233232
}
234233
}
235234
continuation.perform()
236235
}
237236
continuation.wait()
238-
try errorStorageLock.withLock {
239-
if let error = errorStorage.wrappedValue {
237+
try errorStorage.withLockedValue { errorStorage in
238+
if let error = errorStorage {
240239
throw error
241240
}
242241
}

‎Sources/AsyncHTTPClient/HTTPHandler.swift‎

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -47,25 +47,65 @@ extension HTTPClient {
4747
}
4848

4949
@inlinable
50-
func writeChunks<Bytes: Collection>(of bytes: Bytes, maxChunkSize: Int) -> EventLoopFuture<Void>
51-
where Bytes.Element == UInt8 {
52-
let iterator = UnsafeMutableTransferBox(bytes.chunks(ofCount: maxChunkSize).makeIterator())
53-
guard let chunk = iterator.wrappedValue.next() else {
50+
func writeChunks<Bytes: Collection>(
51+
of bytes: Bytes,
52+
maxChunkSize: Int
53+
) -> EventLoopFuture<Void> where Bytes.Element == UInt8 {
54+
// `StreamWriter` is has design issues, for example
55+
// - https://github.com/swift-server/async-http-client/issues/194
56+
// - https://github.com/swift-server/async-http-client/issues/264
57+
// - We're not told the EventLoop the task runs on and the user is free to return whatever EL they
58+
// want.
59+
// One important consideration then is that we must lock around the iterator because we could be hopping
60+
// between threads.
61+
typealias Iterator = EnumeratedSequence<ChunksOfCountCollection<Bytes>>.Iterator
62+
typealias Chunk = (offset: Int, element: ChunksOfCountCollection<Bytes>.Element)
63+
64+
func makeIteratorAndFirstChunk() -> (
65+
iterator: NIOLockedValueBox<Iterator>,
66+
chunk: Chunk
67+
)? {
68+
var iterator = bytes.chunks(ofCount: maxChunkSize).enumerated().makeIterator()
69+
guard let chunk = iterator.next() else {
70+
return nil
71+
}
72+
73+
return (NIOLockedValueBox(iterator), chunk)
74+
}
75+
76+
guard let (iterator, chunk) = makeIteratorAndFirstChunk() else {
5477
return self.write(IOData.byteBuffer(.init()))
5578
}
5679

5780
@Sendable // can't use closure here as we recursively call ourselves which closures can't do
58-
func writeNextChunk(_ chunk: Bytes.SubSequence) -> EventLoopFuture<Void> {
59-
if let nextChunk = iterator.wrappedValue.next() {
60-
return self.write(.byteBuffer(ByteBuffer(bytes: chunk))).flatMap {
61-
writeNextChunk(nextChunk)
62-
}
81+
func writeNextChunk(_ chunk: Chunk, allDone: EventLoopPromise<Void>) {
82+
if let nextElement = iterator.withLockedValue({ $0.next() }) {
83+
let (index, chunk) = nextElement
84+
self.write(.byteBuffer(ByteBuffer(bytes: chunk))).map {
85+
if (index + 1) % 4 == 0 {
86+
// Let's not stack-overflow if the futures insta-complete which they at least in HTTP/2
87+
// mode.
88+
// Also, we must frequently return to the EventLoop because we may get the pause signal
89+
// from another thread. If we fail to do that promptly, we may balloon our body chunks
90+
// into memory.
91+
allDone.futureResult.eventLoop.execute {
92+
writeNextChunk(nextElement, allDone: allDone)
93+
}
94+
} else {
95+
writeNextChunk(nextElement, allDone: allDone)
96+
}
97+
}.cascadeFailure(to: allDone)
6398
} else {
64-
return self.write(.byteBuffer(ByteBuffer(bytes: chunk)))
99+
self.write(.byteBuffer(ByteBuffer(bytes: chunk.element))).cascade(to: allDone)
65100
}
66101
}
67102

68-
return writeNextChunk(chunk)
103+
// HACK (again, we're not told the right EventLoop): Let's write 0 bytes to make the user tell us...
104+
return self.write(.byteBuffer(ByteBuffer())).flatMapWithEventLoop { (_, loop) in
105+
let allDone = loop.makePromise(of: Void.self)
106+
writeNextChunk(chunk, allDone: allDone)
107+
return allDone.futureResult
108+
}
69109
}
70110
}
71111

‎Sources/AsyncHTTPClient/UnsafeTransfer.swift‎

Lines changed: 0 additions & 29 deletions
This file was deleted.

0 commit comments

Comments
 (0)