You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

233 lines
7.4 KiB

4 years ago
  1. from contextlib import suppress
  2. from .base_protocol import BaseProtocol
  3. from .client_exceptions import (ClientOSError, ClientPayloadError,
  4. ServerDisconnectedError, ServerTimeoutError)
  5. from .http import HttpResponseParser
  6. from .streams import EMPTY_PAYLOAD, DataQueue
  7. class ResponseHandler(BaseProtocol, DataQueue):
  8. """Helper class to adapt between Protocol and StreamReader."""
  9. def __init__(self, *, loop=None):
  10. BaseProtocol.__init__(self, loop=loop)
  11. DataQueue.__init__(self, loop=loop)
  12. self._should_close = False
  13. self._payload = None
  14. self._skip_payload = False
  15. self._payload_parser = None
  16. self._reading_paused = False
  17. self._timer = None
  18. self._tail = b''
  19. self._upgraded = False
  20. self._parser = None
  21. self._read_timeout = None
  22. self._read_timeout_handle = None
  23. @property
  24. def upgraded(self):
  25. return self._upgraded
  26. @property
  27. def should_close(self):
  28. if (self._payload is not None and
  29. not self._payload.is_eof() or self._upgraded):
  30. return True
  31. return (self._should_close or self._upgraded or
  32. self.exception() is not None or
  33. self._payload_parser is not None or
  34. len(self) or self._tail)
  35. def force_close(self):
  36. self._should_close = True
  37. def close(self):
  38. transport = self.transport
  39. if transport is not None:
  40. transport.close()
  41. self.transport = None
  42. self._payload = None
  43. self._drop_timeout()
  44. return transport
  45. def is_connected(self):
  46. return self.transport is not None
  47. def connection_lost(self, exc):
  48. self._drop_timeout()
  49. if self._payload_parser is not None:
  50. with suppress(Exception):
  51. self._payload_parser.feed_eof()
  52. try:
  53. uncompleted = self._parser.feed_eof()
  54. except Exception:
  55. uncompleted = None
  56. if self._payload is not None:
  57. self._payload.set_exception(
  58. ClientPayloadError('Response payload is not completed'))
  59. if not self.is_eof():
  60. if isinstance(exc, OSError):
  61. exc = ClientOSError(*exc.args)
  62. if exc is None:
  63. exc = ServerDisconnectedError(uncompleted)
  64. # assigns self._should_close to True as side effect,
  65. # we do it anyway below
  66. self.set_exception(exc)
  67. self._should_close = True
  68. self._parser = None
  69. self._payload = None
  70. self._payload_parser = None
  71. self._reading_paused = False
  72. super().connection_lost(exc)
  73. def eof_received(self):
  74. # should call parser.feed_eof() most likely
  75. self._drop_timeout()
  76. def pause_reading(self):
  77. if not self._reading_paused:
  78. try:
  79. self.transport.pause_reading()
  80. except (AttributeError, NotImplementedError, RuntimeError):
  81. pass
  82. self._reading_paused = True
  83. self._drop_timeout()
  84. def resume_reading(self):
  85. if self._reading_paused:
  86. try:
  87. self.transport.resume_reading()
  88. except (AttributeError, NotImplementedError, RuntimeError):
  89. pass
  90. self._reading_paused = False
  91. self._reschedule_timeout()
  92. def set_exception(self, exc):
  93. self._should_close = True
  94. self._drop_timeout()
  95. super().set_exception(exc)
  96. def set_parser(self, parser, payload):
  97. self._payload = payload
  98. self._payload_parser = parser
  99. self._drop_timeout()
  100. if self._tail:
  101. data, self._tail = self._tail, b''
  102. self.data_received(data)
  103. def set_response_params(self, *, timer=None,
  104. skip_payload=False,
  105. read_until_eof=False,
  106. auto_decompress=True,
  107. read_timeout=None):
  108. self._skip_payload = skip_payload
  109. self._read_timeout = read_timeout
  110. self._reschedule_timeout()
  111. self._parser = HttpResponseParser(
  112. self, self._loop, timer=timer,
  113. payload_exception=ClientPayloadError,
  114. read_until_eof=read_until_eof,
  115. auto_decompress=auto_decompress)
  116. if self._tail:
  117. data, self._tail = self._tail, b''
  118. self.data_received(data)
  119. def _drop_timeout(self):
  120. if self._read_timeout_handle is not None:
  121. self._read_timeout_handle.cancel()
  122. self._read_timeout_handle = None
  123. def _reschedule_timeout(self):
  124. timeout = self._read_timeout
  125. if self._read_timeout_handle is not None:
  126. self._read_timeout_handle.cancel()
  127. if timeout:
  128. self._read_timeout_handle = self._loop.call_later(
  129. timeout, self._on_read_timeout)
  130. else:
  131. self._read_timeout_handle = None
  132. def _on_read_timeout(self):
  133. exc = ServerTimeoutError("Timeout on reading data from socket")
  134. self.set_exception(exc)
  135. if self._payload is not None:
  136. self._payload.set_exception(exc)
  137. def data_received(self, data):
  138. if not data:
  139. return
  140. # custom payload parser
  141. if self._payload_parser is not None:
  142. eof, tail = self._payload_parser.feed_data(data)
  143. if eof:
  144. self._payload = None
  145. self._payload_parser = None
  146. if tail:
  147. self.data_received(tail)
  148. return
  149. else:
  150. if self._upgraded or self._parser is None:
  151. # i.e. websocket connection, websocket parser is not set yet
  152. self._tail += data
  153. else:
  154. # parse http messages
  155. try:
  156. messages, upgraded, tail = self._parser.feed_data(data)
  157. except BaseException as exc:
  158. if self.transport is not None:
  159. # connection.release() could be called BEFORE
  160. # data_received(), the transport is already
  161. # closed in this case
  162. self.transport.close()
  163. # should_close is True after the call
  164. self.set_exception(exc)
  165. return
  166. self._upgraded = upgraded
  167. payload = None
  168. for message, payload in messages:
  169. if message.should_close:
  170. self._should_close = True
  171. self._payload = payload
  172. if self._skip_payload or message.code in (204, 304):
  173. self.feed_data((message, EMPTY_PAYLOAD), 0)
  174. else:
  175. self.feed_data((message, payload), 0)
  176. if payload is not None:
  177. # new message(s) was processed
  178. # register timeout handler unsubscribing
  179. # either on end-of-stream or immediatelly for
  180. # EMPTY_PAYLOAD
  181. if payload is not EMPTY_PAYLOAD:
  182. payload.on_eof(self._drop_timeout)
  183. else:
  184. self._drop_timeout()
  185. if tail:
  186. if upgraded:
  187. self.data_received(tail)
  188. else:
  189. self._tail = tail