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:

Patch:
diff --git a/tornado/web.py b/tornado/web.py
index 132f42dd..d79889fa 100644
--- a/tornado/web.py
+++ b/tornado/web.py
@@ -319,7 +319,10 @@ class RequestHandler(object):
         if reason is not None:
             self._reason = escape.native_str(reason)
         else:
-            self._reason = httputil.responses.get(status_code, "Unknown")
+            try:
+                self._reason = httputil.responses[status_code]
+            except KeyError:
+                raise ValueError("unknown status code %d" % status_code)
 
     def get_status(self):
         """Returns the status code for our response."""
@@ -1558,7 +1561,11 @@ class RequestHandler(object):
             # send a response.
             return
         if isinstance(e, HTTPError):
-            self.send_error(e.status_code, exc_info=sys.exc_info())
+            if e.status_code not in httputil.responses and not e.reason:
+                gen_log.error("Bad HTTP status code: %d", e.status_code)
+                self.send_error(500, exc_info=sys.exc_info())
+            else:
+                self.send_error(e.status_code, exc_info=sys.exc_info())
         else:
             self.send_error(500, exc_info=sys.exc_info())
 
diff --git a/tornado/websocket.py b/tornado/websocket.py
index 69437ee4..7600910c 100644
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -764,10 +764,7 @@ class WebSocketProtocol13(WebSocketProtocol):
             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 @@ class WebSocketProtocol13(WebSocketProtocol):
             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 @@ class WebSocketProtocol13(WebSocketProtocol):
                     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/web_test.py b/tornado/test/web_test.py
index de26c42e..d79ea52c 100644
--- a/tornado/test/web_test.py
+++ b/tornado/test/web_test.py
@@ -917,10 +917,6 @@ class ErrorResponseTest(WebTestCase):
             self.assertEqual(response.code, 503)
             self.assertTrue(b"503: Service Unavailable" in response.body)
 
-            response = self.fetch("/default?status=435")
-            self.assertEqual(response.code, 435)
-            self.assertTrue(b"435: Unknown" in response.body)
-
     def test_write_error(self):
         with ExpectLog(app_log, "Uncaught exception"):
             response = self.fetch("/write_error")
@@ -1493,9 +1489,9 @@ class StatusReasonTest(SimpleHandlerTestCase):
         response = self.fetch("/?code=682&reason=Bar")
         self.assertEqual(response.code, 682)
         self.assertEqual(response.reason, "Bar")
-        response = self.fetch("/?code=682")
-        self.assertEqual(response.code, 682)
-        self.assertEqual(response.reason, "Unknown")
+        with ExpectLog(app_log, 'Uncaught exception'):
+            response = self.fetch("/?code=682")
+        self.assertEqual(response.code, 500)
 
 
 @wsgi_safe
diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index d47a74e6..e0b5573d 100644
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -7,6 +7,7 @@ import traceback
 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 @@ class TestWebSocketHandler(WebSocketHandler):
 
 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 @@ class WebSocketTest(WebSocketBaseTestCase):
         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