Clean up close because of error

hummingbird-project · Mar 22, 2024 · 2d1d5d6 · 2d1d5d6
1 parent d97dfe0
commit 2d1d5d6
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/Sources/HummingbirdWebSocket/WebSocketHandler.swift b/Sources/HummingbirdWebSocket/WebSocketHandler.swift
@@ -86,11 +86,10 @@ actor WebSocketHandler: Sendable {
                                         await webSocket.inbound.send(frameSeq.data)
                                         frameSequence = nil
                                     }
-                                } catch let error as NIOWebSocketError {
-                                    let errorCode = WebSocketErrorCode(error)
-                                    try await self.close(code: errorCode, outbound: webSocket.outbound, context: context)
                                 } catch {
-                                    let errorCode = WebSocketErrorCode.unexpectedServerError
+                                    // catch errors while processing websocket frames so responding close message
+                                    // can be dealt with
+                                    let errorCode = WebSocketErrorCode(error)
                                     try await self.close(code: errorCode, outbound: webSocket.outbound, context: context)
                                 }
                             }
@@ -189,13 +188,15 @@ actor WebSocketHandler: Sendable {
 }
 
 extension WebSocketErrorCode {
-    init(_ error: NIOWebSocketError) {
+    init(_ error: any Error) {
         switch error {
-        case .invalidFrameLength:
+        case NIOWebSocketError.invalidFrameLength:
             self = .messageTooLarge
-        case .fragmentedControlFrame,
-             .multiByteControlFrameLength:
+        case NIOWebSocketError.fragmentedControlFrame,
+             NIOWebSocketError.multiByteControlFrameLength:
             self = .protocolError
+        default:
+            self = .unexpectedServerError
         }
     }
 }