Merge pull request #299 from allenhuang/add_pong_closure_to_websocket

Add pong closure to websocket

Author: damian-kolakowski

Sources/DemoServer.swift

@@ -180,7 +180,10 @@ public func demoServer(_ publicDir: String) -> HttpServer {
         session.writeText(text)
         }, { (session, binary) in
         session.writeBinary(binary)
-    })
+        }, { (session, pong) in
+        // Got a pong frame
+        }
+    )
 
     server.notFoundHandler = { r in
         return .movedPermanently("https://github.com/404")

Sources/WebSockets.swift

@@ -10,7 +10,8 @@ import Foundation
 
 public func websocket(
       _ text: ((WebSocketSession, String) -> Void)?,
-    _ binary: ((WebSocketSession, [UInt8]) -> Void)?) -> ((HttpRequest) -> HttpResponse) {
+    _ binary: ((WebSocketSession, [UInt8]) -> Void)?,
+      _ pong: ((WebSocketSession, [UInt8]) -> Void)?) -> ((HttpRequest) -> HttpResponse) {
     return { r in
         guard r.hasTokenForHeader("upgrade", token: "websocket") else {
             return .badRequest(.text("Invalid value of 'Upgrade' header: \(r.headers["upgrade"] ?? "unknown")"))
@@ -90,6 +91,9 @@ public func websocket(
                         session.writeFrame(ArraySlice(frame.payload), .pong)
                     }
                 case .pong:
+                    if let handlePong = pong {
+                       handlePong(session, frame.payload)
+                    }
                     break
                 }
             }