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/Concurrency/Compression/Compression.swift

@@ -0,0 +1,23 @@
+
+public enum Compression {
+    public struct Algorithm {
+        enum Base {
+            case gzip
+            case deflate
+        }
+
+        private let base: Base
+
+        var window: CInt {
+            switch base {
+            case .deflate:
+                return 15
+            case .gzip:
+                return 15 + 16
+            }
+        }
+
+        public static let gzip = Self(base: .gzip)
+        public static let deflate = Self(base: .deflate)
+    }
+}

Sources/WebSocketKit/Concurrency/Compression/Decompression.swift

@@ -0,0 +1,173 @@
+import CZlib
+
+public enum Decompression {
+
+    public struct Configuration {
+        public var algorithm: Compression.Algorithm
+        public var limit: Limit
+
+        public init(algorithm: Compression.Algorithm, limit: Limit) {
+            self.algorithm = algorithm
+            self.limit = limit
+        }
+    }
+
+    /// Specifies how to limit decompression inflation.
+    public struct Limit: Sendable {
+        private enum Base {
+            case none
+            case size(Int)
+            case ratio(Int)
+        }
+
+        private var limit: Base
+
+        /// No limit will be set.
+        /// - warning: Setting `limit` to `.none` leaves you vulnerable to denial of service attacks.
+        public static let none = Limit(limit: .none)
+        /// Limit will be set on the request body size.
+        public static func size(_ value: Int) -> Limit {
+            return Limit(limit: .size(value))
+        }
+        /// Limit will be set on a ratio between compressed body size and decompressed result.
+        public static func ratio(_ value: Int) -> Limit {
+            return Limit(limit: .ratio(value))
+        }
+
+        func exceeded(compressed: Int, decompressed: Int) -> Bool {
+            switch self.limit {
+            case .none:
+                return false
+            case .size(let allowed):
+                return decompressed > allowed
+            case .ratio(let ratio):
+                return decompressed > compressed * ratio
+            }
+        }
+    }
+
+    public struct DecompressionError: Error, Equatable, CustomStringConvertible {
+
+        private enum Base: Error, Equatable {
+            case limit
+            case inflationError(Int)
+            case initializationError(Int)
+            case invalidTrailingData
+        }
+
+        private var base: Base
+
+        /// The set ``DecompressionLimit`` has been exceeded
+        public static let limit = Self(base: .limit)
+
+        /// 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 initialised.  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 let limit: Limit
+        private var stream = z_stream()
+
+        init(limit: Limit) {
+            self.limit = limit
+        }
+
+        /// Assumes `buffer` is a new empty buffer.
+        mutating func decompress(part: inout ByteBuffer, buffer: inout ByteBuffer) throws {
+            let compressedLength = part.readableBytes
+            var isComplete = false
+
+            while part.readableBytes > 0 && !isComplete {
+                try self.stream.inflatePart(
+                    input: &part,
+                    output: &buffer,
+                    isComplete: &isComplete
+                )
+
+                if self.limit.exceeded(
+                    compressed: compressedLength,
+                    decompressed: buffer.writerIndex + 1
+                ) {
+                    throw DecompressionError.limit
+                }
+            }
+
+            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/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(limit: decompression.limit)
+            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
         }
     }