diff --git a/Sources/NIOCore/ByteBuffer-QUICLengthPrefix.swift b/Sources/NIOCore/ByteBuffer-QUICLengthPrefix.swift index 02c9c5c017..0d8415544d 100644 --- a/Sources/NIOCore/ByteBuffer-QUICLengthPrefix.swift +++ b/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: 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 diff --git a/Tests/NIOCoreTests/ByteBufferQUICLengthPrefixTests.swift b/Tests/NIOCoreTests/ByteBufferQUICLengthPrefixTests.swift index b59f6ea493..6701c73c79 100644 --- a/Tests/NIOCoreTests/ByteBufferQUICLengthPrefixTests.swift +++ b/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) + } }