فهرست منبع

Send the handshake close on client request.
Send .Pong when .Ping is asked.
Properly close the session on error.

Mathieu Barnachon 10 سال پیش
والد
کامیت
6ae62d4327
1فایلهای تغییر یافته به همراه35 افزوده شده و 11 حذف شده
  1. 35 11
      Sources/WebSockets.swift

+ 35 - 11
Sources/WebSockets.swift

@@ -26,18 +26,42 @@ public func websocket(
         }
         let protocolSessionClosure: (Socket -> Void) = { socket in
             let session = WebSocketSession(socket)
-            while let frame = try? session.readFrame() {
-                switch frame.opcode {
-                case .Text:
-                    if let handleText = text {
-                        handleText(session, String.fromUInt8(frame.payload))
+            do {
+                while true {
+                    let frame = try session.readFrame()
+                    switch frame.opcode {
+                    case .Text:
+                        if let handleText = text {
+                            handleText(session, String.fromUInt8(frame.payload))
+                        }
+                    case .Binary:
+                        if let handleBinary = binary {
+                            handleBinary(session, frame.payload)
+                        }
+                    case .Close:
+                        session.writeCloseFrame()
+                    case .Continue:
+                        break
+                    case .Ping:
+                        if frame.payload.count > 125 {
+                            throw WebSocketSession.Error.ProtocolError("payload gretter than 125 octets.")
+                        } else {
+                            session.writeFrame(ArraySlice(frame.payload), .Pong)
+                        }
+                        break
+                    case .Pong:
+                        break
                     }
-                case .Binary:
-                    if let handleBinary = binary {
-                        handleBinary(session, frame.payload)
-                    }
-                default: break
                 }
+            } catch let error {
+                switch error {
+                case WebSocketSession.Error.UnknownOpCode:
+                    print("Unknown Op Code: \(error)")
+                default:
+                    print("Unkown error \(error)")
+                }
+                // If an error occurs, send the close handshake.
+                session.writeCloseFrame()
             }
         }
         let secWebSocketAccept = String.toBase64((secWebSocketKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").SHA1())
@@ -48,7 +72,7 @@ public func websocket(
 
 public class WebSocketSession: Hashable, Equatable  {
     
-    public enum Error: ErrorType { case UnknownOpCode(String), UnMaskedFrame }
+    public enum Error: ErrorType { case UnknownOpCode(String), UnMaskedFrame, ProtocolError(String) }
     public enum OpCode: UInt8 { case Continue = 0x00, Close = 0x08, Ping = 0x09, Pong = 0x0A, Text = 0x01, Binary = 0x02 }
     
     public class Frame {