From: Ben Darnell <ben@bendarnell.com>
Date: Sun, 11 Mar 2018 20:13:51 +0000 (-0400)
Subject: websocket: Make ping() argument optional
X-Git-Tag: v5.1.0b1~42^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=cfd621a30d56f345c179f2d30c51163986ffac33;p=thirdparty%2Ftornado.git

websocket: Make ping() argument optional

Also accept both bytes and str, and add a client-side ping() method.

Fixes #2295
---

diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 38b3211d7..2e14342c1 100644
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -621,6 +621,34 @@ class ClientPeriodicPingTest(WebSocketBaseTestCase):
         # TODO: test that the connection gets closed if ping responses stop.
 
 
+class ManualPingTest(WebSocketBaseTestCase):
+    def get_app(self):
+        class PingHandler(TestWebSocketHandler):
+            def on_ping(self, data):
+                self.write_message(data, binary=isinstance(data, bytes))
+
+        self.close_future = Future()
+        return Application([
+            ('/', PingHandler, dict(close_future=self.close_future)),
+        ])
+
+    @gen_test
+    def test_manual_ping(self):
+        ws = yield self.ws_connect('/')
+
+        self.assertRaises(ValueError, ws.ping, 'a' * 126)
+
+        ws.ping('hello')
+        resp = yield ws.read_message()
+        # on_ping always sees bytes.
+        self.assertEqual(resp, b'hello')
+
+        ws.ping(b'binary hello')
+        resp = yield ws.read_message()
+        self.assertEqual(resp, b'binary hello')
+        yield self.close(ws)
+
+
 class MaxMessageSizeTest(WebSocketBaseTestCase):
     def get_app(self):
         self.close_future = Future()
diff --git a/tornado/websocket.py b/tornado/websocket.py
index df60885e2..629e9417e 100644
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -312,8 +312,23 @@ class WebSocketHandler(tornado.web.RequestHandler):
         """
         raise NotImplementedError
 
-    def ping(self, data):
-        """Send ping frame to the remote end."""
+    def ping(self, data=b''):
+        """Send ping frame to the remote end.
+
+        The data argument allows a small amount of data (up to 125
+        bytes) to be sent as a part of the ping message. Note that not
+        all websocket implementations expose this data to
+        applications.
+
+        Consider using the ``websocket_ping_interval`` application
+        setting instead of sending pings manually.
+
+        .. versionchanged:: 5.1
+
+           The data argument is now optional.
+
+        """
+        data = utf8(data)
         if self.ws_connection is None:
             raise WebSocketClosedError()
         self.ws_connection.write_ping(data)
@@ -756,12 +771,19 @@ class WebSocketProtocol13(WebSocketProtocol):
             **self._get_compressor_options(other_side, agreed_parameters, compression_options))
 
     def _write_frame(self, fin, opcode, data, flags=0):
+        data_len = len(data)
+        if opcode & 0x8:
+            # All control frames MUST have a payload length of 125
+            # bytes or less and MUST NOT be fragmented.
+            if not fin:
+                raise ValueError("control frames may not be fragmented")
+            if data_len > 125:
+                raise ValueError("control frame payloads may not exceed 125 bytes")
         if fin:
             finbit = self.FIN
         else:
             finbit = 0
         frame = struct.pack("B", finbit | opcode | flags)
-        data_len = len(data)
         if self.mask_outgoing:
             mask_bit = 0x80
         else:
@@ -1205,6 +1227,25 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
         else:
             self.read_queue.append(message)
 
+    def ping(self, data=b''):
+        """Send ping frame to the remote end.
+
+        The data argument allows a small amount of data (up to 125
+        bytes) to be sent as a part of the ping message. Note that not
+        all websocket implementations expose this data to
+        applications.
+
+        Consider using the ``ping_interval`` argument to
+        `websocket_connect` instead of sending pings manually.
+
+        .. versionadded:: 5.1
+
+        """
+        data = utf8(data)
+        if self.protocol is None:
+            raise WebSocketClosedError()
+        self.protocol.write_ping(data)
+
     def on_pong(self, data):
         pass