From 2bd5c8febd73aad3995d77896fbbde5af30a1ba6 Mon Sep 17 00:00:00 2001
From: Aric Stewart <aricstewart@gmail.com>
Date: Fri, 10 Apr 2015 11:23:52 -0500
Subject: [PATCH] Add option to turn on auto-pong

This allows the websockify server to respond to heart beat ping
messages sent by clients
---
 websockify/websocket.py      | 14 +++++++++++++-
 websockify/websocketproxy.py |  2 ++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/websockify/websocket.py b/websockify/websocket.py
index de56af3..4d300a3 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,11 @@ 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
 
             self.print_traffic("}")
 
@@ -396,6 +402,11 @@ 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 do_websocket_handshake(self):
         h = self.headers
 
@@ -571,7 +582,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 +607,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..605fe44 100755
--- a/websockify/websocketproxy.py
+++ b/websockify/websocketproxy.py
@@ -368,6 +368,8 @@ 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")
 
     (opts, args) = parser.parse_args()