Преглед на файлове

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 {