Client-Decompression support

Package.swift

@@ -19,6 +19,7 @@ let package = Package(
     ],
     targets: [
         .target(name: "WebSocketKit", dependencies: [
+            "CZlib",
             .product(name: "NIO", package: "swift-nio"),
             .product(name: "NIOCore", package: "swift-nio"),
             .product(name: "NIOConcurrencyHelpers", package: "swift-nio"),
@@ -29,6 +30,12 @@ let package = Package(
             .product(name: "NIOTransportServices", package: "swift-nio-transport-services"),
             .product(name: "Atomics", package: "swift-atomics")
         ]),
+        .target(
+            name: "CZlib",
+            dependencies: [],
+            linkerSettings: [
+                .linkedLibrary("z")
+            ]),
         .testTarget(name: "WebSocketKitTests", dependencies: [
             .target(name: "WebSocketKit"),
         ]),

Sources/CZlib/empty.c

@@ -0,0 +1 @@
+

Sources/CZlib/include/CZlib.h

@@ -0,0 +1,14 @@
+#ifndef C_ZLIB_H
+#define C_ZLIB_H
+
+#include <zlib.h>
+
+static inline int CZlib_inflateInit2(z_streamp strm, int windowBits) {
+    return inflateInit2(strm, windowBits);
+}
+
+static inline Bytef *CZlib_voidPtr_to_BytefPtr(void *in) {
+    return (Bytef *)in;
+}
+
+#endif

Sources/WebSocketKit/Compression/Compression.swift

@@ -0,0 +1,13 @@
+
+enum Compression {
+    enum Algorithm {
+        case deflate
+
+        var window: CInt {
+            switch self {
+            case .deflate:
+                return 15
+            }
+        }
+    }
+}

Sources/WebSocketKit/Compression/Decompression.swift

@@ -0,0 +1,122 @@
+import CZlib
+
+public enum Decompression {
+
+    public struct Configuration {
+        /// For now we only support `deflate`, because it's the main compression
+        /// algorithm for web-sockets (RFC 7692).
+        let algorithm: Compression.Algorithm = .deflate
+
+        private init() { }
+
+        public static let enabled = Configuration()
+    }
+
+    public struct DecompressionError: Error, Equatable, CustomStringConvertible {
+
+        private enum Base: Error, Equatable {
+            case inflationError(Int)
+            case initializationError(Int)
+            case invalidTrailingData
+        }
+
+        private var base: Base
+
+        /// An error occurred when inflating. Error code is included to aid diagnosis.
+        public static var inflationError: (Int) -> Self = {
+            Self(base: .inflationError($0))
+        }
+
+        /// Decoder could not be initialized. Error code is included to aid diagnosis.
+        public static var initializationError: (Int) -> Self = {
+            Self(base: .initializationError($0))
+        }
+
+        /// Decompression completed but there was invalid trailing data behind the compressed data.
+        public static var invalidTrailingData = Self(base: .invalidTrailingData)
+
+        public var description: String {
+            return String(describing: self.base)
+        }
+    }
+
+    struct Decompressor {
+        private var stream = z_stream()
+
+        /// Assumes `buffer` is a new empty buffer.
+        mutating func decompress(part: inout ByteBuffer, buffer: inout ByteBuffer) throws {
+            var isComplete = false
+
+            while part.readableBytes > 0 && !isComplete {
+                try self.stream.inflatePart(
+                    input: &part,
+                    output: &buffer,
+                    isComplete: &isComplete
+                )
+            }
+
+            if part.readableBytes > 0 {
+                throw DecompressionError.invalidTrailingData
+            }
+        }
+
+        mutating func initializeDecoder(encoding: Compression.Algorithm) throws {
+            self.stream.zalloc = nil
+            self.stream.zfree = nil
+            self.stream.opaque = nil
+
+            let rc = CZlib_inflateInit2(&self.stream, encoding.window)
+            guard rc == Z_OK else {
+                throw DecompressionError.initializationError(Int(rc))
+            }
+        }
+
+        mutating func deinitializeDecoder() {
+            inflateEnd(&self.stream)
+        }
+    }
+}
+
+//MARK: - +z_stream
+private extension z_stream {
+    mutating func inflatePart(
+        input: inout ByteBuffer,
+        output: inout ByteBuffer,
+        isComplete: inout Bool
+    ) throws {
+        let minimumCapacity = input.readableBytes * 4
+        try input.readWithUnsafeMutableReadableBytes { pointer in
+            self.avail_in = UInt32(pointer.count)
+            self.next_in = CZlib_voidPtr_to_BytefPtr(pointer.baseAddress!)
+
+            defer {
+                self.avail_in = 0
+                self.next_in = nil
+                self.avail_out = 0
+                self.next_out = nil
+            }
+
+            isComplete = try self.inflatePart(to: &output, minimumCapacity: minimumCapacity)
+
+            return pointer.count - Int(self.avail_in)
+        }
+    }
+
+    private mutating func inflatePart(to buffer: inout ByteBuffer, minimumCapacity: Int) throws -> Bool {
+        var rc = Z_OK
+
+        try buffer.writeWithUnsafeMutableBytes(minimumWritableBytes: minimumCapacity) { pointer in
+            self.avail_out = UInt32(pointer.count)
+            self.next_out = CZlib_voidPtr_to_BytefPtr(pointer.baseAddress!)
+
+            rc = inflate(&self, Z_SYNC_FLUSH)
+            guard rc == Z_OK || rc == Z_STREAM_END else {
+                throw Decompression.DecompressionError.inflationError(Int(rc))
+            }
+
+            return pointer.count - Int(self.avail_out)
+        }
+
+        return rc == Z_STREAM_END
+    }
+}

Sources/WebSocketKit/HTTPInitialRequestHandler.swift

@@ -8,13 +8,15 @@ final class HTTPInitialRequestHandler: ChannelInboundHandler, RemovableChannelHa
     let host: String
     let path: String
     let query: String?
+    let decompression: Decompression.Configuration?
     let headers: HTTPHeaders
     let upgradePromise: EventLoopPromise<Void>
 
-    init(host: String, path: String, query: String?, headers: HTTPHeaders, upgradePromise: EventLoopPromise<Void>) {
+    init(host: String, path: String, query: String?, decompression: Decompression.Configuration?, headers: HTTPHeaders, upgradePromise: EventLoopPromise<Void>) {
         self.host = host
         self.path = path
         self.query = query
+        self.decompression = decompression
         self.headers = headers
         self.upgradePromise = upgradePromise
     }

Sources/WebSocketKit/WebSocket.swift

@@ -25,20 +25,31 @@ public final class WebSocket {
     }
 
     private let channel: Channel
-    private var onTextCallback: (WebSocket, String) -> ()
+
+    private var onTextCallback: ((WebSocket, String) -> ())?
+    private var onTextBufferCallback: (WebSocket, ByteBuffer) -> ()
     private var onBinaryCallback: (WebSocket, ByteBuffer) -> ()
     private var onPongCallback: (WebSocket) -> ()
     private var onPingCallback: (WebSocket) -> ()
+
     private var frameSequence: WebSocketFrameSequence?
     private let type: PeerType
+
+    private var decompressor: Decompression.Decompressor?
+
     private var waitingForPong: Bool
     private var waitingForClose: Bool
     private var scheduledTimeoutTask: Scheduled<Void>?
 
-    init(channel: Channel, type: PeerType) {
+    init(channel: Channel, type: PeerType, decompression: Decompression.Configuration?) throws {
         self.channel = channel
         self.type = type
+        if let decompression = decompression {
+            self.decompressor = Decompression.Decompressor()
+            try self.decompressor?.initializeDecoder(encoding: decompression.algorithm)
+        }
         self.onTextCallback = { _, _ in }
+        self.onTextBufferCallback = { _, _ in }
         self.onBinaryCallback = { _, _ in }
         self.onPongCallback = { _ in }
         self.onPingCallback = { _ in }
@@ -51,6 +62,11 @@ public final class WebSocket {
         self.onTextCallback = callback
     }
 
+    /// The same as `onText`, but with raw data instead of the decoded `String`.
+    public func onTextBuffer(_ callback: @escaping (WebSocket, ByteBuffer) -> ()) {
+        self.onTextBufferCallback = callback
+    }
+
     public func onBinary(_ callback: @escaping (WebSocket, ByteBuffer) -> ()) {
         self.onBinaryCallback = callback
     }
@@ -64,10 +80,10 @@ public final class WebSocket {
     }
 
     /// If set, this will trigger automatic pings on the connection. If ping is not answered before
-    /// the next ping is sent, then the WebSocket will be presumed innactive and will be closed
+    /// the next ping is sent, then the WebSocket will be presumed inactive and will be closed
     /// automatically.
     /// These pings can also be used to keep the WebSocket alive if there is some other timeout
-    /// mechanism shutting down innactive connections, such as a Load Balancer deployed in
+    /// mechanism shutting down inactive connections, such as a Load Balancer deployed in
     /// front of the server.
     public var pingInterval: TimeAmount? {
         didSet {
@@ -233,6 +249,7 @@ public final class WebSocket {
             } else {
                 frameSequence = WebSocketFrameSequence(type: frame.opcode)
             }
+
             // append this frame and update the sequence
             frameSequence.append(frame)
             self.frameSequence = frameSequence
@@ -252,12 +269,27 @@ public final class WebSocket {
 
         // if this frame was final and we have a non-nil frame sequence,
         // output it to the websocket and clear storage
-        if let frameSequence = self.frameSequence, frame.fin {
+        if var frameSequence = self.frameSequence, frame.fin {
             switch frameSequence.type {
             case .binary:
-                self.onBinaryCallback(self, frameSequence.binaryBuffer)
+                if decompressor != nil {
+                    do {
+                        var buffer = ByteBuffer()
+                        try decompressor!.decompress(part: &frameSequence.buffer, buffer: &buffer)
+
+                        self.onBinaryCallback(self, buffer)
+                    } catch {
+                        self.close(code: .protocolError, promise: nil)
+                        return
+                    }
+                } else {
+                    self.onBinaryCallback(self, frameSequence.buffer)
+                }
             case .text:
-                self.onTextCallback(self, frameSequence.textBuffer)
+                if let callback = self.onTextCallback {
+                    callback(self, String(buffer: frameSequence.buffer))
+                }
+                self.onTextBufferCallback(self, frameSequence.buffer)
             case .ping, .pong:
                 assertionFailure("Control frames never have a frameSequence")
             default: break
@@ -293,30 +325,25 @@ public final class WebSocket {
     }
 
     deinit {
+        self.decompressor?.deinitializeDecoder()
         assert(self.isClosed, "WebSocket was not closed before deinit.")
     }
 }
 
 private struct WebSocketFrameSequence {
-    var binaryBuffer: ByteBuffer
-    var textBuffer: String
+    var buffer: ByteBuffer
     var type: WebSocketOpcode
 
     init(type: WebSocketOpcode) {
-        self.binaryBuffer = ByteBufferAllocator().buffer(capacity: 0)
-        self.textBuffer = .init()
+        self.buffer = ByteBufferAllocator().buffer(capacity: 0)
         self.type = type
     }
 
     mutating func append(_ frame: WebSocketFrame) {
-        var data = frame.unmaskedData
         switch type {
-        case .binary:
-            self.binaryBuffer.writeBuffer(&data)
-        case .text:
-            if let string = data.readString(length: data.readableBytes) {
-                self.textBuffer += string
-            }
+        case .binary, .text:
+            var data = frame.unmaskedData
+            self.buffer.writeBuffer(&data)
         default: break
         }
     }

Sources/WebSocketKit/WebSocketClient.swift

@@ -25,13 +25,25 @@ public final class WebSocketClient {
     public struct Configuration {
         public var tlsConfiguration: TLSConfiguration?
         public var maxFrameSize: Int
-
+        public var decompression: Decompression.Configuration?
+
         public init(
             tlsConfiguration: TLSConfiguration? = nil,
             maxFrameSize: Int = 1 << 14
         ) {
             self.tlsConfiguration = tlsConfiguration
             self.maxFrameSize = maxFrameSize
+            self.decompression = nil
+        }
+
+        public init(
+            tlsConfiguration: TLSConfiguration? = nil,
+            maxFrameSize: Int = 1 << 14,
+            decompression: Decompression.Configuration?
+        ) {
+            self.tlsConfiguration = tlsConfiguration
+            self.maxFrameSize = maxFrameSize
+            self.decompression = decompression
         }
     }
 
@@ -69,6 +81,7 @@ public final class WebSocketClient {
                     host: host,
                     path: path,
                     query: query,
+                    decompression: self.configuration.decompression,
                     headers: headers,
                     upgradePromise: upgradePromise
                 )
@@ -82,7 +95,11 @@ public final class WebSocketClient {
                     maxFrameSize: self.configuration.maxFrameSize,
                     automaticErrorHandling: true,
                     upgradePipelineHandler: { channel, req in
-                        return WebSocket.client(on: channel, onUpgrade: onUpgrade)
+                        return WebSocket.client(
+                            on: channel,
+                            decompression: self.configuration.decompression,
+                            onUpgrade: onUpgrade
+                        )
                     }
                 )
 
@@ -130,7 +147,6 @@ public final class WebSocketClient {
         }
     }
 
-
     public func syncShutdown() throws {
         switch self.eventLoopGroupProvider {
         case .shared: