Skip to content

Commit

Permalink
Random WebSocket Request Key (apple#1855)
Browse files Browse the repository at this point in the history
* add randomRequestKey method

* fix some typos

* fix compilation on 5.0

* run generate_linux_tests.rb

* add test with default random number generator

* @inlineable

* base64Encoding @inlineable
  • Loading branch information
dnadoba authored May 24, 2021
1 parent 4ded068 commit 465f87d
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 5 deletions.
13 changes: 10 additions & 3 deletions Sources/NIOWebSocket/Base64.swift
Original file line number Diff line number Diff line change
Expand Up @@ -18,16 +18,18 @@
extension String {

/// Base64 encode a collection of UInt8 to a string, without the use of Foundation.
@inlinable
init<Buffer: Collection>(base64Encoding bytes: Buffer)
where Buffer.Element == UInt8
{
self = Base64.encode(bytes: bytes)
}
}

@usableFromInline
internal struct Base64 {

fileprivate struct Base64 {

@inlinable
static func encode<Buffer: Collection>(bytes: Buffer)
-> String where Buffer.Element == UInt8
{
Expand Down Expand Up @@ -66,6 +68,7 @@ fileprivate struct Base64 {
// MARK: Internal

// The base64 unicode table.
@usableFromInline
static let encodeBase64: [UInt8] = [
UInt8(ascii: "A"), UInt8(ascii: "B"), UInt8(ascii: "C"), UInt8(ascii: "D"),
UInt8(ascii: "E"), UInt8(ascii: "F"), UInt8(ascii: "G"), UInt8(ascii: "H"),
Expand All @@ -86,12 +89,14 @@ fileprivate struct Base64 {
]

static let encodePaddingCharacter: UInt8 = UInt8(ascii: "=")


@usableFromInline
static func encode(alphabet: [UInt8], firstByte: UInt8) -> UInt8 {
let index = firstByte >> 2
return alphabet[Int(index)]
}

@usableFromInline
static func encode(alphabet: [UInt8], firstByte: UInt8, secondByte: UInt8?) -> UInt8 {
var index = (firstByte & 0b00000011) << 4
if let secondByte = secondByte {
Expand All @@ -100,6 +105,7 @@ fileprivate struct Base64 {
return alphabet[Int(index)]
}

@usableFromInline
static func encode(alphabet: [UInt8], secondByte: UInt8?, thirdByte: UInt8?) -> UInt8 {
guard let secondByte = secondByte else {
// No second byte means we are just emitting padding.
Expand All @@ -112,6 +118,7 @@ fileprivate struct Base64 {
return alphabet[Int(index)]
}

@usableFromInline
static func encode(alphabet: [UInt8], thirdByte: UInt8?) -> UInt8 {
guard let thirdByte = thirdByte else {
// No third byte means just padding.
Expand Down
33 changes: 31 additions & 2 deletions Sources/NIOWebSocket/NIOWebSocketClientUpgrader.swift
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,13 @@ public final class NIOWebSocketClientUpgrader: NIOHTTPClientProtocolUpgrader {
private let maxFrameSize: Int
private let automaticErrorHandling: Bool
private let upgradePipelineHandler: (Channel, HTTPResponseHead) -> EventLoopFuture<Void>

public init(requestKey: String,

/// - Parameters:
/// - requestKey: sent to the server in the `Sec-WebSocket-Key` HTTP header. Default is random request key.
/// - maxFrameSize: largest incoming `WebSocketFrame` size in bytes. Default is 16,384 bytes.
/// - automaticErrorHandling: If true, adds `WebSocketProtocolErrorHandler` to the channel pipeline to catch and respond to WebSocket protocol errors. Default is true.
/// - upgradePipelineHandler: called once the upgrade was successful
public init(requestKey: String = randomRequestKey(),
maxFrameSize: Int = 1 << 14,
automaticErrorHandling: Bool = true,
upgradePipelineHandler: @escaping (Channel, HTTPResponseHead) -> EventLoopFuture<Void>) {
Expand Down Expand Up @@ -92,3 +97,27 @@ public final class NIOWebSocketClientUpgrader: NIOHTTPClientProtocolUpgrader {
}
}
}

extension NIOWebSocketClientUpgrader {
/// Generates a random WebSocket Request Key by generating 16 bytes randomly and encoding them as a base64 string as defined in RFC6455 https://tools.ietf.org/html/rfc6455#section-4.1
/// - Parameter generator: the `RandomNumberGenerator` used as a the source of randomness
/// - Returns: base64 encoded request key
@inlinable
public static func randomRequestKey<Generator>(
using generator: inout Generator
) -> String where Generator: RandomNumberGenerator{
var buffer = ByteBuffer()
buffer.reserveCapacity(minimumWritableBytes: 16)
/// we may want to use `randomBytes(count:)` once the proposal is accepted: https://forums.swift.org/t/pitch-requesting-larger-amounts-of-randomness-from-systemrandomnumbergenerator/27226
buffer.writeInteger(UInt64.random(in: UInt64.min...UInt64.max, using: &generator))
buffer.writeInteger(UInt64.random(in: UInt64.min...UInt64.max, using: &generator))
return String(base64Encoding: buffer.readableBytesView)
}
/// Generates a random WebSocket Request Key by generating 16 bytes randomly using the `SystemRandomNumberGenerator` and encoding them as a base64 string as defined in RFC6455 https://tools.ietf.org/html/rfc6455#section-4.1.
/// - Returns: base64 encoded request key
@inlinable
public static func randomRequestKey() -> String {
var generator = SystemRandomNumberGenerator()
return NIOWebSocketClientUpgrader.randomRequestKey(using: &generator)
}
}
1 change: 1 addition & 0 deletions Tests/LinuxMain.swift
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ class LinuxMainRunnerImpl: LinuxMainRunner {
testCase(NIOHTTPServerRequestAggregatorTest.allTests),
testCase(NIOSingleStepByteToMessageDecoderTest.allTests),
testCase(NIOThreadPoolTest.allTests),
testCase(NIOWebSocketClientUpgraderTests.allTests),
testCase(NIOWebSocketFrameAggregatorTests.allTests),
testCase(NonBlockingFileIOTest.allTests),
testCase(PendingDatagramWritesManagerTests.allTests),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
//===----------------------------------------------------------------------===//
//
// This source file is part of the SwiftNIO open source project
//
// Copyright (c) 2017-2018 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
//
//===----------------------------------------------------------------------===//
//
// NIOWebSocketClientUpgraderTests+XCTest.swift
//
import XCTest

///
/// NOTE: This file was generated by generate_linux_tests.rb
///
/// Do NOT edit this file directly as it will be regenerated automatically when needed.
///

extension NIOWebSocketClientUpgraderTests {

@available(*, deprecated, message: "not actually deprecated. Just deprecated to allow deprecated tests (which test deprecated functionality) without warnings")
static var allTests : [(String, (NIOWebSocketClientUpgraderTests) -> () throws -> Void)] {
return [
("testRandomRequestKey", testRandomRequestKey),
("testRandomRequestKeyWithSystemRandomNumberGenerator", testRandomRequestKeyWithSystemRandomNumberGenerator),
]
}
}

42 changes: 42 additions & 0 deletions Tests/NIOWebSocketTests/NIOWebSocketClientUpgraderTests.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
//===----------------------------------------------------------------------===//
//
// 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 XCTest
import NIOWebSocket


/// a mock random number generator which will return the given `numbers` in order
fileprivate struct TestRandomNumberGenerator: RandomNumberGenerator {
var numbers: [UInt64]
var nextRandomNumberIndex: Int
init(numbers: [UInt64], nextRandomNumberIndex: Int = 0) {
self.numbers = numbers
self.nextRandomNumberIndex = nextRandomNumberIndex
}
mutating func next() -> UInt64 {
defer { nextRandomNumberIndex += 1 }
return numbers[nextRandomNumberIndex % numbers.count]
}
}

final class NIOWebSocketClientUpgraderTests: XCTestCase {
func testRandomRequestKey() {
var generator = TestRandomNumberGenerator(numbers: [10, 11])
let requestKey = NIOWebSocketClientUpgrader.randomRequestKey(using: &generator)
XCTAssertEqual(requestKey, "AAAAAAAAAAoAAAAAAAAACw==")
}
func testRandomRequestKeyWithSystemRandomNumberGenerator() {
XCTAssertEqual(NIOWebSocketClientUpgrader.randomRequestKey().count, 24, "request key must be exactly 16 bytes long and this corresponds to 24 characters in base64")
}
}

0 comments on commit 465f87d

Please sign in to comment.