Merge pull request #169 from AricStewart/ping-pong

Implement Ping and Pong heartbeat messages

2 files changed, 40 insertions, 1 deletions

diff --git a/websockify/websocket.py b/websockify/websocket.py
index de56af3..9a1306d 100644
--- a/websockify/websocket.py
+++ b/websockify/websocket.py
@@ -104,6 +104,7 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
         self.handler_id = getattr(server, "handler_id", False)
         self.file_only = getattr(server, "file_only", False)
         self.traffic = getattr(server, "traffic", False)
+        self.auto_pong = getattr(server, "auto_pong", False)
 
         self.logger = getattr(server, "logger", None)
         if self.logger is None:
@@ -364,6 +365,15 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
                     closed = {'code': frame['close_code'],
                               'reason': frame['close_reason']}
                     break
+                elif self.auto_pong and frame['opcode'] == 0x9: # ping
+                    self.print_traffic("} ping %s\n" %
+                        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("}")
 
@@ -396,6 +406,16 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
         buf, h, t = self.encode_hybi(msg, opcode=0x08, base64=False)
         self.request.send(buf)
 
+    def send_pong(self, data=''):
+        """ Send a WebSocket pong frame. """
+        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
 
@@ -571,7 +591,7 @@ class WebSocketServer(object):
             file_only=False,
             run_once=False, timeout=0, idle_timeout=0, traffic=False,
             tcp_keepalive=True, tcp_keepcnt=None, tcp_keepidle=None,
-            tcp_keepintvl=None):
+            tcp_keepintvl=None, auto_pong=False):
 
         # settings
         self.RequestHandlerClass = RequestHandlerClass
@@ -596,6 +616,7 @@ class WebSocketServer(object):
         self.tcp_keepidle   = tcp_keepidle
         self.tcp_keepintvl  = tcp_keepintvl
 
+        self.auto_pong      = auto_pong
         # Make paths settings absolute
         self.cert = os.path.abspath(cert)
         self.key = self.web = self.record = ''
diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py
index 97c5d82..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)
@@ -368,6 +381,10 @@ def websockify_init():
     parser.add_option("--token-source", default=None, metavar="ARG",
                       help="an argument to be passed to the token plugin"
                            "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()
 
@@ -454,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