Add functions for reading and writing quic variable-length encoded in…

…tegers

Sources/NIOCore/ByteBuffer-QUICLengthPrefix.swift

@@ -0,0 +1,171 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 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
+//
+//===----------------------------------------------------------------------===//
+
+#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
+
+extension ByteBuffer {
+    /// Read a QUIC variable-length integer, moving the `readerIndex` appropriately
+    /// If there are not enough bytes, nil is returned
+    ///
+    /// [RFC 9000 § 16](https://www.rfc-editor.org/rfc/rfc9000.html#section-16)
+    public mutating func readQUICVariableLengthInteger() -> Int? {
+        guard let firstByte = self.getInteger(at: self.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.
+            self.moveReaderIndex(forwardBy: 1)
+            return Int(firstByte & ~0xC0)
+        case 0x40:
+            // Length is two bytes long, read the next one.
+            return self.readInteger(as: UInt16.self).map { Int($0 & ~(0xC0 << 8)) }
+        case 0x80:
+            // Length is 4 bytes long.
+            return self.readInteger(as: UInt32.self).map { Int($0 & ~(0xC0 << 24)) }
+        case 0xC0:
+            // Length is 8 bytes long.
+            return self.readInteger(as: UInt64.self).map { Int($0 & ~(0xC0 << 56)) }
+        default:
+            fatalError("Unreachable")
+        }
+    }
+
+    /// Write a QUIC variable-length integer.
+    ///
+    /// [RFC 9000 § 16](https://www.rfc-editor.org/rfc/rfc9000.html#section-16)
+    @discardableResult
+    public mutating func writeQUICVariableLengthInteger(_ value: Int) -> Int {
+        self.writeQUICVariableLengthInteger(UInt(value))
+    }
+
+    /// Write a QUIC variable-length integer.
+    ///
+    /// [RFC 9000 § 16](https://www.rfc-editor.org/rfc/rfc9000.html#section-16)
+    ///
+    /// - Returns: The number of bytes written
+    @discardableResult
+    public mutating func writeQUICVariableLengthInteger(_ value: UInt) -> Int {
+        switch value {
+        case 0..<63:
+            // Easy, store the value. The top two bits are 0 so we don't need to do any masking.
+            return self.writeInteger(UInt8(truncatingIfNeeded: value))
+        case 0..<16383:
+            // Set the top two bit mask, then write the value.
+            let value = UInt16(truncatingIfNeeded: value) | (0x40 << 8)
+            return self.writeInteger(value)
+        case 0..<1_073_741_823:
+            // Set the top two bit mask, then write the value.
+            let value = UInt32(truncatingIfNeeded: value) | (0x80 << 24)
+            return self.writeInteger(value)
+        case 0..<4_611_686_018_427_387_903:
+            // Set the top two bit mask, then write the value.
+            let value = UInt64(truncatingIfNeeded: value) | (0xC0 << 56)
+            return self.writeInteger(value)
+        default:
+            fatalError("Could not write QUIC variable-length integer: outside of valid range")
+        }
+    }
+
+    /// Write a QUIC variable-length integer prefixed buffer.
+    ///
+    /// [RFC 9000 § 16](https://www.rfc-editor.org/rfc/rfc9000.html#section-16)
+    ///
+    /// Write `buffer` prefixed with the length of buffer as a QUIC variable-length integer
+    /// - Parameter buffer: The buffer to be written
+    /// - 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 writeQUICLengthPrefixed(_ buffer: ByteBuffer) -> Int {
+        var written = 0
+        written += self.writeQUICVariableLengthInteger(buffer.readableBytes)
+        written += self.writeImmutableBuffer(buffer)
+        return written
+    }
+
+    /// Prefixes a message written by `writeMessage` with the number of bytes written as a QUIC variable-length integer
+    /// - Parameters:
+    ///     - 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
+    /// - Note: Because the length of the message is not known upfront, 4 bytes will be used for encoding the length even if that may not be necessary. If you know the length of your message, prefer ``writeQUICLengthPrefixed(_:)`` instead
+    @discardableResult
+    @inlinable
+    public mutating func writeQUICLengthPrefixed(
+        writeMessage: (inout ByteBuffer) throws -> Int
+    ) rethrows -> Int {
+        var totalBytesWritten = 0
+
+        let lengthPrefixIndex = self.writerIndex
+        // Write 4 bytes as a placeholder
+        // This is enough for upto 1gb of data
+        // It is very unlikely someone is trying to write more than that, if they are then we will catch it later and do a memmove
+        // Reserving 8 bytes now would require either wasting bytes for the majority of use cases, or doing a memmove for a majority of usecases
+        // This way the majority only waste 0-3 bytes, and never incur a memmove, and the minority get a memmove
+        totalBytesWritten += self.writeBytes([0, 0, 0, 0])
+
+        let startWriterIndex = self.writerIndex
+        let messageLength = try writeMessage(&self)
+        let endWriterIndex = self.writerIndex
+
+        totalBytesWritten += messageLength
+
+        let actualBytesWritten = endWriterIndex - startWriterIndex
+        assert(
+            actualBytesWritten == messageLength,
+            "writeMessage returned \(messageLength) bytes, but actually \(actualBytesWritten) bytes were written, but they should be the same"
+        )
+
+        if messageLength >= 1_073_741_823 {
+            // This message length cannot fit in the 4 bytes which we reserved. We need to make more space
+            // Reserve 4 more bytes
+            totalBytesWritten += self.writeBytes([0, 0, 0, 0])
+            // Move the message forward by 4 bytes
+            self.withUnsafeMutableReadableBytes { pointer in
+                _ = memmove(
+                    // The new position is 4 forward from where the user written message currently begins
+                    pointer.baseAddress!.advanced(by: startWriterIndex + 4),
+                    // This is the position where the user written message currently begins
+                    pointer.baseAddress?.advanced(by: startWriterIndex),
+                    messageLength
+                )
+            }
+            // We now have 8 bytes to use for the length
+            let value = UInt64(truncatingIfNeeded: messageLength) | (0xC0 << 56)
+            self.setInteger(value, at: lengthPrefixIndex)
+        } else {
+            // We must encode the length in a way that uses 4 bytes, even if not necessary, because we can't have leftover 0's after the length, before the data
+            // The only way to use fewer bytes would be to do a memmove or similar
+            // That is unlikely to be worthwhile to save, at most, 3 bytes
+            // Set the top two bit mask, then write the value.
+            let value = UInt32(truncatingIfNeeded: messageLength) | (0x80 << 24)
+            self.setInteger(value, at: lengthPrefixIndex)
+        }
+
+        return totalBytesWritten
+    }
+}

Tests/NIOCoreTests/ByteBufferQUICLengthPrefixTests.swift

@@ -0,0 +1,173 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 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
+import XCTest
+
+final class ByteBufferQUICLengthPrefixTests: XCTestCase {
+    private var buffer = ByteBuffer()
+
+    // MARK: - writeQUICVariableLengthInteger tests
+
+    func testWriteOneByteQUICVariableLengthInteger() {
+        // One byte, ie less than 63, just write out as-is
+        for number in 0..<63 {
+            var buffer = ByteBuffer()
+            buffer.writeQUICVariableLengthInteger(number)
+            XCTAssertEqual(buffer.readInteger(as: UInt8.self), UInt8(number))
+            XCTAssertEqual(buffer.readableBytes, 0)
+        }
+    }
+
+    func testWriteTwoByteQUICVariableLengthInteger() {
+        var buffer = ByteBuffer()
+        buffer.writeQUICVariableLengthInteger(0b00111011_10111101)
+        // 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()
+        buffer.writeQUICVariableLengthInteger(0b00011101_01111111_00111110_01111101)
+        // 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()
+        buffer.writeQUICVariableLengthInteger(0b00000010_00011001_01111100_01011110_11111111_00010100_11101000_10001100)
+        // 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: - readQUICVariableLengthInteger tests
+
+    func testReadEmptyQUICVariableLengthInteger() {
+        var buffer = ByteBuffer()
+        XCTAssertNil(buffer.readQUICVariableLengthInteger())
+    }
+
+    func testWriteReadQUICVariableLengthInteger() {
+        for integer in [37, 15293, 494_878_333, 151_288_809_941_952_652] {
+            var buffer = ByteBuffer()
+            buffer.writeQUICVariableLengthInteger(integer)
+            XCTAssertEqual(buffer.readQUICVariableLengthInteger(), integer)
+        }
+    }
+
+    // MARK: - writeQUICLengthPrefixed with unknown length tests
+
+    func testWriteMessageWithLengthOfZero() throws {
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed { _ in
+            // write nothing
+            0
+        }
+        XCTAssertEqual(bytesWritten, 4)  // we always encode the length as 4 bytes
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), 0)
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWithLengthOfOne() throws {
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed { buffer in
+            buffer.writeString("A")
+        }
+        XCTAssertEqual(bytesWritten, 5)  // 4 for the length + 1 for the 'A'
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), 1)
+        XCTAssertEqual(self.buffer.readString(length: 1), "A")
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWithMultipleWrites() throws {
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed { buffer in
+            buffer.writeString("Hello") + buffer.writeString(" ") + buffer.writeString("World")
+        }
+        XCTAssertEqual(bytesWritten, 15)  // 4 for the length, plus 11 for the string
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), 11)
+        XCTAssertEqual(self.buffer.readString(length: 11), "Hello World")
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWithLengthUsingFull4Bytes() throws {
+        // This is the largest possible length you could encode in 4 bytes
+        let maxLength = 1_073_741_823 - 1
+        let messageWithMaxLength = String(repeating: "A", count: maxLength)
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed { buffer in
+            buffer.writeString(messageWithMaxLength)
+        }
+        XCTAssertEqual(bytesWritten, maxLength + 4)  // 4 for the length plus the message itself
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), maxLength)
+        XCTAssertEqual(self.buffer.readString(length: maxLength), messageWithMaxLength)
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWithLengthUsing8Bytes() throws {
+        // This is the largest possible length you could encode in 4 bytes
+        let maxLength = 1_073_741_823
+        let messageWithMaxLength = String(repeating: "A", count: maxLength)
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed { buffer in
+            buffer.writeString(messageWithMaxLength)
+        }
+        XCTAssertEqual(bytesWritten, maxLength + 8)  // 8 for the length plus the message itself
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), maxLength)
+        XCTAssertEqual(self.buffer.readString(length: maxLength), messageWithMaxLength)
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    // MARK: - writeQUICLengthPrefixed with known length tests
+
+    func testWriteMessageWith1ByteLength() throws {
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed(ByteBuffer(string: "hello"))
+        XCTAssertEqual(bytesWritten, 6)  // The length can be encoded in just 1 byte, followed by 'hello'
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), 5)
+        XCTAssertEqual(self.buffer.readString(length: 5), "hello")
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWith2ByteLength() throws {
+        let length = 100  // this can be anything between 64 and 16383
+        let testString = String(repeating: "A", count: length)
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
+        XCTAssertEqual(bytesWritten, length + 2)  // The length of the string, plus 2 bytes for the length
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), length)
+        XCTAssertEqual(self.buffer.readString(length: length), testString)
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWith4ByteLength() throws {
+        let length = 20_000  // this can be anything between 16384 and 1073741823
+        let testString = String(repeating: "A", count: length)
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
+        XCTAssertEqual(bytesWritten, length + 4)  // The length of the string, plus 4 bytes for the length
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), length)
+        XCTAssertEqual(self.buffer.readString(length: length), testString)
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+
+    func testWriteMessageWith8ByteLength() throws {
+        let length = 1_073_741_824  // this can be anything between 1073741824 and 4611686018427387903
+        let testString = String(repeating: "A", count: length)
+        let bytesWritten = self.buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
+        XCTAssertEqual(bytesWritten, length + 8)  // The length of the string, plus 8 bytes for the length
+        XCTAssertEqual(self.buffer.readQUICVariableLengthInteger(), length)
+        XCTAssertEqual(self.buffer.readString(length: length), testString)
+        XCTAssertTrue(self.buffer.readableBytesView.isEmpty)
+    }
+}