0001-Fix-behavior-of-recv-in-the-CLOSING-state.patch 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261
  1. From 402059e4a46a764632eba8a669f5b012f173ee7b Mon Sep 17 00:00:00 2001
  2. From: Aymeric Augustin <aymeric.augustin@m4x.org>
  3. Date: Tue, 1 May 2018 17:05:05 +0200
  4. Subject: [PATCH] Fix behavior of recv() in the CLOSING state.
  5. The behavior wasn't tested correctly: in some test cases, the connection
  6. had already moved to the CLOSED state, where the close code and reason
  7. are already known.
  8. Refactor half_close_connection_{local,remote} to allow multiple runs of
  9. the event loop while remaining in the CLOSING state. Refactor affected
  10. tests accordingly.
  11. I verified that all tests in the CLOSING state were behaving is intended
  12. by inserting debug statements in recv/send/ping/pong and running:
  13. $ PYTHONASYNCIODEBUG=1 python -m unittest -v websockets.test_protocol.{Client,Server}Tests.test_{recv,send,ping,pong}_on_closing_connection_{local,remote}
  14. Fix #317, #327, #350, #357.
  15. Signed-off-by: Joseph Kogut <joseph.kogut@gmail.com>
  16. ---
  17. websockets/protocol.py | 10 ++---
  18. websockets/test_protocol.py | 78 +++++++++++++++++++++++++++++--------
  19. 2 files changed, 66 insertions(+), 22 deletions(-)
  20. diff --git a/websockets/protocol.py b/websockets/protocol.py
  21. index f8121a1..7583fe9 100644
  22. --- a/websockets/protocol.py
  23. +++ b/websockets/protocol.py
  24. @@ -303,7 +303,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
  25. # Don't yield from self.ensure_open() here because messages could be
  26. # received before the closing frame even if the connection is closing.
  27. - # Wait for a message until the connection is closed
  28. + # Wait for a message until the connection is closed.
  29. next_message = asyncio_ensure_future(
  30. self.messages.get(), loop=self.loop)
  31. try:
  32. @@ -315,15 +315,15 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
  33. next_message.cancel()
  34. raise
  35. - # Now there's no need to yield from self.ensure_open(). Either a
  36. - # message was received or the connection was closed.
  37. -
  38. if next_message in done:
  39. return next_message.result()
  40. else:
  41. next_message.cancel()
  42. if not self.legacy_recv:
  43. - raise ConnectionClosed(self.close_code, self.close_reason)
  44. + assert self.state in [State.CLOSING, State.CLOSED]
  45. + # Wait until the connection is closed to raise
  46. + # ConnectionClosed with the correct code and reason.
  47. + yield from self.ensure_open()
  48. @asyncio.coroutine
  49. def send(self, data):
  50. diff --git a/websockets/test_protocol.py b/websockets/test_protocol.py
  51. index 70348fb..bfd4e3b 100644
  52. --- a/websockets/test_protocol.py
  53. +++ b/websockets/test_protocol.py
  54. @@ -105,7 +105,7 @@ class CommonTests:
  55. self.loop.call_soon(self.loop.stop)
  56. self.loop.run_forever()
  57. - def make_drain_slow(self, delay=3 * MS):
  58. + def make_drain_slow(self, delay=MS):
  59. # Process connection_made in order to initialize self.protocol.writer.
  60. self.run_loop_once()
  61. @@ -174,6 +174,8 @@ class CommonTests:
  62. # Empty the outgoing data stream so we can make assertions later on.
  63. self.assertOneFrameSent(True, OP_CLOSE, close_frame_data)
  64. + assert self.protocol.state is State.CLOSED
  65. +
  66. def half_close_connection_local(self, code=1000, reason='close'):
  67. """
  68. Start a closing handshake but do not complete it.
  69. @@ -181,31 +183,56 @@ class CommonTests:
  70. The main difference with `close_connection` is that the connection is
  71. left in the CLOSING state until the event loop runs again.
  72. + The current implementation returns a task that must be awaited or
  73. + cancelled, else asyncio complains about destroying a pending task.
  74. +
  75. """
  76. close_frame_data = serialize_close(code, reason)
  77. - # Trigger the closing handshake from the local side.
  78. - self.ensure_future(self.protocol.close(code, reason))
  79. + # Trigger the closing handshake from the local endpoint.
  80. + close_task = self.ensure_future(self.protocol.close(code, reason))
  81. self.run_loop_once() # wait_for executes
  82. self.run_loop_once() # write_frame executes
  83. # Empty the outgoing data stream so we can make assertions later on.
  84. self.assertOneFrameSent(True, OP_CLOSE, close_frame_data)
  85. - # Prepare the response to the closing handshake from the remote side.
  86. - self.loop.call_soon(
  87. - self.receive_frame, Frame(True, OP_CLOSE, close_frame_data))
  88. - self.loop.call_soon(self.receive_eof_if_client)
  89. +
  90. + assert self.protocol.state is State.CLOSING
  91. +
  92. + # Complete the closing sequence at 1ms intervals so the test can run
  93. + # at each point even it goes back to the event loop several times.
  94. + self.loop.call_later(
  95. + MS, self.receive_frame, Frame(True, OP_CLOSE, close_frame_data))
  96. + self.loop.call_later(2 * MS, self.receive_eof_if_client)
  97. +
  98. + # This task must be awaited or cancelled by the caller.
  99. + return close_task
  100. def half_close_connection_remote(self, code=1000, reason='close'):
  101. """
  102. - Receive a closing handshake.
  103. + Receive a closing handshake but do not complete it.
  104. The main difference with `close_connection` is that the connection is
  105. left in the CLOSING state until the event loop runs again.
  106. """
  107. + # On the server side, websockets completes the closing handshake and
  108. + # closes the TCP connection immediately. Yield to the event loop after
  109. + # sending the close frame to run the test while the connection is in
  110. + # the CLOSING state.
  111. + if not self.protocol.is_client:
  112. + self.make_drain_slow()
  113. +
  114. close_frame_data = serialize_close(code, reason)
  115. - # Trigger the closing handshake from the remote side.
  116. + # Trigger the closing handshake from the remote endpoint.
  117. self.receive_frame(Frame(True, OP_CLOSE, close_frame_data))
  118. - self.receive_eof_if_client()
  119. + self.run_loop_once() # read_frame executes
  120. + # Empty the outgoing data stream so we can make assertions later on.
  121. + self.assertOneFrameSent(True, OP_CLOSE, close_frame_data)
  122. +
  123. + assert self.protocol.state is State.CLOSING
  124. +
  125. + # Complete the closing sequence at 1ms intervals so the test can run
  126. + # at each point even it goes back to the event loop several times.
  127. + self.loop.call_later(2 * MS, self.receive_eof_if_client)
  128. def process_invalid_frames(self):
  129. """
  130. @@ -335,11 +362,13 @@ class CommonTests:
  131. self.assertEqual(data, b'tea')
  132. def test_recv_on_closing_connection_local(self):
  133. - self.half_close_connection_local()
  134. + close_task = self.half_close_connection_local()
  135. with self.assertRaises(ConnectionClosed):
  136. self.loop.run_until_complete(self.protocol.recv())
  137. + self.loop.run_until_complete(close_task) # cleanup
  138. +
  139. def test_recv_on_closing_connection_remote(self):
  140. self.half_close_connection_remote()
  141. @@ -421,24 +450,29 @@ class CommonTests:
  142. self.assertNoFrameSent()
  143. def test_send_on_closing_connection_local(self):
  144. - self.half_close_connection_local()
  145. + close_task = self.half_close_connection_local()
  146. with self.assertRaises(ConnectionClosed):
  147. self.loop.run_until_complete(self.protocol.send('foobar'))
  148. +
  149. self.assertNoFrameSent()
  150. + self.loop.run_until_complete(close_task) # cleanup
  151. +
  152. def test_send_on_closing_connection_remote(self):
  153. self.half_close_connection_remote()
  154. with self.assertRaises(ConnectionClosed):
  155. self.loop.run_until_complete(self.protocol.send('foobar'))
  156. - self.assertOneFrameSent(True, OP_CLOSE, serialize_close(1000, 'close'))
  157. +
  158. + self.assertNoFrameSent()
  159. def test_send_on_closed_connection(self):
  160. self.close_connection()
  161. with self.assertRaises(ConnectionClosed):
  162. self.loop.run_until_complete(self.protocol.send('foobar'))
  163. +
  164. self.assertNoFrameSent()
  165. # Test the ping coroutine.
  166. @@ -466,24 +500,29 @@ class CommonTests:
  167. self.assertNoFrameSent()
  168. def test_ping_on_closing_connection_local(self):
  169. - self.half_close_connection_local()
  170. + close_task = self.half_close_connection_local()
  171. with self.assertRaises(ConnectionClosed):
  172. self.loop.run_until_complete(self.protocol.ping())
  173. +
  174. self.assertNoFrameSent()
  175. + self.loop.run_until_complete(close_task) # cleanup
  176. +
  177. def test_ping_on_closing_connection_remote(self):
  178. self.half_close_connection_remote()
  179. with self.assertRaises(ConnectionClosed):
  180. self.loop.run_until_complete(self.protocol.ping())
  181. - self.assertOneFrameSent(True, OP_CLOSE, serialize_close(1000, 'close'))
  182. +
  183. + self.assertNoFrameSent()
  184. def test_ping_on_closed_connection(self):
  185. self.close_connection()
  186. with self.assertRaises(ConnectionClosed):
  187. self.loop.run_until_complete(self.protocol.ping())
  188. +
  189. self.assertNoFrameSent()
  190. # Test the pong coroutine.
  191. @@ -506,24 +545,29 @@ class CommonTests:
  192. self.assertNoFrameSent()
  193. def test_pong_on_closing_connection_local(self):
  194. - self.half_close_connection_local()
  195. + close_task = self.half_close_connection_local()
  196. with self.assertRaises(ConnectionClosed):
  197. self.loop.run_until_complete(self.protocol.pong())
  198. +
  199. self.assertNoFrameSent()
  200. + self.loop.run_until_complete(close_task) # cleanup
  201. +
  202. def test_pong_on_closing_connection_remote(self):
  203. self.half_close_connection_remote()
  204. with self.assertRaises(ConnectionClosed):
  205. self.loop.run_until_complete(self.protocol.pong())
  206. - self.assertOneFrameSent(True, OP_CLOSE, serialize_close(1000, 'close'))
  207. +
  208. + self.assertNoFrameSent()
  209. def test_pong_on_closed_connection(self):
  210. self.close_connection()
  211. with self.assertRaises(ConnectionClosed):
  212. self.loop.run_until_complete(self.protocol.pong())
  213. +
  214. self.assertNoFrameSent()
  215. # Test the protocol's logic for acknowledging pings with pongs.
  216. --
  217. 2.17.0