From 731fd20796437d476de63134d698e56b82779c7a Mon Sep 17 00:00:00 2001 From: Aric Stewart Date: Fri, 10 Apr 2015 11:23:52 -0500 Subject: [PATCH 1/2] 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() From 52adf957b3f20c0f43baea360baf6bec8eacd210 Mon Sep 17 00:00:00 2001 From: Aric Stewart Date: Fri, 10 Apr 2015 13:14:31 -0500 Subject: [PATCH 2/2] 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