remove decompression limit, as maxFrameSize already covers that

vapor · Jan 11, 2023 · e90c3d7 · e90c3d7
1 parent f1cd76e
commit e90c3d7
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 54 deletions.
diff --git a/Sources/WebSocketKit/Concurrency/Compression/Decompression.swift b/Sources/WebSocketKit/Concurrency/Compression/Decompression.swift
@@ -6,45 +6,10 @@ public enum Decompression {
         /// `deflate` is the main compression algorithm for web-sockets (RFC 7692),
         /// and for now we only support `deflate`.
         let algorithm: Compression.Algorithm = .deflate
-        public var limit: Limit
 
-        public init(limit: Limit) {
-            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))
-        }
+        private init() { }
 
-        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 static let enabled = Configuration()
     }
 
     public struct DecompressionError: Error, Equatable, CustomStringConvertible {
@@ -80,16 +45,10 @@ public enum Decompression {
     }
 
     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 {
@@ -98,13 +57,6 @@ public enum Decompression {
                     output: &buffer,
                     isComplete: &isComplete
                 )
-
-                if self.limit.exceeded(
-                    compressed: compressedLength,
-                    decompressed: buffer.writerIndex + 1
-                ) {
-                    throw DecompressionError.limit
-                }
             }
 
             if part.readableBytes > 0 {

diff --git a/Sources/WebSocketKit/HTTPInitialRequestHandler.swift b/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
     }

diff --git a/Sources/WebSocketKit/WebSocket.swift b/Sources/WebSocketKit/WebSocket.swift
@@ -45,7 +45,7 @@ public final class WebSocket {
         self.channel = channel
         self.type = type
         if let decompression = decompression {
-            self.decompressor = Decompression.Decompressor(limit: decompression.limit)
+            self.decompressor = Decompression.Decompressor()
             try self.decompressor?.initializeDecoder(encoding: decompression.algorithm)
         }
         self.onTextCallback = { _, _ in }

diff --git a/Sources/WebSocketKit/WebSocketClient.swift b/Sources/WebSocketKit/WebSocketClient.swift
@@ -81,6 +81,7 @@ public final class WebSocketClient {
                     host: host,
                     path: path,
                     query: query,
+                    decompression: self.configuration.decompression,
                     headers: headers,
                     upgradePromise: upgradePromise
                 )

diff --git a/Tests/WebSocketKitTests/CompressionTests.swift b/Tests/WebSocketKitTests/CompressionTests.swift
@@ -4,7 +4,7 @@ import XCTest
 class DecompressionTests: XCTestCase {
 
     func testDeflateDecompression() throws {
-        var decompressor = Decompression.Decompressor(limit: .none)
+        var decompressor = Decompression.Decompressor()
         try decompressor.initializeDecoder(encoding: .deflate)
         defer { decompressor.deinitializeDecoder() }
 

diff --git a/Tests/WebSocketKitTests/WebSocketKitTests.swift b/Tests/WebSocketKitTests/WebSocketKitTests.swift
@@ -45,7 +45,7 @@ final class WebSocketKitTests: XCTestCase {
 
         let closePromise = elg.next().makePromise(of: Void.self)
         var configuration = WebSocketClient.Configuration()
-        configuration.decompression = .init(limit: .none)
+        configuration.decompression = .enabled
         try WebSocket.connect(to: "ws://localhost:\(port)", configuration: configuration, on: elg) { ws in
             var receivedDeflatedStrings: [String] = []
             ws.onBinary { ws, buffer in