Skip to content

Commit

Permalink
Record what reserved bits are valid
Browse files Browse the repository at this point in the history
Throw protocol error on receiving a frame with a reserved bit set that isnt used.
  • Loading branch information
adam-fowler committed Nov 28, 2024
1 parent 1d60de6 commit b22fb91
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 0 deletions.
3 changes: 3 additions & 0 deletions Sources/WSCompression/PerMessageDeflateExtension.swift
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,9 @@ struct PerMessageDeflateExtension: WebSocketExtension {
}
return frame
}

/// Reserved bits extension uses
var reservedBits: WebSocketFrame.ReservedBits { .rsv1 }
}

extension WebSocketExtensionFactory {
Expand Down
7 changes: 7 additions & 0 deletions Sources/WSCore/Extensions/WebSocketExtension.swift
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,13 @@ public protocol WebSocketExtension: Sendable {
func processReceivedFrame(_ frame: WebSocketFrame, context: WebSocketExtensionContext) async throws -> WebSocketFrame
/// Process frame about to be sent to websocket
func processFrameToSend(_ frame: WebSocketFrame, context: WebSocketExtensionContext) async throws -> WebSocketFrame
/// Reserved bits extension uses
var reservedBits: WebSocketFrame.ReservedBits { get }
/// shutdown extension
func shutdown() async
}

extension WebSocketExtension {
/// Reserved bits extension uses (default is none)
public var reservedBits: WebSocketFrame.ReservedBits { .init() }
}
15 changes: 15 additions & 0 deletions Sources/WSCore/WebSocketHandler.swift
Original file line number Diff line number Diff line change
Expand Up @@ -71,11 +71,16 @@ public struct WebSocketCloseFrame: Sendable {
let extensions: [any WebSocketExtension]
let autoPing: AutoPingSetup
let validateUTF8: Bool
let reservedBits: WebSocketFrame.ReservedBits

@_spi(WSInternal) public init(extensions: [any WebSocketExtension], autoPing: AutoPingSetup, validateUTF8: Bool) {
self.extensions = extensions
self.autoPing = autoPing
self.validateUTF8 = validateUTF8
// store reserved bits used by this handler
self.reservedBits = extensions.reduce(.init()) { partialResult, `extension` in
partialResult.union(`extension`.reservedBits)
}
}
}

Expand Down Expand Up @@ -301,6 +306,16 @@ public struct WebSocketCloseFrame: Sendable {
}

func receivedClose(_ frame: WebSocketFrame) async throws {
guard frame.reservedBits.isEmpty else {
try await self.sendClose(code: .protocolError, reason: nil)
// Only server should initiate a connection close. Clients should wait for the
// server to close the connection when it receives the WebSocket close packet
// See https://www.rfc-editor.org/rfc/rfc6455#section-7.1.1
if self.type == .server {
self.outbound.finish()
}
return
}
switch self.stateMachine.receivedClose(frameData: frame.unmaskedData, validateUTF8: self.configuration.validateUTF8) {
case .sendClose(let errorCode):
try await self.sendClose(code: errorCode, reason: nil)
Expand Down
3 changes: 3 additions & 0 deletions Sources/WSCore/WebSocketInboundStream.swift
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ public final class WebSocketInboundStream: AsyncSequence, Sendable {
case .pong:
try await self.handler.onPong(frame)
case .text, .binary, .continuation:
guard self.handler.configuration.reservedBits.contains(frame.reservedBits) else {
throw WebSocketHandler.InternalError.close(.protocolError)
}
// apply extensions
var frame = frame
for ext in self.handler.configuration.extensions.reversed() {
Expand Down

0 comments on commit b22fb91

Please sign in to comment.