nats.swift/Sources/Nats/Extensions/Data+Parser.swift

// Copyright 2024 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import Foundation

extension Data {
    private static let cr = UInt8(ascii: "\r")
    private static let lf = UInt8(ascii: "\n")
    private static let crlf = Data([cr, lf])
    private static var currentNum = 0
    private static var errored = false
    internal static let versionLinePrefix = "NATS/1.0"

    func removePrefix(_ prefix: Data) -> Data {
        guard self.starts(with: prefix) else { return self }
        return self.dropFirst(prefix.count)
    }

    func split(
        separator: Data, maxSplits: Int = .max, omittingEmptySubsequences: Bool = true
    )
        -> [Data]
    {
        var chunks: [Data] = []
        var start = startIndex
        var end = startIndex
        var splitsCount = 0

        while end < count {
            if splitsCount >= maxSplits {
                break
            }
            if self[start..<end].elementsEqual(separator) {
                if !omittingEmptySubsequences || start != end {
                    chunks.append(self[start..<end])
                }
                start = index(end, offsetBy: separator.count)
                end = start
                splitsCount += 1
                continue
            }
            end = index(after: end)
        }

        if start <= endIndex {
            if !omittingEmptySubsequences || start != endIndex {
                chunks.append(self[start..<endIndex])
            }
        }

        return chunks
    }

    func getMessageType() -> NatsOperation? {
        guard self.count > 2 else { return nil }
        for operation in NatsOperation.allOperations() {
            if self.starts(with: operation.rawBytes) {
                return operation
            }
        }
        return nil
    }

    func starts(with bytes: [UInt8]) -> Bool {
        guard self.count >= bytes.count else { return false }
        return self.prefix(bytes.count).elementsEqual(bytes)
    }

    internal mutating func prepend(_ other: Data) {
        self = other + self
    }

    internal func parseOutMessages() throws -> (ops: [ServerOp], remainder: Data?) {
        var serverOps = [ServerOp]()
        var startIndex = self.startIndex
        var remainder: Data?

        while startIndex < self.endIndex {
            var nextLineStartIndex: Int
            var lineData: Data
            if let range = self[startIndex...].range(of: Data.crlf) {
                let lineEndIndex = range.lowerBound
                nextLineStartIndex =
                    self.index(range.upperBound, offsetBy: 0, limitedBy: self.endIndex)
                    ?? self.endIndex
                lineData = self[startIndex..<lineEndIndex]
            } else {
                remainder = self[startIndex..<self.endIndex]
                break
            }
            if lineData.count == 0 {
                startIndex = nextLineStartIndex
                continue
            }

            let serverOp = try ServerOp.parse(from: lineData)

            // if it's a message, get the full payload and add to returned data
            if case .message(var msg) = serverOp {
                if msg.length == 0 {
                    serverOps.append(serverOp)
                } else {
                    var payload = Data()
                    let payloadEndIndex = nextLineStartIndex + msg.length
                    let payloadStartIndex = nextLineStartIndex
                    // include crlf in the expected payload leangth
                    if payloadEndIndex + Data.crlf.count > endIndex {
                        remainder = self[startIndex..<self.endIndex]
                        break
                    }
                    payload.append(self[payloadStartIndex..<payloadEndIndex])
                    msg.payload = payload
                    startIndex =
                        self.index(
                            payloadEndIndex, offsetBy: Data.crlf.count, limitedBy: self.endIndex)
                        ?? self.endIndex
                    serverOps.append(.message(msg))
                    continue
                }
                //TODO(jrm): Add HMSG handling here too.
            } else if case .hMessage(var msg) = serverOp {
                if msg.length == 0 {
                    serverOps.append(serverOp)
                } else {
                    let headersStartIndex = nextLineStartIndex
                    let headersEndIndex = nextLineStartIndex + msg.headersLength
                    let payloadStartIndex = headersEndIndex
                    let payloadEndIndex = nextLineStartIndex + msg.length

                    var payload: Data?
                    if msg.length > msg.headersLength {
                        payload = Data()
                    }
                    var headers = NatsHeaderMap()

                    // if the whole msg length (including training crlf) is longer
                    // than the remaining chunk, break and return the remainder
                    if payloadEndIndex + Data.crlf.count > endIndex {
                        remainder = self[startIndex..<self.endIndex]
                        break
                    }

                    let headersData = self[headersStartIndex..<headersEndIndex]
                    if let headersString = String(data: headersData, encoding: .utf8) {
                        headers = try NatsHeaderMap(from: headersString)
                    }
                    msg.status = headers.status
                    msg.description = headers.description
                    msg.headers = headers

                    if var payload = payload {
                        payload.append(self[payloadStartIndex..<payloadEndIndex])
                        msg.payload = payload
                    }

                    startIndex =
                        self.index(
                            payloadEndIndex, offsetBy: Data.crlf.count, limitedBy: self.endIndex)
                        ?? self.endIndex
                    serverOps.append(.hMessage(msg))
                    continue
                }

            } else {
                // otherwise, just add this server op to the result
                serverOps.append(serverOp)
            }
            startIndex = nextLineStartIndex

        }

        return (serverOps, remainder)
    }
}