Просмотр исходного кода

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

Mathieu Barnachon 10 лет назад
Родитель
Сommit
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 {