From 52adf957b3f20c0f43baea360baf6bec8eacd210 Mon Sep 17 00:00:00 2001
From: Aric Stewart <aricstewart@gmail.com>
Date: Fri, 10 Apr 2015 13:14:31 -0500
Subject: [PATCH] Add ping heartbeat option

With this option we will ping the client with a heartbeat.
Helpful for connection through firewalls that may disable inactive
connections.
---
 websockify/websocket.py      |  9 +++++++++
 websockify/websocketproxy.py | 16 ++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/websockify/websocket.py b/websockify/websocket.py
index 4d300a3..9a1306d 100644
--- a/websockify/websocket.py
+++ b/websockify/websocket.py
@@ -370,6 +370,10 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
                         repr(frame['payload']))
                     self.send_pong(frame['payload'])
                     return [], False
+                elif frame['opcode'] == 0xA: # pong
+                    self.print_traffic("} pong %s\n" %
+                        repr(frame['payload']))
+                    return [], False
 
             self.print_traffic("}")
 
@@ -407,6 +411,11 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
         buf, h, t = self.encode_hybi(s2b(data), opcode=0x0A, base64=False)
         self.request.send(buf)
 
+    def send_ping(self, data=''):
+        """ Send a WebSocket ping frame. """
+        buf, h, t = self.encode_hybi(s2b(data), opcode=0x09, base64=False)
+        self.request.send(buf)
+
     def do_websocket_handshake(self):
         h = self.headers
 
diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py
index 605fe44..94454c0 100755
--- a/websockify/websocketproxy.py
+++ b/websockify/websocketproxy.py
@@ -111,9 +111,21 @@ Traffic Legend:
         tqueue = []
         rlist = [self.request, target]
 
+        if self.server.heartbeat:
+            now = time.time()
+            self.heartbeat = now + self.server.heartbeat
+        else:
+            self.heartbeat = None
+
         while True:
             wlist = []
 
+            if self.heartbeat is not None:
+                now = time.time()
+                if now > self.heartbeat:
+                    self.heartbeat = now + self.server.heartbeat
+                    self.send_ping()
+
             if tqueue: wlist.append(target)
             if cqueue or c_pend: wlist.append(self.request)
             ins, outs, excepts = select(rlist, wlist, [], 1)
@@ -180,6 +192,7 @@ class WebSocketProxy(websocket.WebSocketServer):
         self.wrap_mode      = kwargs.pop('wrap_mode', None)
         self.unix_target    = kwargs.pop('unix_target', None)
         self.ssl_target     = kwargs.pop('ssl_target', None)
+        self.heartbeat      = kwargs.pop('heartbeat', None)
 
         token_plugin = kwargs.pop('token_plugin', None)
         token_source = kwargs.pop('token_source', None)
@@ -370,6 +383,8 @@ def websockify_init():
                            "on instantiation")
     parser.add_option("--auto-pong", action="store_true",
             help="Automatically respond to ping frames with a pong")
+    parser.add_option("--heartbeat", type=int, default=0,
+            help="send a ping to the client every HEARTBEAT seconds")
 
     (opts, args) = parser.parse_args()
 
@@ -456,6 +471,7 @@ class LibProxyServer(ForkingMixIn, HTTPServer):
         self.ssl_target     = kwargs.pop('ssl_target', None)
         self.token_plugin   = kwargs.pop('token_plugin', None)
         self.token_source   = kwargs.pop('token_source', None)
+        self.heartbeat      = kwargs.pop('heartbeat', None)
 
         self.token_plugin = None
         self.daemon = False