From 5931d913b4ea250891a0b582f1f8b2901b868c79 Mon Sep 17 00:00:00 2001
From: Ben Darnell <ben@bendarnell.com>
Date: Sat, 15 Apr 2017 12:53:28 -0400
Subject: [PATCH] websocket: Fix use of render_string in websocket handlers

PR #1917 caused websocket connections to be "finished" early, which
broke the use of render_string by setting self.ui to None.

Fixes #2010
---
 tornado/test/websocket_test.py | 24 +++++++++++++++++++++++-
 tornado/web.py                 |  3 +++
 tornado/websocket.py           | 10 ++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 7bdca8773..d47a74e65 100644
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -8,6 +8,7 @@ from tornado.concurrent import Future
 from tornado import gen
 from tornado.httpclient import HTTPError, HTTPRequest
 from tornado.log import gen_log, app_log
+from tornado.template import DictLoader
 from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog
 from tornado.test.util import unittest, skipBefore35, exec_test
 from tornado.web import Application, RequestHandler
@@ -130,6 +131,11 @@ class CoroutineOnMessageHandler(TestWebSocketHandler):
         self.write_message(message)
 
 
+class RenderMessageHandler(TestWebSocketHandler):
+    def on_message(self, message):
+        self.write_message(self.render_string('message.html', message=message))
+
+
 class WebSocketBaseTestCase(AsyncHTTPTestCase):
     @gen.coroutine
     def ws_connect(self, path, **kwargs):
@@ -168,7 +174,15 @@ class WebSocketTest(WebSocketBaseTestCase):
              dict(close_future=self.close_future)),
             ('/coroutine', CoroutineOnMessageHandler,
              dict(close_future=self.close_future)),
-        ])
+            ('/render', RenderMessageHandler,
+             dict(close_future=self.close_future)),
+        ], template_loader=DictLoader({
+            'message.html': '<b>{{ message }}</b>',
+        }))
+
+    def tearDown(self):
+        super(WebSocketTest, self).tearDown()
+        RequestHandler._template_loaders.clear()
 
     def test_http_request(self):
         # WS server, HTTP client.
@@ -219,6 +233,14 @@ class WebSocketTest(WebSocketBaseTestCase):
         self.assertEqual(response, u'hello \u00e9')
         yield self.close(ws)
 
+    @gen_test
+    def test_render_message(self):
+        ws = yield self.ws_connect('/render')
+        ws.write_message('hello')
+        response = yield ws.read_message()
+        self.assertEqual(response, '<b>hello</b>')
+        yield self.close(ws)
+
     @gen_test
     def test_error_in_on_message(self):
         ws = yield self.ws_connect('/error_in_on_message')
diff --git a/tornado/web.py b/tornado/web.py
index 8ff52e9ce..d79889fa3 100644
--- a/tornado/web.py
+++ b/tornado/web.py
@@ -993,6 +993,9 @@ class RequestHandler(object):
         self._log()
         self._finished = True
         self.on_finish()
+        self._break_cycles()
+
+    def _break_cycles(self):
         # Break up a reference cycle between this handler and the
         # _ui_module closures to allow for faster GC on CPython.
         self.ui = None
diff --git a/tornado/websocket.py b/tornado/websocket.py
index ce13d262b..69437ee4e 100644
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -434,6 +434,16 @@ class WebSocketHandler(tornado.web.RequestHandler):
         if not self._on_close_called:
             self._on_close_called = True
             self.on_close()
+            self._break_cycles()
+
+    def _break_cycles(self):
+        # WebSocketHandlers call finish() early, but we don't want to
+        # break up reference cycles (which makes it impossible to call
+        # self.render_string) until after we've really closed the
+        # connection (if it was established in the first place,
+        # indicated by status code 101).
+        if self.get_status() != 101 or self._on_close_called:
+            super(WebSocketHandler, self)._break_cycles()
 
     def send_error(self, *args, **kwargs):
         if self.stream is None:
-- 
2.47.2