Add WebSocket.maximumMessageSize (#3)

shareup · Oct 31, 2020 · b15957b · b15957b
1 parent 98375bd
commit b15957b
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 7 deletions.
diff --git a/Package.resolved b/Package.resolved
@@ -15,17 +15,17 @@
         "repositoryURL": "https://github.com/shareup/synchronized.git",
         "state": {
           "branch": null,
-          "revision": "5f3312cf35cde0ddd1196ca9c44190b1859426e0",
-          "version": "2.0.0"
+          "revision": "e8bbab660d61b4f78fd94c3640f1a03886c326ea",
+          "version": "2.1.0"
         }
       },
       {
         "package": "WebSocketProtocol",
         "repositoryURL": "https://github.com/shareup/websocket-protocol.git",
         "state": {
           "branch": null,
-          "revision": "607dad12e8cc26575cab8d188d401a9b5520c051",
-          "version": "2.1.0"
+          "revision": "e3cc3105c75d070e0e6cd82568036b67a103ff5a",
+          "version": "2.2.0"
         }
       }
     ]

diff --git a/Package.swift b/Package.swift
@@ -12,8 +12,8 @@ let package = Package(
             targets: ["WebSocket"]),
     ],
     dependencies: [
-        .package(name: "Synchronized", url: "https://github.com/shareup/synchronized.git", from: "2.0.0"),
-        .package(name: "WebSocketProtocol", url: "https://github.com/shareup/websocket-protocol.git", from: "2.0.0"),
+        .package(name: "Synchronized", url: "https://github.com/shareup/synchronized.git", from: "2.1.0"),
+        .package(name: "WebSocketProtocol", url: "https://github.com/shareup/websocket-protocol.git", from: "2.2.0"),
         .package(name: "swift-nio", url: "https://github.com/apple/swift-nio.git", from: "2.0.0"),
     ],
     targets: [

diff --git a/Sources/WebSocket/WebSocket.swift b/Sources/WebSocket/WebSocket.swift
@@ -24,6 +24,15 @@ public final class WebSocket: WebSocketProtocol {
         }
     }
 
+    /// The maximum number of bytes to buffer before the receive call fails with an error.
+    /// Default: 1 MiB
+    public var maximumMessageSize: Int = 1024 * 1024 {
+        didSet { sync {
+            guard let (_, task) = state.webSocketSessionAndTask else { return }
+            task.maximumMessageSize = maximumMessageSize
+        } }
+    }
+
     public var isOpen: Bool { sync {
         guard case .open = state else { return false }
         return true
@@ -76,6 +85,7 @@ public final class WebSocket: WebSocketProtocol {
                 let delegate = WebSocketDelegate(onOpen: onOpen, onClose: onClose, onCompletion: onCompletion)
                 let session = URLSession(configuration: .default, delegate: delegate, delegateQueue: delegateQueue)
                 let task = session.webSocketTask(with: url)
+                task.maximumMessageSize = maximumMessageSize
                 state = .connecting(session, task, delegate)
                 task.resume()
                 receiveFromWebSocket()

diff --git a/Sources/WebSocket/WebSocketMessage+URLSessionWebSocketTask.swift b/Sources/WebSocket/WebSocketMessage+URLSessionWebSocketTask.swift
@@ -5,7 +5,7 @@ extension WebSocketMessage {
     init(_ message: URLSessionWebSocketTask.Message) {
         switch message {
         case .data(let data):
-            self = .data(data)
+            self = .binary(data)
         case .string(let string):
             self = .text(string)
         @unknown default: