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

…mizable encodings for the length (#2867)

Add functions for reading and writing length-prefixed data with
customizable encodings for the length, particularly for quic
variable-length integers (RFC 9000)

### Motivation:

Many protocols require us to write data and then prefix that data with
its length. But each protocol has a different way of encoding the
length. This PR introduces general purpose functions which can be
extended for different encoding strategies

### Modifications:

Create a new protocol which defines how to encode an integer
Implement this protocol for QUIC

Provide functions on bytebuffer for writing length-prefixed buffers,
strings or bytes

Sources/NIOCore/ByteBuffer-binaryEncodedLengthPrefix.swift

@@ -0,0 +1,315 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/// Describes a way to encode and decode an integer as bytes.
+/// For more information, see <doc:ByteBuffer-lengthPrefix>
+///
+public protocol NIOBinaryIntegerEncodingStrategy {
+ /// Read an integer from a buffer.
+ /// If there are not enough bytes to read an integer of this encoding, return nil, and do not move the reader index.
+ /// If the the full integer can be read, move the reader index to after the integer, and return the integer.
+ /// - Parameters:
+ /// - as: The type of integer to be read.
+ /// - buffer: The buffer to read from.
+ /// - Returns: The integer that was read, or nil if it was not possible to read it.
+ func readInteger<IntegerType: FixedWidthInteger>(
+ as: IntegerType.Type,
+ from buffer: inout ByteBuffer
+ ) -> IntegerType?
+
+ /// Write an integer to a buffer. Move the writer index to after the written integer.
+ /// - Parameters:
+ /// - integer: The integer to write.
+ /// - buffer: The buffer to write to.
+ /// - Returns: The number of bytes used to write the integer.
+ func writeInteger<IntegerType: FixedWidthInteger>(
+ _ integer: IntegerType,
+ to buffer: inout ByteBuffer
+ ) -> Int
+
+ /// An estimate of the number of bytes required to write integers using this strategy.
+ /// Callers may use this to reserve bytes before writing the integer.
+ /// If the actual bytes used by the write function is more or less than this, it may be necessary to shuffle bytes.
+ /// Therefore, an accurate prediction here will improve performance.
+ /// This function will be called from ``ByteBuffer/writeLengthPrefixed(strategy:writeData:)``
+ var requiredBytesHint: Int { get }
+
+ /// Write an integer to a buffer. Move the writer index to after the written integer.
+ /// This function will be called when an integer needs to be written, and some capacity has already been reserved for it.
+ /// Implementers should consider using a less efficient encoding, if possible, to fit exactly within the reserved capacity.
+ /// Otherwise, the caller will need to shift bytes to reconcile the difference.
+ /// It is up to the implementer to find the balance between performance and size.
+ /// - Parameters:
+ /// - integer: The integer to write
+ /// - reservedCapacity: The capacity already reserved for writing this integer
+ /// - buffer: The buffer to write into.
+ /// - Returns: The number of bytes used to write the integer.
+ func writeInteger(
+ _ integer: Int,
+ reservedCapacity: Int,
+ to buffer: inout ByteBuffer
+ ) -> Int
+}
+
+extension NIOBinaryIntegerEncodingStrategy {
+ @inlinable
+ public var requiredBytesHint: Int { 1 }
+
+ @inlinable
+ public func writeInteger<IntegerType: FixedWidthInteger>(
+ _ integer: IntegerType,
+ reservedCapacity: 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: NIOBinaryIntegerEncodingStrategy, Integer: FixedWidthInteger>(
+ as: Integer.Type = Integer.self,
+ strategy: Strategy
+ ) -> Integer? {
+ strategy.readInteger(as: Integer.self, from: &self)
+ }
+
+ /// Write a binary encoded integer.
+ ///
+ /// - Returns: The number of bytes written.
+ @discardableResult
+ @inlinable
+ public mutating func writeEncodedInteger<
+ Integer: FixedWidthInteger,
+ Strategy: NIOBinaryIntegerEncodingStrategy
+ >(
+ _ value: Integer,
+ strategy: Strategy
+ ) -> Int {
+ strategy.writeInteger(value, to: &self)
+ }
+
+ /// Prefixes bytes written by `writeData` with the number of bytes written.
+ /// The number of bytes written is encoded using `strategy`.
+ ///
+ /// - Note: This function works by reserving the number of bytes suggested by `strategy` 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.
+ /// - writeData: A closure that takes a buffer, writes some data to it, and returns the number of bytes written.
+ /// - Returns: Number of total bytes written. This is the length of the written data + the number of bytes used to write the length before it.
+ @discardableResult
+ @inlinable
+ public mutating func writeLengthPrefixed<Strategy: NIOBinaryIntegerEncodingStrategy>(
+ strategy: Strategy,
+ writeData: (_ buffer: 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 reservedCapacity = strategy.requiredBytesHint
+ self.writeRepeatingByte(0, count: reservedCapacity)
+
+ /// The index at which we start writing the data originally. We may later move the data if the reserved space for the length wasn't right
+ let originalDataStartIndex = self.writerIndex
+ /// The length of the data written
+ let dataLength: Int
+ do {
+ dataLength = try writeData(&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 data originally. We may later move the data if the reserved space for the length wasn't right
+ let originalDataEndIndex = self.writerIndex
+
+ // Quick check to make sure the user didn't do something silly
+ precondition(
+ originalDataEndIndex - originalDataStartIndex == dataLength,
+ "writeData returned \(dataLength) bytes, but actually \(originalDataEndIndex - originalDataStartIndex) bytes were written. They must be the same."
+ )
+
+ // We write the length after the data 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.writeInteger(
+ dataLength,
+ reservedCapacity: reservedCapacity,
+ to: &self
+ )
+
+ switch actualIntegerLength {
+ case reservedCapacity:
+ // Good, exact match, swap the values and then "delete" the trailing bytes by moving the index back
+ self._moveBytes(from: originalDataEndIndex, to: lengthPrefixIndex, size: actualIntegerLength)
+ self.moveWriterIndex(to: originalDataEndIndex)
+ case ..<reservedCapacity:
+ // 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: originalDataEndIndex, to: lengthPrefixIndex, size: actualIntegerLength)
+ let newDataStartIndex = lengthPrefixIndex + actualIntegerLength
+ self._moveBytes(
+ from: originalDataStartIndex,
+ to: newDataStartIndex,
+ size: dataLength
+ )
+ self.moveWriterIndex(to: newDataStartIndex + dataLength)
+ case reservedCapacity...:
+ // 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 - reservedCapacity
+ self._createSpace(before: lengthPrefixIndex, requiredSpace: extraSpaceNeeded)
+
+ // Clean up the indices.
+ let newDataEndIndex = originalDataEndIndex + extraSpaceNeeded
+ // We wrote the length after the data, so we have to move those bytes to the space at the front
+ self._moveBytes(from: newDataEndIndex, to: lengthPrefixIndex, size: actualIntegerLength)
+ self.moveWriterIndex(to: newDataEndIndex)
+ 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.
+ @inlinable
+ public mutating func readLengthPrefixedSlice<Strategy: NIOBinaryIntegerEncodingStrategy>(
+ 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
+ @inlinable
+ public mutating func writeLengthPrefixedBuffer<
+ Strategy: NIOBinaryIntegerEncodingStrategy
+ >(
+ _ buffer: ByteBuffer,
+ strategy: Strategy
+ ) -> Int {
+ self.reserveCapacity(minimumWritableBytes: buffer.readableBytes + strategy.requiredBytesHint)
+ 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
+ @inlinable
+ public mutating func writeLengthPrefixedString<
+ Strategy: NIOBinaryIntegerEncodingStrategy
+ >(
+ _ string: String,
+ strategy: Strategy
+ ) -> Int {
+ // writeString always writes the String as UTF8 bytes, without a null-terminator
+ // So the length will be the utf8 count
+ self.reserveCapacity(minimumWritableBytes: string.utf8.count + strategy.requiredBytesHint)
+ var written = 0
+ 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.
+ @discardableResult
+ @inlinable
+ public mutating func writeLengthPrefixedBytes<
+ Bytes: Sequence,
+ Strategy: NIOBinaryIntegerEncodingStrategy
+ >(
+ _ bytes: Bytes,
+ strategy: Strategy
+ ) -> Int
+ where Bytes.Element == UInt8 {
+ let numberOfBytes = bytes.withContiguousStorageIfAvailable { b in
+ UnsafeRawBufferPointer(b).count
+ }
+ if let numberOfBytes {
+ self.reserveCapacity(minimumWritableBytes: numberOfBytes + strategy.requiredBytesHint)
+ 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) {
+ precondition(index >= self.readerIndex)
+ let bytesToMove = self.writerIndex - index
+
+ // Add the required number of bytes to the end first
+ self.writeRepeatingByte(0, count: requiredSpace)
+ // Move the data forward by that many bytes, to make space at the front
+ // The precondition above makes this safe: our indices are in the valid range, so we can safely use them here
+ try! self.copyBytes(at: index, to: index + requiredSpace, length: 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
+ try! self.copyBytes(at: source, to: destination, length: size)
+ }
+}