Add writeQUICLengthPrefixedString and writeQUICLengthPrefixedBytes

Sources/NIOCore/ByteBuffer-QUICLengthPrefix.swift

@@ -85,26 +85,11 @@ extension ByteBuffer {
         }
     }
 
-    /// 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
+    /// - 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 ``writeQUICLengthPrefixedBuffer(_:)`` instead
     @discardableResult
     @inlinable
     public mutating func writeQUICLengthPrefixed(
@@ -135,7 +120,7 @@ extension ByteBuffer {
         if messageLength >= 1_073_741_823 {
             // This message length cannot fit in the 4 bytes which we reserved. We need to make more space before the message
             self._createSpace(before: startWriterIndex, length: messageLength, requiredSpace: 4)
-            totalBytesWritten += 4 // the 4 bytes we just reserved
+            totalBytesWritten += 4  // the 4 bytes we just reserved
             // We now have 8 bytes to use for the length
             let value = UInt64(truncatingIfNeeded: messageLength) | (0xC0 << 56)
             self.setInteger(value, at: lengthPrefixIndex)
@@ -152,6 +137,67 @@ extension ByteBuffer {
     }
 }
 
+// MARK: - Helpers for writing QUIC variable-length prefixed things
+
+extension ByteBuffer {
+    /// 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 writeQUICLengthPrefixedBuffer(_ buffer: ByteBuffer) -> Int {
+        var written = 0
+        written += self.writeQUICVariableLengthInteger(buffer.readableBytes)
+        written += self.writeImmutableBuffer(buffer)
+        return written
+    }
+
+    /// Write a QUIC variable-length integer prefixed string.
+    ///
+    /// [RFC 9000 § 16](https://www.rfc-editor.org/rfc/rfc9000.html#section-16)
+    ///
+    /// Write `string` prefixed with the length of string as a QUIC variable-length integer
+    /// - Parameter string: The string to be written
+    /// - 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 writeQUICLengthPrefixedString(_ string: String) -> 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.writeQUICVariableLengthInteger(string.utf8.count)
+        written += self.writeString(string)
+        return written
+    }
+
+    /// Write a QUIC variable-length integer prefixed collection of bytes.
+    ///
+    /// [RFC 9000 § 16](https://www.rfc-editor.org/rfc/rfc9000.html#section-16)
+    ///
+    /// Write `bytes` prefixed with the number of bytes as a QUIC variable-length integer
+    /// - Parameter bytes: The bytes to be written
+    /// - Returns: The total bytes written. This is the bytes needed to write the length, plus the length of the string itself
+    @inlinable
+    public mutating func writeQUICLengthPrefixedBytes<Bytes: Sequence>(_ bytes: Bytes) -> Int
+    where Bytes.Element == UInt8 {
+        let numberOfBytes = bytes.withContiguousStorageIfAvailable { b in
+            UnsafeRawBufferPointer(b).count
+        }
+        if let numberOfBytes {
+            var written = 0
+            written += self.writeQUICVariableLengthInteger(numberOfBytes)
+            written += self.writeBytes(bytes)
+            return written
+        } else {
+            return self.writeQUICLengthPrefixed { buffer in
+                buffer.writeBytes(bytes)
+            }
+        }
+    }
+}
+
 extension ByteBuffer {
     /// Creates `requiredSpace` bytes of free space before `index` by moving the `length` bytes after `index` forward by `requiredSpace`
     /// e.g. given [a, b, c, d, e, f, g, h, i, j] and calling this function with (before: 4, length: 6, requiredSpace: 2) would result in

Tests/NIOCoreTests/ByteBufferQUICLengthPrefixTests.swift

@@ -77,7 +77,7 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {
         }
     }
 
-    // MARK: - writeQUICLengthPrefixed with unknown length tests
+    // MARK: - writeQUICLengthPrefixed tests
 
     func testWriteMessageWithLengthOfZero() throws {
         var buffer = ByteBuffer()
@@ -140,11 +140,11 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {
         XCTAssertTrue(buffer.readableBytesView.isEmpty)
     }
 
-    // MARK: - writeQUICLengthPrefixed with known length tests
+    // MARK: - writeQUICLengthPrefixedBuffer tests
 
     func testWriteMessageWith1ByteLength() throws {
         var buffer = ByteBuffer()
-        let bytesWritten = buffer.writeQUICLengthPrefixed(ByteBuffer(string: "hello"))
+        let bytesWritten = buffer.writeQUICLengthPrefixedBuffer(ByteBuffer(string: "hello"))
         XCTAssertEqual(bytesWritten, 6)  // The length can be encoded in just 1 byte, followed by 'hello'
         XCTAssertEqual(buffer.readQUICVariableLengthInteger(), 5)
         XCTAssertEqual(buffer.readString(length: 5), "hello")
@@ -155,7 +155,7 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {
         var buffer = ByteBuffer()
         let length = 100  // this can be anything between 64 and 16383
         let testString = String(repeating: "A", count: length)
-        let bytesWritten = buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
+        let bytesWritten = buffer.writeQUICLengthPrefixedBuffer(ByteBuffer(string: testString))
         XCTAssertEqual(bytesWritten, length + 2)  // The length of the string, plus 2 bytes for the length
         XCTAssertEqual(buffer.readQUICVariableLengthInteger(), length)
         XCTAssertEqual(buffer.readString(length: length), testString)
@@ -166,7 +166,7 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {
         var buffer = ByteBuffer()
         let length = 20_000  // this can be anything between 16384 and 1073741823
         let testString = String(repeating: "A", count: length)
-        let bytesWritten = buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
+        let bytesWritten = buffer.writeQUICLengthPrefixedBuffer(ByteBuffer(string: testString))
         XCTAssertEqual(bytesWritten, length + 4)  // The length of the string, plus 4 bytes for the length
         XCTAssertEqual(buffer.readQUICVariableLengthInteger(), length)
         XCTAssertEqual(buffer.readString(length: length), testString)
@@ -177,10 +177,40 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {
         var buffer = ByteBuffer()
         let length = 1_073_741_824  // this can be anything between 1073741824 and 4611686018427387903
         let testString = String(repeating: "A", count: length)
-        let bytesWritten = buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
+        let bytesWritten = buffer.writeQUICLengthPrefixedBuffer(ByteBuffer(string: testString))
         XCTAssertEqual(bytesWritten, length + 8)  // The length of the string, plus 8 bytes for the length
         XCTAssertEqual(buffer.readQUICVariableLengthInteger(), length)
         XCTAssertEqual(buffer.readString(length: length), testString)
         XCTAssertTrue(buffer.readableBytesView.isEmpty)
     }
+
+    // MARK: - writeQUICLengthPrefixedString tests
+
+    func testWriteQUICLengthPrefixedString() throws {
+        var buffer = ByteBuffer()
+        let testString = "Hello World"  // length = 11
+        let bytesWritten = buffer.writeQUICLengthPrefixedString(testString)
+        XCTAssertEqual(bytesWritten, 12)
+
+        let length = buffer.readQUICVariableLengthInteger()
+        XCTAssertEqual(length, 11)
+
+        XCTAssertEqual(buffer.readString(length: 11), testString)
+        XCTAssertTrue(buffer.readableBytesView.isEmpty)
+    }
+
+    // MARK: - writeQUICLengthPrefixedBytes tests
+
+    func testWriteQUICLengthPrefixedBytes() throws {
+        var buffer = ByteBuffer()
+        let testBytes: [UInt8] = [1, 2, 3, 4, 5]  // length = 5
+        let bytesWritten = buffer.writeQUICLengthPrefixedBytes(testBytes)
+        XCTAssertEqual(bytesWritten, 6)
+
+        let length = buffer.readQUICVariableLengthInteger()
+        XCTAssertEqual(length, 5)
+
+        XCTAssertEqual(buffer.readBytes(length: 5), testBytes)
+        XCTAssertTrue(buffer.readableBytesView.isEmpty)
+    }
 }