From fb87d5e49e81a0e6ca5670af240442bfa31ccdaa Mon Sep 17 00:00:00 2001 From: hamzahrmalik Date: Wed, 2 Oct 2024 10:54:42 +0100 Subject: [PATCH] Add functions for reading and writing length-prefixed data with customizable 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 --- ...ByteBuffer-binaryEncodedLengthPrefix.swift | 315 +++++++++++++++++ ...yteBuffer-quicBinaryEncodingStrategy.swift | 145 ++++++++ .../Docs.docc/ByteBuffer-lengthPrefix.md | 117 ++++++ Sources/NIOCore/Docs.docc/index.md | 1 + ...BufferBinaryEncodedLengthPrefixTests.swift | 332 ++++++++++++++++++ ...ufferQUICBinaryEncodingStrategyTests.swift | 136 +++++++ 6 files changed, 1046 insertions(+) create mode 100644 Sources/NIOCore/ByteBuffer-binaryEncodedLengthPrefix.swift create mode 100644 Sources/NIOCore/ByteBuffer-quicBinaryEncodingStrategy.swift create mode 100644 Sources/NIOCore/Docs.docc/ByteBuffer-lengthPrefix.md create mode 100644 Tests/NIOCoreTests/ByteBufferBinaryEncodedLengthPrefixTests.swift create mode 100644 Tests/NIOCoreTests/ByteBufferQUICBinaryEncodingStrategyTests.swift diff --git a/Sources/NIOCore/ByteBuffer-binaryEncodedLengthPrefix.swift b/Sources/NIOCore/ByteBuffer-binaryEncodedLengthPrefix.swift new file mode 100644 index 0000000000..88bc56d7ce --- /dev/null +++ b/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 +/// +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( + 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( + _ 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( + _ 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( + 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: 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 ..( + 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) + } +} diff --git a/Sources/NIOCore/ByteBuffer-quicBinaryEncodingStrategy.swift b/Sources/NIOCore/ByteBuffer-quicBinaryEncodingStrategy.swift new file mode 100644 index 0000000000..e5fa983ac1 --- /dev/null +++ b/Sources/NIOCore/ByteBuffer-quicBinaryEncodingStrategy.swift @@ -0,0 +1,145 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +extension ByteBuffer { + /// A ``NIOBinaryIntegerEncodingStrategy`` which encodes bytes as defined in RFC 9000 § 16 + public struct QUICBinaryEncodingStrategy: NIOBinaryIntegerEncodingStrategy { + /// All possible values for how many bytes a QUIC encoded integer can be + public enum IntegerLength: Int, Sendable { + case one = 1 + case two = 2 + case four = 4 + case eight = 8 + } + /// An estimate of the bytes required to write integers using this strategy + public var requiredBytesHint: Int + + /// Note: Prefer to use the APIs directly on ByteBuffer such as ``ByteBuffer/writeEncodedInteger(_:strategy:)`` and pass `.quic` rather than directly initialising an instance of this strategy + /// - Parameter requiredBytesHint: An estimate of the bytes required to write integers using this strategy. This parameter is only relevant if calling ``ByteBuffer/writeLengthPrefixed(strategy:writeData:)`` + @inlinable + public init(requiredBytesHint: IntegerLength) { + self.requiredBytesHint = requiredBytesHint.rawValue + } + + @inlinable + public func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + guard let firstByte = buffer.getInteger(at: buffer.readerIndex, as: UInt8.self) else { + return nil + } + + // Look at the first two bits to work out the length, then read that, mask off the top two bits, and + // extend to integer. + switch firstByte & 0xC0 { + case 0x00: + // Easy case. + buffer.moveReaderIndex(forwardBy: 1) + return IntegerType(firstByte & ~0xC0) + case 0x40: + // Length is two bytes long, read the next one. + return buffer.readInteger(as: UInt16.self).map { IntegerType($0 & ~(0xC0 << 8)) } + case 0x80: + // Length is 4 bytes long. + return buffer.readInteger(as: UInt32.self).map { IntegerType($0 & ~(0xC0 << 24)) } + case 0xC0: + // Length is 8 bytes long. + return buffer.readInteger(as: UInt64.self).map { IntegerType($0 & ~(0xC0 << 56)) } + default: + fatalError("Unreachable") + } + } + + /// Calculates the minimum number of bytes needed to encode an integer using this strategy + /// - Parameter integer: The integer to be encoded + /// - Returns: The number of bytes needed to encode it + public static func bytesNeededForInteger(_ integer: IntegerType) -> Int { + // We must cast the integer to UInt64 here + // Otherwise, an integer can fall through to the default case + // E.g., if someone calls this function with UInt8.max (which is 255), they would not hit the first case (0..<63) + // The second case cannot be represented at all in UInt8, because 16383 is too big + // Swift will end up creating the 16383 literal as 0, and thus we will fall all the way through to the default + switch UInt64(integer) { + case 0..<63: + return 1 + case 0..<16383: + return 2 + case 0..<1_073_741_823: + return 4 + case 0..<4_611_686_018_427_387_903: + return 8 + default: + fatalError("QUIC variable-length integer outside of valid range") + } + } + + @inlinable + public func writeInteger( + _ integer: IntegerType, + to buffer: inout ByteBuffer + ) -> Int { + self.writeInteger(integer, reservedCapacity: 0, to: &buffer) + } + + @inlinable + public func writeInteger( + _ integer: IntegerType, + reservedCapacity: Int, + to buffer: inout ByteBuffer + ) -> Int { + if reservedCapacity > 8 { + fatalError("Reserved space for QUIC encoded integer must be at most 8 bytes") + } + // Use more space than necessary in order to fill the reserved space + // This will avoid a memmove + // If the needed space is more than the reserved, we can't avoid the move + switch max(reservedCapacity, Self.bytesNeededForInteger(integer)) { + case 1: + // Easy, store the value. The top two bits are 0 so we don't need to do any masking. + return buffer.writeInteger(UInt8(truncatingIfNeeded: integer)) + case 2: + // Set the top two bit mask, then write the value. + let value = UInt16(truncatingIfNeeded: integer) | (0x40 << 8) + return buffer.writeInteger(value) + case 4: + // Set the top two bit mask, then write the value. + let value = UInt32(truncatingIfNeeded: integer) | (0x80 << 24) + return buffer.writeInteger(value) + case 8: + // Set the top two bit mask, then write the value. + let value = UInt64(truncatingIfNeeded: integer) | (0xC0 << 56) + return buffer.writeInteger(value) + default: + fatalError("Unreachable") + } + } + } +} + +extension NIOBinaryIntegerEncodingStrategy where Self == ByteBuffer.QUICBinaryEncodingStrategy { + @inlinable + /// Encodes bytes as defined in RFC 9000 § 16 + /// - Parameter requiredBytesHint: An estimate of the bytes required to write integers using this strategy. This parameter is only relevant if calling ``ByteBuffer/writeLengthPrefixed(strategy:writeData:)`` + /// - Returns: An instance of ``ByteBuffer/QUICBinaryEncodingStrategy`` + public static func quic( + requiredBytesHint: ByteBuffer.QUICBinaryEncodingStrategy.IntegerLength + ) -> ByteBuffer.QUICBinaryEncodingStrategy { + ByteBuffer.QUICBinaryEncodingStrategy(requiredBytesHint: requiredBytesHint) + } + + @inlinable + /// Encodes bytes as defined in RFC 9000 § 16 + public static var quic: ByteBuffer.QUICBinaryEncodingStrategy { .quic(requiredBytesHint: .four) } +} diff --git a/Sources/NIOCore/Docs.docc/ByteBuffer-lengthPrefix.md b/Sources/NIOCore/Docs.docc/ByteBuffer-lengthPrefix.md new file mode 100644 index 0000000000..34ba804694 --- /dev/null +++ b/Sources/NIOCore/Docs.docc/ByteBuffer-lengthPrefix.md @@ -0,0 +1,117 @@ +# Writing length-prefixed data in ByteBuffer + +This article explains how to write data prefixed with a length, where the length could be encoded in various ways. + +## Overview + +We often need to write some data prefixed by its length. Sometimes, this may simply be a fixed width integer. But many +protocols encode the length differently, depending on how big it is. For example, the QUIC protocol uses variable-length +integer encodings, in which smaller numbers can be encoded in fewer bytes. + +We have added functions to help with reading and writing data which is prefixed with lengths encoded by various +strategies. + +## ``NIOBinaryIntegerEncodingStrategy`` protocol + +The first building block is a protocol which describes how to encode and decode an integer. + +An implementation of this protocol is needed for any encoding strategy. One example is the ``ByteBuffer/QUICBinaryEncodingStrategy``. + +This protocol only has two requirements which don't have default implementations: + +- `readInteger`: Reads an integer from the `ByteBuffer` using this encoding. Implementations will read as many bytes as + they need to, according to their wire format, and move the reader index accordingly +- `writeInteger`: Write an integer to the `ByteBuffer` using this encoding. Implementations will write as many bytes as + they need to, according to their wire format, and move the writer index accordingly. + +Note that implementations of this protocol need to either: + +- Encode the length of the integer into the integer itself when writing, so it knows how many bytes to read when + reading. This is what QUIC does. +- Always use the same length, e.g. a simple strategy which always writes the integer as a `UInt64`. + +## Extensions on ``ByteBuffer`` + +To provide a more user-friendly API, we have added extensions on `ByteBuffer` for writing integers with a +chosen ``NIOBinaryIntegerEncodingStrategy``. These are ``ByteBuffer/writeEncodedInteger(_:strategy:)`` +and ``ByteBuffer/readEncodedInteger(as:strategy:)``. + +## Reading and writing length-prefixed data + +We added further APIs on ByteBuffer for reading data, strings and buffers which are written with a length prefix. These +APIs first read an integer using a chosen encoding strategy. The integer then dictates how many bytes of data are read +starting from after the integer. + +Similarly, there are APIs which take data, write its length using the provided strategy, and then write the data itself. + +## Writing complex data with a length-prefix + +Consider the scenario where we want to write multiple pieces of data with a length-prefix, but it is difficult or +complex to work out the total length of that data. + +We decided to add the following API to ByteBuffer: + +```swift +/// - 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. +public mutating func writeLengthPrefixed( + strategy: Strategy, + writeData: (_ buffer: inout ByteBuffer) throws -> Int +) rethrows -> Int +``` + +Users could use the function as follows: + +```swift +myBuffer.writeLengthPrefixed(strategy: .quic) { buffer in + buffer.writeString("something") + buffer.writeSomethingComplex(something) +} +``` + +Writing the implementation of `writeLengthPrefixed` presents a challenge. We need to write the length _before_ the +data. But we do not know the length until the data is written. + +Ideally, we would reserve some number of bytes, then call the `writeData` closure, and then go back and write the length +in the reserved space. However, we would not even know how many bytes of space to reserve, because the number of bytes +needed to write an integer will depend on the integer! + +The solution we landed on is the following: + +- Added ``NIOBinaryIntegerEncodingStrategy/requiredBytesHint``. This allows strategies to provide an estimate of how + many bytes they need for encoding a length +- Using this property, reserve the estimated number of bytes +- Call the `writeData` closure to write the data +- Go back to the reserved space to write the length + - If the length ends up needing fewer bytes than we had reserved, shuffle the data back to close the gap + - If the length ends up needing more bytes than we had reserved, shuffle the data forward to make space + +This code will be most performant when the `requiredBytesHint` is exactly correct, because it will avoid needing to +shuffle any bytes. With that in mind, we can actually make one more optimisation: when we call the `writeInteger` function +on a strategy, we can tell the strategy that we have already reserved some number of bytes. Some encoding strategies +will be able to adjust the way they encode such that they can use exactly that many bytes. + +We added the following function to the ``NIOBinaryIntegerEncodingStrategy`` protocol. This is optional to implement, and +will default to simply calling the existing ``NIOBinaryIntegerEncodingStrategy/writeInteger(_:to:)`` function. + +```swift +/// - 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 +``` + +Many strategies will not be able to do anything useful with the additional `reservedCapacity` parameter. For example, in +ASN1, there is only one possible encoding for a given integer. However, some protocols, such as QUIC, do allow less +efficient encodings. E.g. it is valid in QUIC to encode the number `6` using 4 bytes, even though it could be encoded +using just 1. Such encoding strategies need to make a decision here: they can either use the less efficient +encoding (and therefore use more bytes to encode the integer than would otherwise be necessary), or they can use the +more efficient encoding (and therefore suffer a performance penalty as the bytes need to be shuffled). diff --git a/Sources/NIOCore/Docs.docc/index.md b/Sources/NIOCore/Docs.docc/index.md index 7cf9164783..65e595732d 100644 --- a/Sources/NIOCore/Docs.docc/index.md +++ b/Sources/NIOCore/Docs.docc/index.md @@ -14,6 +14,7 @@ More specialized modules provide concrete implementations of many of the abstrac ### Articles - +- ### Event Loops and Event Loop Groups diff --git a/Tests/NIOCoreTests/ByteBufferBinaryEncodedLengthPrefixTests.swift b/Tests/NIOCoreTests/ByteBufferBinaryEncodedLengthPrefixTests.swift new file mode 100644 index 0000000000..6c12625221 --- /dev/null +++ b/Tests/NIOCoreTests/ByteBufferBinaryEncodedLengthPrefixTests.swift @@ -0,0 +1,332 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import NIOCore +import XCTest + +/// A strategy which just writes integers as UInt8. Enforces the integer must be a particular number to aid testing. Forbids reads +struct UInt8WritingTestStrategy: NIOBinaryIntegerEncodingStrategy { + let expectedWrite: Int + + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + XCTFail("This should not be called") + return 1 + } + + func writeInteger(_ integer: IntegerType, to buffer: inout ByteBuffer) -> Int { + XCTAssertEqual(Int(integer), self.expectedWrite) + return buffer.writeInteger(UInt8(integer)) + } + + func writeInteger(_ integer: Int, reservedCapacity: Int, to buffer: inout ByteBuffer) -> Int { + XCTFail("This should not be called") + return 1 + } +} + +// A which reads a single UInt8 for the length. Forbids writes +struct UInt8ReadingTestStrategy: NIOBinaryIntegerEncodingStrategy { + let expectedRead: UInt8 + + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + let value = buffer.readInteger(as: UInt8.self) + XCTAssertEqual(value, self.expectedRead) + return value.flatMap(IntegerType.init) + } + + func writeInteger(_ integer: IntegerType, to buffer: inout ByteBuffer) -> Int { + XCTFail("This should not be called") + return 1 + } + + func writeInteger(_ integer: Int, reservedCapacity: Int, to buffer: inout ByteBuffer) -> Int { + XCTFail("This should not be called") + return 1 + } + + var requiredBytesHint: Int { 1 } +} + +final class ByteBufferBinaryEncodedLengthPrefixTests: XCTestCase { + // MARK: - simple readEncodedInteger and writeEncodedInteger tests + + func testReadWriteEncodedInteger() { + struct TestStrategy: NIOBinaryIntegerEncodingStrategy { + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + 10 + } + + func writeInteger( + _ integer: IntegerType, + to buffer: inout ByteBuffer + ) -> Int { + XCTAssertEqual(integer, 10) + return 1 + } + + func writeInteger( + _ integer: Int, + reservedCapacity: Int, + to buffer: inout ByteBuffer + ) -> Int { + XCTFail("This should not be called") + return 1 + } + } + + // This should just call down to the strategy function + var buffer = ByteBuffer() + XCTAssertEqual(buffer.readEncodedInteger(strategy: TestStrategy()), 10) + XCTAssertEqual(buffer.writeEncodedInteger(10, strategy: TestStrategy()), 1) + } + + // MARK: - writeLengthPrefixed tests + + func testWriteLengthPrefixedFitsInReservedCapacity() { + struct TestStrategy: NIOBinaryIntegerEncodingStrategy { + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: IntegerType, + to buffer: inout ByteBuffer + ) -> Int { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: Int, + reservedCapacity: Int, + to buffer: inout ByteBuffer + ) -> Int { + XCTAssertEqual(Int(integer), 4) + XCTAssertEqual(reservedCapacity, 1) + return buffer.writeInteger(UInt8(integer)) + } + + var requiredBytesHint: Int { 1 } + } + + var buffer = ByteBuffer() + buffer.writeLengthPrefixed(strategy: TestStrategy()) { writer in + writer.writeString("test") + } + + XCTAssertEqual(buffer.readableBytes, 5) + XCTAssertEqual(buffer.readBytes(length: 5), [4] + "test".utf8) + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } + + func testWriteLengthPrefixedNeedsMoreThanReservedCapacity() { + struct TestStrategy: NIOBinaryIntegerEncodingStrategy { + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: IntegerType, + to buffer: inout ByteBuffer + ) -> Int { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: Int, + reservedCapacity: Int, + to buffer: inout ByteBuffer + ) -> Int { + XCTAssertEqual(Int(integer), 4) + XCTAssertEqual(reservedCapacity, 1) + // We use 8 bytes, but only one was reserved + return buffer.writeInteger(UInt64(integer)) + } + + var requiredBytesHint: Int { 1 } + } + + var buffer = ByteBuffer() + buffer.writeLengthPrefixed(strategy: TestStrategy()) { writer in + writer.writeString("test") + } + + // The strategy above uses 8 bytes for encoding the length. The data is 4, making a total of 12 bytes written + XCTAssertEqual(buffer.readableBytes, 12) + XCTAssertEqual(buffer.readBytes(length: 12), [0, 0, 0, 0, 0, 0, 0, 4] + "test".utf8) + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } + + func testWriteLengthPrefixedNeedsLessThanReservedCapacity() { + struct TestStrategy: NIOBinaryIntegerEncodingStrategy { + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: IntegerType, + to buffer: inout ByteBuffer + ) -> Int { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: Int, + reservedCapacity: Int, + to buffer: inout ByteBuffer + ) -> Int { + XCTAssertEqual(Int(integer), 4) + XCTAssertEqual(reservedCapacity, 8) + return buffer.writeInteger(UInt8(integer)) + } + + var requiredBytesHint: Int { 8 } + } + + var buffer = ByteBuffer() + buffer.writeLengthPrefixed(strategy: TestStrategy()) { writer in + writer.writeString("test") + } + + // The strategy above reserves 8 bytes, but only uses 1 + // The implementation will take care of removing the 7 spare bytes for us + XCTAssertEqual(buffer.readableBytes, 5) + XCTAssertEqual(buffer.readBytes(length: 5), [4] + "test".utf8) + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } + + func testWriteLengthPrefixedThrowing() { + // A strategy which fails the test if anything is called + struct NeverCallStrategy: NIOBinaryIntegerEncodingStrategy { + func readInteger( + as: IntegerType.Type, + from buffer: inout ByteBuffer + ) -> IntegerType? { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: IntegerType, + to buffer: inout ByteBuffer + ) -> Int { + XCTFail("This should not be called") + return 1 + } + + func writeInteger( + _ integer: Int, + reservedCapacity: Int, + to buffer: inout ByteBuffer + ) -> Int { + XCTFail("This should not be called") + return 1 + } + + var requiredBytesHint: Int { 1 } + } + + struct TestError: Error {} + + var buffer = ByteBuffer() + do { + try buffer.writeLengthPrefixed(strategy: NeverCallStrategy()) { _ in + throw TestError() + } + XCTFail("Expected call to throw") + } catch { + // Nothing should have happened, buffer should still be empty + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } + } + + // MARK: - readLengthPrefixed tests + + func testReadLengthPrefixedSlice() { + var buffer = ByteBuffer() + buffer.writeBytes([5, 1, 2, 3, 4, 5]) + let slice = buffer.readLengthPrefixedSlice(strategy: UInt8ReadingTestStrategy(expectedRead: 5)) + XCTAssertEqual(slice?.readableBytesView, [1, 2, 3, 4, 5]) + } + + func testReadLengthPrefixedSliceInsufficientBytes() { + var buffer = ByteBuffer() + buffer.writeBytes([5, 1, 2, 3]) // We put a length of 5, followed by only 3 bytes + let slice = buffer.readLengthPrefixedSlice(strategy: UInt8ReadingTestStrategy(expectedRead: 5)) + XCTAssertNil(slice) + // The original buffer reader index should NOT move + XCTAssertEqual(buffer.readableBytesView, [5, 1, 2, 3]) + } + + // MARK: - writeLengthPrefixed* tests + + func testWriteVariableLengthPrefixedString() { + var buffer = ByteBuffer() + let strategy = UInt8WritingTestStrategy(expectedWrite: 11) + let testString = "Hello World" // length = 11 + let bytesWritten = buffer.writeLengthPrefixedString(testString, strategy: strategy) + XCTAssertEqual(bytesWritten, 11 + 1) // we use 1 byte to write the length + + XCTAssertEqual(buffer.readableBytes, 12) + XCTAssertEqual(buffer.readBytes(length: 12), [11] + testString.utf8) + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } + + func testWriteVariableLengthPrefixedBytes() { + var buffer = ByteBuffer() + let strategy = UInt8WritingTestStrategy(expectedWrite: 10) + let testBytes = [UInt8](repeating: 1, count: 10) + let bytesWritten = buffer.writeLengthPrefixedBytes(testBytes, strategy: strategy) + XCTAssertEqual(bytesWritten, 10 + 1) // we use 1 byte to write the length + + XCTAssertEqual(buffer.readableBytes, 11) + XCTAssertEqual(buffer.readBytes(length: 11), [10, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]) + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } + + func testWriteVariableLengthPrefixedBuffer() { + var buffer = ByteBuffer() + let strategy = UInt8WritingTestStrategy(expectedWrite: 4) + let testBuffer = ByteBuffer(string: "test") + let bytesWritten = buffer.writeLengthPrefixedBuffer(testBuffer, strategy: strategy) + XCTAssertEqual(bytesWritten, 4 + 1) // we use 1 byte to write the length + + XCTAssertEqual(buffer.readableBytes, 5) + XCTAssertEqual(buffer.readBytes(length: 5), [4] + "test".utf8) + XCTAssertTrue(buffer.readableBytesView.isEmpty) + } +} diff --git a/Tests/NIOCoreTests/ByteBufferQUICBinaryEncodingStrategyTests.swift b/Tests/NIOCoreTests/ByteBufferQUICBinaryEncodingStrategyTests.swift new file mode 100644 index 0000000000..9c9eeacc6f --- /dev/null +++ b/Tests/NIOCoreTests/ByteBufferQUICBinaryEncodingStrategyTests.swift @@ -0,0 +1,136 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import XCTest + +@testable import NIOCore + +final class ByteBufferQUICBinaryEncodingStrategyTests: XCTestCase { + // MARK: - writeEncodedInteger tests + + func testWriteOneByteQUICVariableLengthInteger() { + // One byte, ie less than 63, just write out as-is + for number in 0..<63 { + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger(number, to: &buffer) + XCTAssertEqual(bytesWritten, 1) + // The number is written exactly as is + XCTAssertEqual(buffer.readInteger(as: UInt8.self), UInt8(number)) + XCTAssertEqual(buffer.readableBytes, 0) + } + } + + func testWriteBigUInt8() { + // This test case specifically tests the scenario where 2 bytes are needed, but the number being written is UInt8. + // A naive implementation of the quic variable length integer encoder might check whether the number is in + // the range of 64..<16383, to determine that it should be written with 2 bytes. + // However, constructing such a range on a UInt8 would actually construct 64..<0, because 16383 can't be represented as UInt8. + // So this test makes sure we didn't make that mistake + let number: UInt8 = .max + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger(number, to: &buffer) + XCTAssertEqual(bytesWritten, 2) + XCTAssertEqual(buffer.readInteger(as: UInt16.self), 0b01000000_11111111) + XCTAssertEqual(buffer.readableBytes, 0) + } + + func testWriteTwoByteQUICVariableLengthInteger() { + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger(0b00111011_10111101, to: &buffer) + XCTAssertEqual(bytesWritten, 2) + // We need to mask the first 2 bits with 01 to indicate this is a 2 byte integer + // Final result 0b01111011_10111101 + XCTAssertEqual(buffer.readInteger(as: UInt16.self), 0b01111011_10111101) + XCTAssertEqual(buffer.readableBytes, 0) + } + + func testWriteFourByteQUICVariableLengthInteger() { + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger(0b00011101_01111111_00111110_01111101, to: &buffer) + XCTAssertEqual(bytesWritten, 4) + // 2 bit mask is 10 for 4 bytes so this becomes 0b10011101_01111111_00111110_01111101 + XCTAssertEqual(buffer.readInteger(as: UInt32.self), 0b10011101_01111111_00111110_01111101) + XCTAssertEqual(buffer.readableBytes, 0) + } + + func testWriteEightByteQUICVariableLengthInteger() { + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger( + 0b00000010_00011001_01111100_01011110_11111111_00010100_11101000_10001100, + to: &buffer + ) + XCTAssertEqual(bytesWritten, 8) + // 2 bit mask is 11 for 8 bytes so this becomes 0b11000010_00011001_01111100_01011110_11111111_00010100_11101000_10001100 + XCTAssertEqual( + buffer.readInteger(as: UInt64.self), + 0b11000010_00011001_01111100_01011110_11111111_00010100_11101000_10001100 + ) + XCTAssertEqual(buffer.readableBytes, 0) + } + + // MARK: - writeEncodedIntegerWithReservedCapacity tests + + func testWriteOneByteQUICVariableLengthIntegerWithTwoBytesReserved() { + // We only need one byte but the encoder will use 2 because we reserved 2 + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger(0b00000001, reservedCapacity: 2, to: &buffer) + XCTAssertEqual(bytesWritten, 2) + XCTAssertEqual(buffer.readInteger(as: UInt16.self), UInt16(0b01000000_00000001)) + XCTAssertEqual(buffer.readableBytes, 0) + } + + func testRoundtripWithReservedCapacity() { + // This test makes sure that a number encoded with more space than necessary can still be decoded as normal + for reservedCapacity in [0, 1, 2, 4, 8] { + for testNumber in [0, 63, 15293, 494_878_333, 151_288_809_941_952_652] { + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + let bytesWritten = strategy.writeInteger( + testNumber, + reservedCapacity: reservedCapacity, + to: &buffer + ) + let minRequiredBytes = ByteBuffer.QUICBinaryEncodingStrategy.bytesNeededForInteger(testNumber) + // If the reserved capacity is higher than the min required, use the reserved number + let expectedUsedBytes = max(minRequiredBytes, reservedCapacity) + XCTAssertEqual(bytesWritten, expectedUsedBytes) + XCTAssertEqual(strategy.readInteger(as: UInt64.self, from: &buffer), UInt64(testNumber)) + XCTAssertEqual(buffer.readableBytes, 0) + } + } + } + + // MARK: - readEncodedInteger tests + + func testReadEmptyQUICVariableLengthInteger() { + var buffer = ByteBuffer() + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + XCTAssertNil(strategy.readInteger(as: Int.self, from: &buffer)) + } + + func testWriteReadQUICVariableLengthInteger() { + let strategy = ByteBuffer.QUICBinaryEncodingStrategy.quic + for integer in [37, 15293, 494_878_333, 151_288_809_941_952_652] { + var buffer = ByteBuffer() + _ = strategy.writeInteger(integer, to: &buffer) + XCTAssertEqual(strategy.readInteger(as: Int.self, from: &buffer), integer) + } + } +}