From a36579bb1238cd017f00a1d00ac82b8ae19b9699 Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Fri, 10 Feb 2017 15:11:51 +0800
Subject: [PATCH] Send ping packets to keep connection alive

Fix #126

Idle WebSocket connections tend to be closed after some period of time.
This commit enables the Butterfly server to send ping packets
periodically in order to keep the connection alive.

A new option `keepalive_interval` is also introduced for users to
specify the interval to send `ping` packets. By default it is 30
seconds.
---
 butterfly.server.py |  2 ++
 butterfly/routes.py | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/butterfly.server.py b/butterfly.server.py
index b825996..6ba58c3 100755
--- a/butterfly.server.py
+++ b/butterfly.server.py
@@ -40,6 +40,8 @@ tornado.options.define("unminified", default=False,
 
 tornado.options.define("host", default='localhost', help="Server host")
 tornado.options.define("port", default=57575, type=int, help="Server port")
+tornado.options.define("keepalive_interval", default=30, type=int,
+                        help="Interval between ping packets sent from server to client (in seconds)")
 tornado.options.define("one_shot", default=False,
                        help="Run a one-shot instance. Quit at term close")
 tornado.options.define("shell", help="Shell to execute at login")
diff --git a/butterfly/routes.py b/butterfly/routes.py
index dd301c5..b63cb4b 100644
--- a/butterfly/routes.py
+++ b/butterfly/routes.py
@@ -17,7 +17,9 @@
 
 
 import os
+import struct
 import sys
+import time
 import tornado.options
 import tornado.process
 import tornado.escape
@@ -133,10 +135,15 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
     # Session history
     history = {}
 
+    # Keepalive timer
+    keepalive_timer = None
+
     def open(self, user, path, session):
+        opts = tornado.options.options
         self.session = session
         self.closed = False
         self.secure_user = None
+        self.keepalive_timer = tornado.ioloop.PeriodicCallback(self.send_ping, opts.keepalive_interval * 1000)
 
         # Prevent cross domain
         if self.request.headers['Origin'] not in (
@@ -155,7 +162,6 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
         self.set_nodelay(True)
 
         socket = utils.Socket(self.ws_connection.stream.socket)
-        opts = tornado.options.options
 
         if not opts.unsecure:
             user = utils.parse_cert(
@@ -206,6 +212,8 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
         else:
             self._terminal = terminal
 
+        self.keepalive_timer.start()
+
     @property
     def user_sessions(self):
         """Return the dict session of socket lists"""
@@ -278,6 +286,8 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
     def on_close(self):
         if self.closed:
             return
+        if self.keepalive_timer != None:
+            self.keepalive_timer.stop()
         self.closed = True
         self.log.info('Websocket closed %r' % self)
         TermWebSocket.sockets.remove(self)
@@ -297,6 +307,15 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
                     for user, sessions in TermWebSocket.terminals.items()])):
             sys.exit(0)
 
+    def send_ping(self):
+        t = int(time.time())
+        frame = struct.pack('<I', t) # A ping frame based on time
+        self.log.info("Sending ping frame %s" % t)
+        try:
+            self.ping(frame)
+        except tornado.websocket.WebSocketClosedError:
+            self.keepalive_timer.stop()
+
 
 @url(r'/sessions/list.json')
 class SessionsList(Route):