Repository: vapor/vapor
File: Sources/Vapor/HTTP/Server/HTTPServerResponseEncoder.swift (Swift)

Code before:
        context: ChannelHandlerContext,
        handler: HTTPServerResponseEncoder,
        promise: EventLoopPromise<Void>?,
        count: Int?
    ) {
        self.contextBox = .init(context, eventLoop: context.eventLoop)
        self.handlerBox = .init(handler, eventLoop: context.eventLoop)
        self.promise = promise
        self.count = count
        self.currentCount = .init(0)
        self.isComplete = .init(false)
        self.eventLoop = context.eventLoop
    }
    
    func write(_ result: BodyStreamResult, promise: EventLoopPromise<Void>?) {
        switch result {
        case .buffer(let buffer):
            self.contextBox.value.writeAndFlush(self.handlerBox.value.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: promise)
            self.currentCount.wrappingIncrement(by: buffer.readableBytes, ordering: .sequentiallyConsistent)
            if let count = self.count, self.currentCount.load(ordering: .sequentiallyConsistent) > count {
                self.promise?.fail(Error.tooManyBytes)
                promise?.fail(Error.notEnoughBytes)
            }
        case .end:
            self.isComplete.store(true, ordering: .relaxed)
            if let count = self.count, self.currentCount.load(ordering: .sequentiallyConsistent) != count {
                self.promise?.fail(Error.notEnoughBytes)
                promise?.fail(Error.notEnoughBytes)
            }
            self.contextBox.value.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
            self.contextBox.value.writeAndFlush(self.handlerBox.value.wrapOutboundOut(.end(nil)), promise: promise)
            self.promise?.succeed(())
        case .error(let error):
            self.isComplete.store(true, ordering: .relaxed)
            self.contextBox.value.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
            self.contextBox.value.writeAndFlush(self.handlerBox.value.wrapOutboundOut(.end(nil)), promise: promise)
            self.promise?.fail(error)
        }
    }

    deinit {
        assert(self.isComplete.load(ordering: .sequentiallyConsistent), "Response body stream writer deinitialized before .end or .error was sent.")
    }
}


Review comment (bug):
```suggestion
            if let count = self.count, self.currentCount.load(ordering: .sequentiallyConsistent) < count {
```
Preexisting semantic bug - "too many bytes" is not the same as "not enough bytes".

Code after:
        context: ChannelHandlerContext,
        handler: HTTPServerResponseEncoder,
        promise: EventLoopPromise<Void>?,
        count: Int?
    ) {
        self.contextBox = .init(context, eventLoop: context.eventLoop)
        self.handlerBox = .init(handler, eventLoop: context.eventLoop)
        self.promise = promise
        self.count = count
        self.currentCount = .init(0)
        self.isComplete = .init(false)
        self.eventLoop = context.eventLoop
    }
    
    func write(_ result: BodyStreamResult, promise: EventLoopPromise<Void>?) {
        switch result {
        case .buffer(let buffer):
            self.contextBox.value.writeAndFlush(self.handlerBox.value.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: promise)
            if let count = self.count, self.currentCount.wrappingIncrementThenLoad(by: buffer.readableBytes, ordering: .sequentiallyConsistent) > count {
                self.promise?.fail(Error.tooManyBytes)
                promise?.fail(Error.notEnoughBytes)
            }
        case .end:
            self.isComplete.store(true, ordering: .sequentiallyConsistent)
            if let count = self.count, self.currentCount.load(ordering: .sequentiallyConsistent) < count {
                self.promise?.fail(Error.notEnoughBytes)
                promise?.fail(Error.notEnoughBytes)
            }
            self.contextBox.value.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
            self.contextBox.value.writeAndFlush(self.handlerBox.value.wrapOutboundOut(.end(nil)), promise: promise)
            self.promise?.succeed(())
        case .error(let error):
            self.isComplete.store(true, ordering: .relaxed)
            self.contextBox.value.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
            self.contextBox.value.writeAndFlush(self.handlerBox.value.wrapOutboundOut(.end(nil)), promise: promise)
            self.promise?.fail(error)
        }
    }

    deinit {
        assert(self.isComplete.load(ordering: .sequentiallyConsistent), "Response body stream writer deinitialized before .end or .error was sent.")
    }
}