Skip to content

Commit 39af842

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

File tree

6 files changed

+285
-80
lines changed

6 files changed

+285
-80
lines changed

‎Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift‎

Lines changed: 37 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ 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 final class Transaction: Sendable {
2323
let logger: Logger
2424

2525
let request: HTTPClientRequest.Prepared
@@ -28,8 +28,8 @@ import NIOSSL
2828
let preferredEventLoop: EventLoop
2929
let requestOptions: RequestOptions
3030

31-
private let stateLock = NIOLock()
32-
private var state: StateMachine
31+
// TODO: Rewrite using `NIOLoopBound` to avoid the unnecessary lock
32+
private let state: NIOLockedValueBox<StateMachine>
3333

3434
init(
3535
request: HTTPClientRequest.Prepared,
@@ -44,7 +44,7 @@ import NIOSSL
4444
self.logger = logger
4545
self.connectionDeadline = connectionDeadline
4646
self.preferredEventLoop = preferredEventLoop
47-
self.state = StateMachine(responseContinuation)
47+
self.state = NIOLockedValueBox(StateMachine(responseContinuation))
4848
}
4949

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

6363
switch writeAction {
@@ -99,30 +99,33 @@ import NIOSSL
9999

100100
struct BreakTheWriteLoopError: Swift.Error {}
101101

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

109110
case .writeAndWait(let executor):
110111
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
111-
self.state.waitForRequestBodyDemand(continuation: continuation)
112-
self.stateLock.unlock()
112+
self.state.unsafe.withValueAssumingLockIsAcquired({ state in
113+
state.waitForRequestBodyDemand(continuation: continuation)
114+
})
115+
self.state.unsafe.unlock()
113116

114117
executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
115118
}
116119

117120
case .fail:
118-
self.stateLock.unlock()
121+
self.state.unsafe.unlock()
119122
throw BreakTheWriteLoopError()
120123
}
121124
}
122125

123126
private func requestBodyStreamFinished() {
124-
let finishAction = self.stateLock.withLock {
125-
self.state.finishRequestBodyStream()
127+
let finishAction = self.state.withLockedValue { state in
128+
state.finishRequestBodyStream()
126129
}
127130

128131
switch finishAction {
@@ -150,8 +153,8 @@ extension Transaction: HTTPSchedulableRequest {
150153
var requiredEventLoop: EventLoop? { nil }
151154

152155
func requestWasQueued(_ scheduler: HTTPRequestScheduler) {
153-
self.stateLock.withLock {
154-
self.state.requestWasQueued(scheduler)
156+
self.state.withLockedValue { state in
157+
state.requestWasQueued(scheduler)
155158
}
156159
}
157160
}
@@ -165,8 +168,8 @@ extension Transaction: HTTPExecutableRequest {
165168
// MARK: Request
166169

167170
func willExecuteRequest(_ executor: HTTPRequestExecutor) {
168-
let action = self.stateLock.withLock {
169-
self.state.willExecuteRequest(executor)
171+
let action = self.state.withLockedValue { state in
172+
state.willExecuteRequest(executor)
170173
}
171174

172175
switch action {
@@ -183,8 +186,8 @@ extension Transaction: HTTPExecutableRequest {
183186
func requestHeadSent() {}
184187

185188
func resumeRequestBodyStream() {
186-
let action = self.stateLock.withLock {
187-
self.state.resumeRequestBodyStream()
189+
let action = self.state.withLockedValue { state in
190+
state.resumeRequestBodyStream()
188191
}
189192

190193
switch action {
@@ -214,16 +217,16 @@ extension Transaction: HTTPExecutableRequest {
214217
}
215218

216219
func pauseRequestBodyStream() {
217-
self.stateLock.withLock {
218-
self.state.pauseRequestBodyStream()
220+
self.state.withLockedValue { state in
221+
state.pauseRequestBodyStream()
219222
}
220223
}
221224

222225
// MARK: Response
223226

224227
func receiveResponseHead(_ head: HTTPResponseHead) {
225-
let action = self.stateLock.withLock {
226-
self.state.receiveResponseHead(head, delegate: self)
228+
let action = self.state.withLockedValue { state in
229+
state.receiveResponseHead(head, delegate: self)
227230
}
228231

229232
switch action {
@@ -243,8 +246,8 @@ extension Transaction: HTTPExecutableRequest {
243246
}
244247

245248
func receiveResponseBodyParts(_ buffer: CircularBuffer<ByteBuffer>) {
246-
let action = self.stateLock.withLock {
247-
self.state.receiveResponseBodyParts(buffer)
249+
let action = self.state.withLockedValue { state in
250+
state.receiveResponseBodyParts(buffer)
248251
}
249252
switch action {
250253
case .none:
@@ -260,8 +263,8 @@ extension Transaction: HTTPExecutableRequest {
260263
}
261264

262265
func succeedRequest(_ buffer: CircularBuffer<ByteBuffer>?) {
263-
let succeedAction = self.stateLock.withLock {
264-
self.state.succeedRequest(buffer)
266+
let succeedAction = self.state.withLockedValue { state in
267+
state.succeedRequest(buffer)
265268
}
266269
switch succeedAction {
267270
case .finishResponseStream(let source, let finalResponse):
@@ -276,8 +279,8 @@ extension Transaction: HTTPExecutableRequest {
276279
}
277280

278281
func fail(_ error: Error) {
279-
let action = self.stateLock.withLock {
280-
self.state.fail(error)
282+
let action = self.state.withLockedValue { state in
283+
state.fail(error)
281284
}
282285
self.performFailAction(action)
283286
}
@@ -304,8 +307,8 @@ extension Transaction: HTTPExecutableRequest {
304307
}
305308

306309
func deadlineExceeded() {
307-
let action = self.stateLock.withLock {
308-
self.state.deadlineExceeded()
310+
let action = self.state.withLockedValue { state in
311+
state.deadlineExceeded()
309312
}
310313
self.performDeadlineExceededAction(action)
311314
}
@@ -329,8 +332,8 @@ extension Transaction: HTTPExecutableRequest {
329332
extension Transaction: NIOAsyncSequenceProducerDelegate {
330333
@usableFromInline
331334
func produceMore() {
332-
let action = self.stateLock.withLock {
333-
self.state.produceMore()
335+
let action = self.state.withLockedValue { state in
336+
state.produceMore()
334337
}
335338
switch action {
336339
case .none:

‎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: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -49,23 +49,52 @@ extension HTTPClient {
4949
@inlinable
5050
func writeChunks<Bytes: Collection>(of bytes: Bytes, maxChunkSize: Int) -> EventLoopFuture<Void>
5151
where Bytes.Element == UInt8 {
52-
let iterator = UnsafeMutableTransferBox(bytes.chunks(ofCount: maxChunkSize).makeIterator())
53-
guard let chunk = iterator.wrappedValue.next() else {
52+
// `StreamWriter` is has design issues, for example
53+
// - https://github.com/swift-server/async-http-client/issues/194
54+
// - https://github.com/swift-server/async-http-client/issues/264
55+
// - We're not told the EventLoop the task runs on and the user is free to return whatever EL they
56+
// want.
57+
// One important consideration then is that we must lock around the iterator because we could be hopping
58+
// between threads.
59+
let iterator = NIOLockedValueBox(
60+
bytes.chunks(ofCount: maxChunkSize).enumerated().makeIterator()
61+
)
62+
guard let chunk = iterator.withLockedValue({ $0.next() }) else {
5463
return self.write(IOData.byteBuffer(.init()))
5564
}
5665

5766
@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-
}
67+
func writeNextChunk(
68+
_ chunk: (offset: Int, element: ChunksOfCountCollection<Bytes>.Element),
69+
allDone: EventLoopPromise<Void>
70+
) {
71+
if let nextElement = iterator.withLockedValue({ $0.next() }) {
72+
let (index, chunk) = nextElement
73+
self.write(.byteBuffer(ByteBuffer(bytes: chunk))).map {
74+
if (index + 1) % 4 == 0 {
75+
// Let's not stack-overflow if the futures insta-complete which they at least in HTTP/2
76+
// mode.
77+
// Also, we must frequently return to the EventLoop because we may get the pause signal
78+
// from another thread. If we fail to do that promptly, we may balloon our body chunks
79+
// into memory.
80+
allDone.futureResult.eventLoop.execute {
81+
writeNextChunk(nextElement, allDone: allDone)
82+
}
83+
} else {
84+
writeNextChunk(nextElement, allDone: allDone)
85+
}
86+
}.cascadeFailure(to: allDone)
6387
} else {
64-
return self.write(.byteBuffer(ByteBuffer(bytes: chunk)))
88+
self.write(.byteBuffer(ByteBuffer(bytes: chunk.element))).cascade(to: allDone)
6589
}
6690
}
6791

68-
return writeNextChunk(chunk)
92+
// HACK (again, we're not told the right EventLoop): Let's write 0 bytes to make the user tell us...
93+
return self.write(.byteBuffer(ByteBuffer())).flatMapWithEventLoop { (_, loop) in
94+
let allDone = loop.makePromise(of: Void.self)
95+
writeNextChunk(chunk, allDone: allDone)
96+
return allDone.futureResult
97+
}
6998
}
7099
}
71100

‎Sources/AsyncHTTPClient/UnsafeTransfer.swift‎

Lines changed: 0 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1 @@
1-
//===----------------------------------------------------------------------===//
2-
//
3-
// This source file is part of the AsyncHTTPClient open source project
4-
//
5-
// Copyright (c) 2022 Apple Inc. and the AsyncHTTPClient project authors
6-
// Licensed under Apache License v2.0
7-
//
8-
// See LICENSE.txt for license information
9-
// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
10-
//
11-
// SPDX-License-Identifier: Apache-2.0
12-
//
13-
//===----------------------------------------------------------------------===//
141

15-
/// ``UnsafeMutableTransferBox`` can be used to make non-`Sendable` values `Sendable` and mutable.
16-
/// It can be used to capture local mutable values in a `@Sendable` closure and mutate them from within the closure.
17-
/// As the name implies, the usage of this is unsafe because it disables the sendable checking of the compiler and does not add any synchronisation.
18-
@usableFromInline
19-
final class UnsafeMutableTransferBox<Wrapped> {
20-
@usableFromInline
21-
var wrappedValue: Wrapped
22-
23-
@inlinable
24-
init(_ wrappedValue: Wrapped) {
25-
self.wrappedValue = wrappedValue
26-
}
27-
}
28-
29-
extension UnsafeMutableTransferBox: @unchecked Sendable {}

0 commit comments

Comments
 (0)