Repository: tornadoweb/tornado

Issue:
websocket `yield connection.write_message()` throws BadYieldError
At present code such as:

        yield connection.write_message(data)

will throw BadYieldError if called on a closed connection.

Documentation for WebSocketHandler.write_message says it was changed in 4.3 so that it " Returns a `.Future` which can be used for flow control."  In fact it can return either a Future or None. It actually returns the result from the method WebSocketProtocol13._write_frame, and that ends with:

        try:
            return self.stream.write(frame)
        except StreamClosedError:
            self._abort()

so in the case where the stream is closed it will return None. It should instead use gen.Return() to send back None as the result.:

        try:
            return self.stream.write(frame)
        except StreamClosedError:
            self._abort()
        raise gen.Return()

    


Discussion:

Resolving patch:
diff --git a/tornado/websocket.py b/tornado/websocket.py
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -764,10 +764,7 @@ def _write_frame(self, fin, opcode, data, flags=0):
             data = mask + _websocket_mask(mask, data)
         frame += data
         self._wire_bytes_out += len(frame)
-        try:
-            return self.stream.write(frame)
-        except StreamClosedError:
-            self._abort()
+        return self.stream.write(frame)
 
     def write_message(self, message, binary=False):
         """Sends the given message to the client of this Web Socket."""
@@ -951,7 +948,10 @@ def _handle_message(self, opcode, data):
             self.close(self.handler.close_code)
         elif opcode == 0x9:
             # Ping
-            self._write_frame(True, 0xA, data)
+            try:
+                self._write_frame(True, 0xA, data)
+            except StreamClosedError:
+                self._abort()
             self._run_callback(self.handler.on_ping, data)
         elif opcode == 0xA:
             # Pong
@@ -972,7 +972,10 @@ def close(self, code=None, reason=None):
                     close_data = struct.pack('>H', code)
                 if reason is not None:
                     close_data += utf8(reason)
-                self._write_frame(True, 0x8, close_data)
+                try:
+                    self._write_frame(True, 0x8, close_data)
+                except StreamClosedError:
+                    self._abort()
             self.server_terminated = True
         if self.client_terminated:
             if self._waiting is not None:


Test patch:
diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -7,6 +7,7 @@
 from tornado.concurrent import Future
 from tornado import gen
 from tornado.httpclient import HTTPError, HTTPRequest
+from tornado.iostream import StreamClosedError
 from tornado.log import gen_log, app_log
 from tornado.template import DictLoader
 from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog
@@ -50,7 +51,10 @@ def on_close(self):
 
 class EchoHandler(TestWebSocketHandler):
     def on_message(self, message):
-        self.write_message(message, isinstance(message, bytes))
+        try:
+            self.write_message(message, isinstance(message, bytes))
+        except StreamClosedError:
+            pass
 
 
 class ErrorInOnMessageHandler(TestWebSocketHandler):
@@ -327,6 +331,14 @@ def test_client_close_reason(self):
         self.assertEqual(code, 1001)
         self.assertEqual(reason, 'goodbye')
 
+    @gen_test
+    def test_write_after_close(self):
+        ws = yield self.ws_connect('/close_reason')
+        msg = yield ws.read_message()
+        self.assertIs(msg, None)
+        with self.assertRaises(StreamClosedError):
+            ws.write_message('hello')
+
     @gen_test
     def test_async_prepare(self):
         # Previously, an async prepare method triggered a bug that would