Merge pull request #319 from marcc-orange/wsConnect

 Websocket: Add connected/disconnected closures

README.md

@@ -51,9 +51,9 @@ server.start()
 ### How to WebSockets ?
 ```swift
 let server = HttpServer()
-server["/websocket-echo"] = websocket({ (session, text) in
+server["/websocket-echo"] = websocket(text: { session, text in
   session.writeText(text)
-}, { (session, binary) in
+}, binary: { session, binary in
   session.writeBinary(binary)
 })
 server.start()

Sources/DemoServer.swift

@@ -176,14 +176,17 @@ public func demoServer(_ publicDir: String) -> HttpServer {
         })
     }
     
-    server["/websocket-echo"] = websocket({ (session, text) in
+    server["/websocket-echo"] = websocket(text: { (session, text) in
         session.writeText(text)
-        }, { (session, binary) in
+    }, binary: { (session, binary) in
         session.writeBinary(binary)
-        }, { (session, pong) in
+    }, pong: { (session, pong) in
         // Got a pong frame
-        }
-    )
+    }, connected: { (session) in
+        // New client connected
+    }, disconnected: { (session) in
+        // Client disconnected
+    })
     
     server.notFoundHandler = { r in
         return .movedPermanently("https://github.com/404")

Sources/WebSockets.swift

@@ -7,11 +7,19 @@
 
 import Foundation
 
+@available(*, deprecated, message: "Use websocket(text:binary:pong:connected:disconnected:) instead.")
+public func websocket(_ text: @escaping (WebSocketSession, String) -> Void,
+                      _ binary: @escaping (WebSocketSession, [UInt8]) -> Void,
+                      _ pong: @escaping (WebSocketSession, [UInt8]) -> Void) -> ((HttpRequest) -> HttpResponse) {
+    return websocket(text: text, binary: binary, pong: pong)
+}
 
 public func websocket(
-      _ text: ((WebSocketSession, String) -> Void)?,
-    _ binary: ((WebSocketSession, [UInt8]) -> Void)?,
-      _ pong: ((WebSocketSession, [UInt8]) -> Void)?) -> ((HttpRequest) -> HttpResponse) {
+    text: ((WebSocketSession, String) -> Void)? = nil,
+    binary: ((WebSocketSession, [UInt8]) -> Void)? = nil,
+    pong: ((WebSocketSession, [UInt8]) -> Void)? = nil,
+    connected: ((WebSocketSession) -> Void)? = nil,
+    disconnected: ((WebSocketSession) -> Void)? = nil) -> ((HttpRequest) -> HttpResponse) {
     return { r in
         guard r.hasTokenForHeader("upgrade", token: "websocket") else {
             return .badRequest(.text("Invalid value of 'Upgrade' header: \(r.headers["upgrade"] ?? "unknown")"))
@@ -105,6 +113,8 @@ public func websocket(
                 }
             }
             
+            connected?(session)
+            
             do {
                 try read()
             } catch let error {
@@ -126,6 +136,8 @@ public func websocket(
                 // If an error occurs, send the close handshake.
                 session.writeCloseFrame()
             }
+            
+            disconnected?(session)
         }
         guard let secWebSocketAccept = String.toBase64((secWebSocketKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").sha1()) else {
             return HttpResponse.internalServerError