Skip to content

Commit

Permalink
code review
Browse files Browse the repository at this point in the history
  • Loading branch information
@hamzahrmalik
hamzahrmalik committed Sep 2, 2024
1 parent 28dba9b commit 5c508e1
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 66 deletions.
53 changes: 29 additions & 24 deletions Sources/NIOCore/ByteBuffer-QUICLengthPrefix.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
//
// This source file is part of the SwiftNIO open source project
//
// Copyright (c) 2021 Apple Inc. and the SwiftNIO project authors
// Copyright (c) 2024 Apple Inc. and the SwiftNIO project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
Expand Down Expand Up @@ -57,21 +57,13 @@ extension ByteBuffer {
}
}

/// 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 {
public mutating func writeQUICVariableLengthInteger<Integer: FixedWidthInteger>(_ value: Integer) -> 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.
Expand Down Expand Up @@ -135,25 +127,15 @@ extension ByteBuffer {
totalBytesWritten += messageLength

let actualBytesWritten = endWriterIndex - startWriterIndex
assert(
precondition(
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
)
}
// 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
// We now have 8 bytes to use for the length
let value = UInt64(truncatingIfNeeded: messageLength) | (0xC0 << 56)
self.setInteger(value, at: lengthPrefixIndex)
Expand All @@ -169,3 +151,26 @@ extension ByteBuffer {
return totalBytesWritten
}
}

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
/// [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, length: Int, requiredSpace: Int) {
// Add the required number of bytes to the end first
self.writeBytes([UInt8](repeatElement(0, count: requiredSpace)))
// Move the message forward by that many bytes, to make space at the front
self.withUnsafeMutableReadableBytes { pointer in
_ = memmove(
// The new position is 4 forward from where the user written message currently begins
pointer.baseAddress!.advanced(by: index + requiredSpace),
// This is the position where the user written message currently begins
pointer.baseAddress!.advanced(by: index),
length
)
}
}
}
97 changes: 55 additions & 42 deletions Tests/NIOCoreTests/ByteBufferQUICLengthPrefixTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
//
// This source file is part of the SwiftNIO open source project
//
// Copyright (c) 2021 Apple Inc. and the SwiftNIO project authors
// Copyright (c) 2024 Apple Inc. and the SwiftNIO project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
Expand All @@ -16,23 +16,23 @@ 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)
let bytesWritten = buffer.writeQUICVariableLengthInteger(number)
XCTAssertEqual(bytesWritten, 1)
XCTAssertEqual(buffer.readInteger(as: UInt8.self), UInt8(number))
XCTAssertEqual(buffer.readableBytes, 0)
}
}

func testWriteTwoByteQUICVariableLengthInteger() {
var buffer = ByteBuffer()
buffer.writeQUICVariableLengthInteger(0b00111011_10111101)
let bytesWritten = buffer.writeQUICVariableLengthInteger(0b00111011_10111101)
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)
Expand All @@ -41,15 +41,19 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {

func testWriteFourByteQUICVariableLengthInteger() {
var buffer = ByteBuffer()
buffer.writeQUICVariableLengthInteger(0b00011101_01111111_00111110_01111101)
let bytesWritten = buffer.writeQUICVariableLengthInteger(0b00011101_01111111_00111110_01111101)
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()
buffer.writeQUICVariableLengthInteger(0b00000010_00011001_01111100_01011110_11111111_00010100_11101000_10001100)
let bytesWritten = buffer.writeQUICVariableLengthInteger(
0b00000010_00011001_01111100_01011110_11111111_00010100_11101000_10001100
)
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),
Expand All @@ -76,98 +80,107 @@ final class ByteBufferQUICLengthPrefixTests: XCTestCase {
// MARK: - writeQUICLengthPrefixed with unknown length tests

func testWriteMessageWithLengthOfZero() throws {
let bytesWritten = self.buffer.writeQUICLengthPrefixed { _ in
var buffer = ByteBuffer()
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), 0)
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWithLengthOfOne() throws {
let bytesWritten = self.buffer.writeQUICLengthPrefixed { buffer in
var buffer = ByteBuffer()
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), 1)
XCTAssertEqual(buffer.readString(length: 1), "A")
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWithMultipleWrites() throws {
let bytesWritten = self.buffer.writeQUICLengthPrefixed { buffer in
var buffer = ByteBuffer()
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), 11)
XCTAssertEqual(buffer.readString(length: 11), "Hello World")
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWithLengthUsingFull4Bytes() throws {
var buffer = ByteBuffer()
// 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
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), maxLength)
XCTAssertEqual(buffer.readString(length: maxLength), messageWithMaxLength)
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWithLengthUsing8Bytes() throws {
var buffer = ByteBuffer()
// 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
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), maxLength)
XCTAssertEqual(buffer.readString(length: maxLength), messageWithMaxLength)
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

// MARK: - writeQUICLengthPrefixed with known length tests

func testWriteMessageWith1ByteLength() throws {
let bytesWritten = self.buffer.writeQUICLengthPrefixed(ByteBuffer(string: "hello"))
var buffer = ByteBuffer()
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), 5)
XCTAssertEqual(buffer.readString(length: 5), "hello")
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWith2ByteLength() throws {
var buffer = ByteBuffer()
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))
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), length)
XCTAssertEqual(buffer.readString(length: length), testString)
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWith4ByteLength() throws {
var buffer = ByteBuffer()
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))
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), length)
XCTAssertEqual(buffer.readString(length: length), testString)
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}

func testWriteMessageWith8ByteLength() throws {
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 = self.buffer.writeQUICLengthPrefixed(ByteBuffer(string: testString))
let bytesWritten = 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)
XCTAssertEqual(buffer.readQUICVariableLengthInteger(), length)
XCTAssertEqual(buffer.readString(length: length), testString)
XCTAssertTrue(buffer.readableBytesView.isEmpty)
}
}

0 comments on commit 5c508e1

Please sign in to comment.