swift-nio-extras/Sources/NIOExtras/JSONRPCFraming+ContentLengthHeader.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the SwiftNIO open source project
//
// Copyright (c) 2019-2021 Apple Inc. and the SwiftNIO project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//

import NIOCore

extension NIOJSONRPCFraming {
    /// ``ContentLengthHeaderFrameEncoder`` is responsible for emitting JSON-RPC wire protocol with 'Content-Length'
    /// HTTP-like headers as used by for example by LSP (Language Server Protocol).
    public final class ContentLengthHeaderFrameEncoder: ChannelOutboundHandler {
        /// We'll get handed one message through the `Channel` of this type and will encode into `OutboundOut`
        public typealias OutboundIn = ByteBuffer
        /// Outbound data will be encoded into a `ByteBuffer`.
        public typealias OutboundOut = ByteBuffer

        private var scratchBuffer: ByteBuffer!

        public init() {}

        /// Called when this `ChannelHandler` is added to the `ChannelPipeline`.
        ///
        /// - parameters:
        ///     - context: The `ChannelHandlerContext` which this `ChannelHandler` belongs to.
        public func handlerAdded(context: ChannelHandlerContext) {
            self.scratchBuffer = context.channel.allocator.buffer(capacity: 512)
        }

        /// Called to request a write operation.  Writes write protocol header and then the message.
        /// - parameters:
        ///     - context: The `ChannelHandlerContext` which this `ChannelHandler` belongs to.
        ///     - data: The data to write through the `Channel`, wrapped in a `NIOAny`.
        ///     - promise: The `EventLoopPromise` which should be notified once the operation completes, or nil if no notification should take place.
        public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
            let data = self.unwrapOutboundIn(data)
            // Step 1, clear the target buffer (note, we are re-using it so if we get lucky we don't need to
            // allocate at all.
            self.scratchBuffer.clear()

            // Step 2, write the wire protocol for the header.
            self.scratchBuffer.writeStaticString("Content-Length: ")
            self.scratchBuffer.writeString(String(data.readableBytes, radix: 10))
            self.scratchBuffer.writeStaticString("\r\n\r\n")

            // Step 3, send header and the raw message through the `Channel`.
            if data.readableBytes > 0 {
                context.write(self.wrapOutboundOut(self.scratchBuffer), promise: nil)
                context.write(self.wrapOutboundOut(data), promise: promise)
            } else {
                context.write(self.wrapOutboundOut(self.scratchBuffer), promise: promise)
            }
        }
    }

    /// ``ContentLengthHeaderFrameDecoder`` is responsible for parsing JSON-RPC wire protocol with 'Content-Length'
    /// HTTP-like headers as used by for example by LSP (Language Server Protocol).
    public struct ContentLengthHeaderFrameDecoder: ByteToMessageDecoder {
        /// We're emitting one `ByteBuffer` corresponding exactly to one full payload, no headers etc.
        public typealias InboundOut = ByteBuffer

        /// `ContentLengthHeaderFrameDecoder` is a simple state machine.
        private enum State {
            /// Waiting for the end of the header block or a new header field
            case waitingForHeaderNameOrHeaderBlockEnd
            /// Waiting for a header value
            case waitingForHeaderValue(name: String)
            /// Waiting for the payload of a given size.
            case waitingForPayload(length: Int)
        }

        /// A ``DecodingError`` is sent through the pipeline if anything went wrong.
        public enum DecodingError: Error, Equatable {
            /// Missing 'Content-Length' header.
            case missingContentLengthHeader

            /// The value of the 'Content-Length' header was illegal, for example a negative number.
            case illegalContentLengthHeaderValue(String)
        }

        public init() {}

        // We start waiting for a header field (or the end of a header block).
        private var state: State = .waitingForHeaderNameOrHeaderBlockEnd
        private var payloadLength: Int? = nil

        // Finishes a header block, most of the time that's very straighforward but we need to verify a few
        // things here.
        private mutating func processHeaderBlockEnd(context: ChannelHandlerContext) throws -> DecodingState {
            if let payloadLength = self.payloadLength {
                if payloadLength == 0 {
                    // special case, we're not actually waiting for anything if it's 0 bytes...
                    self.state = .waitingForHeaderNameOrHeaderBlockEnd
                    self.payloadLength = nil
                    context.fireChannelRead(self.wrapInboundOut(context.channel.allocator.buffer(capacity: 0)))
                    return .continue
                }
                // cool, let's just shift to the `.waitingForPayload` state and continue.
                self.state = .waitingForPayload(length: payloadLength)
                self.payloadLength = nil
                return .continue
            } else {
                // this happens if we reached the end of the header block but we haven't seen the Content-Length
                // header, that's an error. It will be sent through the `Channel` and decoder won't be called
                // again.
                throw DecodingError.missingContentLengthHeader
            }
        }

        /// Decode the data in the supplied `buffer`.
        /// `decode` will be invoked whenever there is more data available (or if we return `.continue`).
        /// - parameters:
        ///     - context: Calling context.
        ///     - buffer: The data to decode.
        /// - returns: Status describing need for more data or otherwise.
        public mutating func decode(context: ChannelHandlerContext, buffer: inout ByteBuffer) throws -> DecodingState {
            switch self.state {
            case .waitingForHeaderNameOrHeaderBlockEnd:
                // Given that we're waiting for the end of a header block or a new header field, it's sensible to
                // check if this might be the end of the block.
                if buffer.readableBytesView.starts(with: "\r\n".utf8) {
                    buffer.moveReaderIndex(forwardBy: 2) // skip \r\n\r\n
                    return try self.processHeaderBlockEnd(context: context)
                }

                // Given that this is not the end of a header block, it must be a new header field. A new header field
                // must always have a colon (or we don't have enough data).
                if let colonIndex = buffer.readableBytesView.firstIndex(of: UInt8(ascii: ":")) {
                    let headerName = buffer.readString(length: colonIndex - buffer.readableBytesView.startIndex)!
                    buffer.moveReaderIndex(forwardBy: 1) // skip the colon
                    self.state = .waitingForHeaderValue(name: headerName.trimmed().lowercased())
                    return .continue
                }

                return .needMoreData
            case .waitingForHeaderValue(name: let headerName):
                // Cool, we're waiting for a header value (ie. we're after the colon).

                // Let's not bother unless we found the whole thing
                guard let newlineIndex = buffer.readableBytesView.firstIndex(of: UInt8(ascii: "\n")) else {
                    return .needMoreData
                }

                // Is this a header we actually care about?
                if headerName == "content-length" {
                    // Yes, let's parse the int.
                    let headerValue = buffer.readString(length: newlineIndex - buffer.readableBytesView.startIndex + 1)!
                    if let length = UInt32(headerValue.trimmed()) { // anything more than 4GB or negative doesn't make sense
                        self.payloadLength = .init(length)
                    } else {
                        throw DecodingError.illegalContentLengthHeaderValue(headerValue)
                    }
                } else {
                    // Nope, let's just skip over it
                    buffer.moveReaderIndex(forwardBy: newlineIndex - buffer.readableBytesView.startIndex + 1)
                }

                // but in any case, we're now waiting for a new header or the end of the header block again.
                self.state = .waitingForHeaderNameOrHeaderBlockEnd
                return .continue
            case .waitingForPayload(length: let length):
                // That's the easiest case, let's just wait until we have enough data.
                if let payload = buffer.readSlice(length: length) {
                    // Cool, we got enough data, let's go back waiting for a new header block.
                    self.state = .waitingForHeaderNameOrHeaderBlockEnd
                    // And send what we found through the pipeline.
                    context.fireChannelRead(self.wrapInboundOut(payload))
                    return .continue
                } else {
                    return .needMoreData
                }
            }
        }

        /// Decode all remaining data.
        /// Invoked when the `Channel` is being brought down.
        /// Reports error through `ByteToMessageDecoderError.leftoverDataWhenDone` if not all data is consumed.
        /// - parameters:
        ///     - context: Calling context.
        ///     - buffer: Buffer of data to decode.
        ///     - seenEOF: If the end of file has been seen.
        ///     - returns: .needMoreData always as all data should be consumed.
        public mutating func decodeLast(context: ChannelHandlerContext,
                                        buffer: inout ByteBuffer,
                                        seenEOF: Bool) throws -> DecodingState {
            // Last chance to decode anything.
            while try self.decode(context: context, buffer: &buffer) == .continue {}

            if buffer.readableBytes > 0 {
                // Oops, there are leftovers that don't form a full message, we could ignore those but it doesn't hurt to send
                // an error.
                throw ByteToMessageDecoderError.leftoverDataWhenDone(buffer)
            }
            return .needMoreData
        }
    }
}

extension String {
    func trimmed() -> Substring {
        guard let firstElementIndex = self.firstIndex(where: { !$0.isWhitespace }) else {
            return Substring("")
        }

        let lastElementIndex = self.reversed().firstIndex(where: { !$0.isWhitespace })!
        return self[firstElementIndex ..< lastElementIndex.base]
    }
}



#if swift(>=5.6)
@available(*, unavailable)
extension NIOJSONRPCFraming.ContentLengthHeaderFrameDecoder: Sendable {}

@available(*, unavailable)
extension NIOJSONRPCFraming.ContentLengthHeaderFrameEncoder: Sendable {}
#endif