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

Code before:
    
    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let req = self.unwrapInboundIn(data)
        
        // check if request is upgrade
        let connectionHeaders = Set(req.headers[canonicalForm: "connection"].map { $0.lowercased() })
        if connectionHeaders.contains("upgrade") {
            let buffer = UpgradeBufferHandler()
            _ = context.channel.pipeline.addHandler(buffer, position: .before(self.httpRequestDecoder))
            self.upgradeState = .pending(req, buffer)
        }
        
        context.fireChannelRead(data)
    }
    
    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
        let res = self.unwrapOutboundIn(data)
        
        // check upgrade
        switch self.upgradeState {
        case .pending(let req, let buffer):
            self.upgradeState = .upgraded
            if res.status == .switchingProtocols, let upgrader = res.upgrader {
                switch upgrader {
                case .webSocket(let configuration):
                    let webSocketUpgrader = NIOWebSocketServerUpgrader(maxFrameSize: configuration.maxFrameSize, automaticErrorHandling: false, shouldUpgrade: { channel, _ in
                        return channel.eventLoop.makeSucceededFuture([:])
                    }, upgradePipelineHandler: { channel, req in
                        return WebSocket.server(on: channel, onUpgrade: configuration.onUpgrade)
                    })

                    var head = HTTPRequestHead(
                        version: req.version,
                        method: req.method,
                        uri: req.url.string
                    )
                    head.headers = req.headers

                    webSocketUpgrader.buildUpgradeResponse(
                        channel: context.channel,
                        upgradeRequest: head,
                        initialResponseHeaders: [:]
                    ).map { headers in
                        res.headers = headers
                        context.write(self.wrapOutboundOut(res), promise: promise)
                    }.flatMap {
                        let handlers: [RemovableChannelHandler] = [self] + self.httpHandlers
                        return .andAllComplete(handlers.map { handler in
                            return context.pipeline.removeHandler(handler)
                        }, on: context.eventLoop)
                    }.flatMap {

Review comment (suggestion):
This `automaticErrorHandling` change is actually from https://github.com/vapor/vapor/pull/2193. I don't think Github supports "stacked" pull requests... These two would conflict if they aren't stacked like I've done.

Code after:
    
    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let req = self.unwrapInboundIn(data)
        
        // check if request is upgrade
        let connectionHeaders = Set(req.headers[canonicalForm: "connection"].map { $0.lowercased() })
        if connectionHeaders.contains("upgrade") {
            let buffer = UpgradeBufferHandler()
            _ = context.channel.pipeline.addHandler(buffer, position: .before(self.httpRequestDecoder))
            self.upgradeState = .pending(req, buffer)
        }
        
        context.fireChannelRead(data)
    }
    
    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
        let res = self.unwrapOutboundIn(data)
        
        // check upgrade
        switch self.upgradeState {
        case .pending(let req, let buffer):
            self.upgradeState = .upgraded
            if res.status == .switchingProtocols, let upgrader = res.upgrader {
                switch upgrader {
                case .webSocket(let maxFrameSize, let onUpgrade):
                    let maxFrameSizeBytes: Int
                    switch maxFrameSize {
                    case .`default`:
                        maxFrameSizeBytes = 1 << 14
                    case .override(let bytes):
                        maxFrameSizeBytes = bytes
                    }
                    let webSocketUpgrader = NIOWebSocketServerUpgrader(maxFrameSize: maxFrameSizeBytes, automaticErrorHandling: false, shouldUpgrade: { channel, _ in
                        return channel.eventLoop.makeSucceededFuture([:])
                    }, upgradePipelineHandler: { channel, req in
                        return WebSocket.server(on: channel, onUpgrade: onUpgrade)
                    })

                    var head = HTTPRequestHead(
                        version: req.version,
                        method: req.method,
                        uri: req.url.string
                    )
                    head.headers = req.headers

                    webSocketUpgrader.buildUpgradeResponse(
                        channel: context.channel,
                        upgradeRequest: head,
                        initialResponseHeaders: [:]
                    ).map { headers in
                        res.headers = headers