Add functions for reading and writing length-prefixed data with custo…

…misable encodings for the length

Sources/NIOCore/ByteBuffer-binaryEncodedLengthPrefix.swift

@@ -0,0 +1,313 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2024 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
+//
+//===----------------------------------------------------------------------===//
+
+#if os(Windows)
+import ucrt
+#elseif canImport(Darwin)
+import Darwin
+#elseif canImport(Glibc)
+import Glibc
+#elseif canImport(Musl)
+import Musl
+#elseif canImport(Bionic)
+import Bionic
+#else
+#error("The Byte Buffer module was unable to identify your C library.")
+#endif
+
+/// Describes a way to encode and decode an integer as bytes
+///
+public protocol BinaryIntegerEncodingStrategy {
+ /// Read an integer from a buffer. The reader index should be moved accordingly
+ /// - Returns: The integer
+ func readInteger<IntegerType: FixedWidthInteger>(
+ as: IntegerType.Type,
+ from buffer: inout ByteBuffer
+ ) -> IntegerType?
+
+ /// Write an integer to a buffer. The writer index should be moved accordingly
+ func writeInteger<IntegerType: FixedWidthInteger>(
+ _ integer: IntegerType,
+ to buffer: inout ByteBuffer
+ ) -> Int
+
+ /// When writing an integer using this strategy, how many bytes we should reserve
+ /// If the actual bytes used by the write function is more or less than this, it will be necessary to shuffle bytes
+ /// Therefore, an accurate prediction here will improve performance
+ var reservedSpaceForInteger: Int { get }
+
+ /// Write an integer to a buffer. The writer index should be moved accordingly
+ /// This function takes a `reservedSpace` parameter which is the number of bytes we already reserved for writing this integer
+ /// If you use write more or less than this many bytes, we wll need to move bytes around to make that possible
+ /// Therefore, you may decide to encode differently to use more bytes than you actually would need, if it is possible for you to do so
+ /// That way, you waste the reserved bytes, but improve writing performance
+ func writeIntegerWithReservedSpace(
+ _ integer: Int,
+ reservedSpace: Int,
+ to buffer: inout ByteBuffer
+ ) -> Int
+}
+
+extension BinaryIntegerEncodingStrategy {
+ @inlinable
+ public var reservedSpaceForInteger: Int { 1 }
+
+ @inlinable
+ public func writeIntegerWithReservedSpace<IntegerType: FixedWidthInteger>(
+ _ integer: IntegerType,
+ reservedSpace: Int,
+ to buffer: inout ByteBuffer
+ ) -> Int {
+ self.writeInteger(integer, to: &buffer)
+ }
+}
+
+extension ByteBuffer {
+ /// Read a binary encoded integer, moving the `readerIndex` appropriately
+ /// If there are not enough bytes, nil is returned
+ @inlinable
+ public mutating func readEncodedInteger<Strategy: BinaryIntegerEncodingStrategy>(_ strategy: Strategy) -> Int? {
+ strategy.readInteger(as: Int.self, from: &self)
+ }
+
+ /// Write a binary encoded integer.
+ ///
+ /// - Returns: The number of bytes written
+ @discardableResult
+ @inlinable
+ public mutating func writeEncodedInteger<
+ Integer: FixedWidthInteger,
+ Strategy: BinaryIntegerEncodingStrategy
+ >(
+ _ value: Integer,
+ strategy: Strategy
+ ) -> Int {
+ strategy.writeInteger(value, to: &self)
+ }
+
+ /// Prefixes a message written by `writeMessage` with the number of bytes written using `strategy`
+ ///
+ /// Implementation note: This function works by reserving the number of bytes suggested by ``BinaryIntegerEncodingStrategy.reservedSpace`` before the data
+ /// It then writes the data, and then goes back to write the length
+ /// If the reserved capacity turns out to be too little or too much, then the data will be shifted
+ /// Therefore, this function is most performant if the strategy is able to use the same number of bytes that it reserved
+ ///
+ /// - Parameters:
+ /// - strategy: The strategy to use for encoding the length
+ /// - writeMessage: A closure that takes a buffer, writes a message to it and returns the number of bytes written
+ /// - Returns: Number of total bytes written. This is the length of the written message + the number of bytes used to write the length before it
+ @discardableResult
+ @inlinable
+ public mutating func writeLengthPrefixed<Strategy: BinaryIntegerEncodingStrategy>(
+ strategy: Strategy,
+ writeMessage: (inout ByteBuffer) throws -> Int
+ ) rethrows -> Int {
+ /// The index at which we write the length
+ let lengthPrefixIndex = self.writerIndex
+ /// The space which we reserve for writing the length
+ let reservedSpace = strategy.reservedSpaceForInteger
+ self.writeRepeatingByte(0, count: reservedSpace)
+
+ /// The index at which we start writing the message originally. We may later move the message if the reserved space for the length wasn't right
+ let originalMessageStartIndex = self.writerIndex
+ /// The length of the message written
+ let messageLength: Int
+ do {
+ messageLength = try writeMessage(&self)
+ } catch {
+ // Clean up our write so that it as if we never did it.
+ self.moveWriterIndex(to: lengthPrefixIndex)
+ throw error
+ }
+ /// The index at the end of the written message originally. We may later move the message if the reserved space for the length wasn't right
+ let originalMessageEndIndex = self.writerIndex
+
+ // Quick check to make sure the user didn't do something silly
+ precondition(
+ originalMessageEndIndex - originalMessageStartIndex == messageLength,
+ "writeMessage returned \(messageLength) bytes, but actually \(originalMessageEndIndex - originalMessageStartIndex) bytes were written. They should be the same"
+ )
+
+ // We write the length after the message to begin with. We will move it later
+
+ /// The actual number of bytes used to write the length written. The user may write more or fewer bytes than what we reserved
+ let actualIntegerLength = strategy.writeIntegerWithReservedSpace(
+ messageLength,
+ reservedSpace: reservedSpace,
+ to: &self
+ )
+
+ switch actualIntegerLength {
+ case reservedSpace:
+ // Good, exact match, swap the values and then "delete" the trailing bytes by moving the index back
+ self._moveBytes(from: originalMessageEndIndex, to: lengthPrefixIndex, size: actualIntegerLength)
+ self.moveWriterIndex(to: originalMessageEndIndex)
+ case ..<reservedSpace:
+ // We wrote fewer bytes. We now have to move the length bytes from the end, and
+ // _then_ shrink the rest of the buffer onto it.
+ self._moveBytes(from: originalMessageEndIndex, to: lengthPrefixIndex, size: actualIntegerLength)
+ let newMessageStartIndex = lengthPrefixIndex + actualIntegerLength
+ self._moveBytes(
+ from: originalMessageStartIndex,
+ to: newMessageStartIndex,
+ size: messageLength
+ )
+ self.moveWriterIndex(to: newMessageStartIndex + messageLength)
+ case reservedSpace...:
+ // We wrote more bytes. We now have to create enough space. Once we do, we have the same
+ // implementation as the matching case.
+ let extraSpaceNeeded = actualIntegerLength - reservedSpace
+ self._createSpace(before: lengthPrefixIndex, requiredSpace: extraSpaceNeeded)
+
+ // Clean up the indices.
+ let newMessageEndIndex = originalMessageEndIndex + extraSpaceNeeded
+ // We wrote the length after the message, so we have to move those bytes to the space at the front
+ self._moveBytes(from: newMessageEndIndex, to: lengthPrefixIndex, size: actualIntegerLength)
+ self.moveWriterIndex(to: newMessageEndIndex)
+ default:
+ fatalError("Unreachable")
+ }
+
+ let totalBytesWritten = self.writerIndex - lengthPrefixIndex
+ return totalBytesWritten
+ }
+
+ /// Reads a slice which is prefixed with a length. The length will be read using `strategy`, and then that many bytes will be read to create a slice
+ /// - Returns: The slice, if there are enough bytes to read it fully. In this case, the readerIndex will move to after the slice
+ /// If there are not enough bytes to read the full slice, the readerIndex will stay unchanged
+ public mutating func readLengthPrefixedSlice<Strategy: BinaryIntegerEncodingStrategy>(
+ _ strategy: Strategy
+ ) -> ByteBuffer? {
+ let originalReaderIndex = self.readerIndex
+ guard let length = strategy.readInteger(as: Int.self, from: &self), let slice = self.readSlice(length: length)
+ else {
+ self.moveReaderIndex(to: originalReaderIndex)
+ return nil
+ }
+ return slice
+ }
+}
+
+// MARK: - Helpers for writing length-prefixed things
+
+extension ByteBuffer {
+ /// Write the length of `buffer` using `strategy`. Then write the buffer
+ /// - Parameters:
+ /// - buffer: The buffer to be written
+ /// - strategy: The encoding strategy to use
+ /// - Returns: The total bytes written. This is the bytes needed to write the length, plus the length of the buffer itself
+ @discardableResult
+ public mutating func writeLengthPrefixedBuffer<
+ Strategy: BinaryIntegerEncodingStrategy
+ >(
+ _ buffer: ByteBuffer,
+ strategy: Strategy
+ ) -> Int {
+ var written = 0
+ written += self.writeEncodedInteger(buffer.readableBytes, strategy: strategy)
+ written += self.writeImmutableBuffer(buffer)
+ return written
+ }
+
+ /// Write the length of `string` using `strategy`. Then write the string
+ /// - Parameters:
+ /// - string: The string to be written
+ /// - strategy: The encoding strategy to use
+ /// - Returns: The total bytes written. This is the bytes needed to write the length, plus the length of the string itself
+ @discardableResult
+ public mutating func writeLengthPrefixedString<
+ Strategy: BinaryIntegerEncodingStrategy
+ >(
+ _ string: String,
+ strategy: Strategy
+ ) -> Int {
+ var written = 0
+ // writeString always writes the String as UTF8 bytes, without a null-terminator
+ // So the length will be the utf8 count
+ written += self.writeEncodedInteger(string.utf8.count, strategy: strategy)
+ written += self.writeString(string)
+ return written
+ }
+
+ /// Write the length of `bytes` using `strategy`. Then write the bytes
+ /// - Parameters:
+ /// - bytes: The bytes to be written
+ /// - strategy: The encoding strategy to use
+ /// - Returns: The total bytes written. This is the bytes needed to write the length, plus the length of the bytes themselves
+ @inlinable
+ public mutating func writeLengthPrefixedBytes<
+ Bytes: Sequence,
+ Strategy: BinaryIntegerEncodingStrategy
+ >(
+ _ bytes: Bytes,
+ strategy: Strategy
+ ) -> Int
+ where Bytes.Element == UInt8 {
+ let numberOfBytes = bytes.withContiguousStorageIfAvailable { b in
+ UnsafeRawBufferPointer(b).count
+ }
+ if let numberOfBytes {
+ var written = 0
+ written += self.writeEncodedInteger(numberOfBytes, strategy: strategy)
+ written += self.writeBytes(bytes)
+ return written
+ } else {
+ return self.writeLengthPrefixed(strategy: strategy) { buffer in
+ buffer.writeBytes(bytes)
+ }
+ }
+ }
+}
+
+extension ByteBuffer {
+ /// Creates `requiredSpace` bytes of free space immediately before `index`
+ /// e.g. given [a, b, c, d, e, f, g, h, i, j] and calling this function with (before: 4, requiredSpace: 2) would result in
+ /// [a, b, c, d, 0, 0, e, f, g, h, i, j]
+ /// 2 extra bytes of space were created before index 4 (the letter e)
+ /// The total bytes written will be equal to `requiredSpace`, and the writer index will be moved accordingly
+ @usableFromInline
+ mutating func _createSpace(before index: Int, requiredSpace: Int) {
+ let bytesToMove = self.writerIndex - index
+
+ // Add the required number of bytes to the end first
+ self.writeRepeatingByte(0, count: requiredSpace)
+ // Move the message forward by that many bytes, to make space at the front
+ self.withVeryUnsafeMutableBytes { pointer in
+ _ = memmove(
+ // The new position is forward from where the user written message currently begins
+ pointer.baseAddress!.advanced(by: index + requiredSpace),
+ // This is the position where the user written message currently begins
+ pointer.baseAddress!.advanced(by: index),
+ bytesToMove
+ )
+ }
+ }
+
+ /// Move the `size` bytes starting from `source` to `destination`
+ /// `source` and `destination` must both be within the writable range
+ @usableFromInline
+ mutating func _moveBytes(from source: Int, to destination: Int, size: Int) {
+ precondition(source >= self.readerIndex && destination < self.writerIndex && source >= destination)
+ precondition(source + size <= self.writerIndex)
+
+ // The precondition above makes this safe: our indices are in the valid range, so we can safely use them here
+ self.withVeryUnsafeMutableBytes { pointer in
+ _ = memmove(
+ pointer.baseAddress!.advanced(by: destination),
+ pointer.baseAddress!.advanced(by: source),
+ size
+ )
+ }
+ }
+}