codekingpro/portable-devtools

1import asyncio
2from contextlib import suppress
3from typing import Any, Optional, Tuple, Union
4
5from .base_protocol import BaseProtocol
6from .client_exceptions import (
7    ClientConnectionError,
8    ClientOSError,
9    ClientPayloadError,
10    ServerDisconnectedError,
11    SocketTimeoutError,
12)
13from .helpers import (
14    _EXC_SENTINEL,
15    EMPTY_BODY_STATUS_CODES,
16    BaseTimerContext,
17    set_exception,
18    set_result,
19)
20from .http import HttpResponseParser, RawResponseMessage
21from .http_exceptions import HttpProcessingError
22from .streams import EMPTY_PAYLOAD, DataQueue, StreamReader
23
24
25class ResponseHandler(BaseProtocol, DataQueue[Tuple[RawResponseMessage, StreamReader]]):
26    """Helper class to adapt between Protocol and StreamReader."""
27
28    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
29        BaseProtocol.__init__(self, loop=loop)
30        DataQueue.__init__(self, loop)
31
32        self._should_close = False
33
34        self._payload: Optional[StreamReader] = None
35        self._skip_payload = False
36        self._payload_parser = None
37
38        self._timer = None
39
40        self._tail = b""
41        self._upgraded = False
42        self._parser: Optional[HttpResponseParser] = None
43
44        self._read_timeout: Optional[float] = None
45        self._read_timeout_handle: Optional[asyncio.TimerHandle] = None
46
47        self._timeout_ceil_threshold: Optional[float] = 5
48
49        self._closed: Union[None, asyncio.Future[None]] = None
50        self._connection_lost_called = False
51
52    @property
53    def closed(self) -> Union[None, asyncio.Future[None]]:
54        """Future that is set when the connection is closed.
55
56        This property returns a Future that will be completed when the connection
57        is closed. The Future is created lazily on first access to avoid creating
58        futures that will never be awaited.
59
60        Returns:
61            - A Future[None] if the connection is still open or was closed after
62              this property was accessed
63            - None if connection_lost() was already called before this property
64              was ever accessed (indicating no one is waiting for the closure)
65        """
66        if self._closed is None and not self._connection_lost_called:
67            self._closed = self._loop.create_future()
68        return self._closed
69
70    @property
71    def upgraded(self) -> bool:
72        return self._upgraded
73
74    @property
75    def should_close(self) -> bool:
76        return bool(
77            self._should_close
78            or (self._payload is not None and not self._payload.is_eof())
79            or self._upgraded
80            or self._exception is not None
81            or self._payload_parser is not None
82            or self._buffer
83            or self._tail
84        )
85
86    def force_close(self) -> None:
87        self._should_close = True
88
89    def close(self) -> None:
90        self._exception = None  # Break cyclic references
91        transport = self.transport
92        if transport is not None:
93            transport.close()
94            self.transport = None
95            self._payload = None
96            self._drop_timeout()
97
98    def abort(self) -> None:
99        self._exception = None  # Break cyclic references
100        transport = self.transport
101        if transport is not None:
102            transport.abort()
103            self.transport = None
104            self._payload = None
105            self._drop_timeout()
106
107    def is_connected(self) -> bool:
108        return self.transport is not None and not self.transport.is_closing()
109
110    def connection_lost(self, exc: Optional[BaseException]) -> None:
111        self._connection_lost_called = True
112        self._drop_timeout()
113
114        original_connection_error = exc
115        reraised_exc = original_connection_error
116
117        connection_closed_cleanly = original_connection_error is None
118
119        if self._closed is not None:
120            # If someone is waiting for the closed future,
121            # we should set it to None or an exception. If
122            # self._closed is None, it means that
123            # connection_lost() was called already
124            # or nobody is waiting for it.
125            if connection_closed_cleanly:
126                set_result(self._closed, None)
127            else:
128                assert original_connection_error is not None
129                set_exception(
130                    self._closed,
131                    ClientConnectionError(
132                        f"Connection lost: {original_connection_error !s}",
133                    ),
134                    original_connection_error,
135                )
136
137        if self._payload_parser is not None:
138            with suppress(Exception):  # FIXME: log this somehow?
139                self._payload_parser.feed_eof()
140
141        uncompleted = None
142        if self._parser is not None:
143            try:
144                uncompleted = self._parser.feed_eof()
145            except Exception as underlying_exc:
146                if self._payload is not None:
147                    client_payload_exc_msg = (
148                        f"Response payload is not completed: {underlying_exc !r}"
149                    )
150                    if not connection_closed_cleanly:
151                        client_payload_exc_msg = (
152                            f"{client_payload_exc_msg !s}. "
153                            f"{original_connection_error !r}"
154                        )
155                    set_exception(
156                        self._payload,
157                        ClientPayloadError(client_payload_exc_msg),
158                        underlying_exc,
159                    )
160
161        if not self.is_eof():
162            if isinstance(original_connection_error, OSError):
163                reraised_exc = ClientOSError(*original_connection_error.args)
164            if connection_closed_cleanly:
165                reraised_exc = ServerDisconnectedError(uncompleted)
166            # assigns self._should_close to True as side effect,
167            # we do it anyway below
168            underlying_non_eof_exc = (
169                _EXC_SENTINEL
170                if connection_closed_cleanly
171                else original_connection_error
172            )
173            assert underlying_non_eof_exc is not None
174            assert reraised_exc is not None
175            self.set_exception(reraised_exc, underlying_non_eof_exc)
176
177        self._should_close = True
178        self._parser = None
179        self._payload = None
180        self._payload_parser = None
181        self._reading_paused = False
182
183        super().connection_lost(reraised_exc)
184
185    def eof_received(self) -> None:
186        # should call parser.feed_eof() most likely
187        self._drop_timeout()
188
189    def pause_reading(self) -> None:
190        super().pause_reading()
191        self._drop_timeout()
192
193    def resume_reading(self) -> None:
194        super().resume_reading()
195        self._reschedule_timeout()
196
197    def set_exception(
198        self,
199        exc: BaseException,
200        exc_cause: BaseException = _EXC_SENTINEL,
201    ) -> None:
202        self._should_close = True
203        self._drop_timeout()
204        super().set_exception(exc, exc_cause)
205
206    def set_parser(self, parser: Any, payload: Any) -> None:
207        # TODO: actual types are:
208        #   parser: WebSocketReader
209        #   payload: WebSocketDataQueue
210        # but they are not generi enough
211        # Need an ABC for both types
212        self._payload = payload
213        self._payload_parser = parser
214
215        self._drop_timeout()
216
217        if self._tail:
218            data, self._tail = self._tail, b""
219            self.data_received(data)
220
221    def set_response_params(
222        self,
223        *,
224        timer: Optional[BaseTimerContext] = None,
225        skip_payload: bool = False,
226        read_until_eof: bool = False,
227        auto_decompress: bool = True,
228        read_timeout: Optional[float] = None,
229        read_bufsize: int = 2**16,
230        timeout_ceil_threshold: float = 5,
231        max_line_size: int = 8190,
232        max_field_size: int = 8190,
233        max_headers: int = 128,
234    ) -> None:
235        self._skip_payload = skip_payload
236
237        self._read_timeout = read_timeout
238
239        self._timeout_ceil_threshold = timeout_ceil_threshold
240
241        self._parser = HttpResponseParser(
242            self,
243            self._loop,
244            read_bufsize,
245            timer=timer,
246            payload_exception=ClientPayloadError,
247            response_with_body=not skip_payload,
248            read_until_eof=read_until_eof,
249            auto_decompress=auto_decompress,
250            max_line_size=max_line_size,
251            max_field_size=max_field_size,
252            max_headers=max_headers,
253        )
254
255        if self._tail:
256            data, self._tail = self._tail, b""
257            self.data_received(data)
258
259    def _drop_timeout(self) -> None:
260        if self._read_timeout_handle is not None:
261            self._read_timeout_handle.cancel()
262            self._read_timeout_handle = None
263
264    def _reschedule_timeout(self) -> None:
265        timeout = self._read_timeout
266        if self._read_timeout_handle is not None:
267            self._read_timeout_handle.cancel()
268
269        if timeout:
270            self._read_timeout_handle = self._loop.call_later(
271                timeout, self._on_read_timeout
272            )
273        else:
274            self._read_timeout_handle = None
275
276    def start_timeout(self) -> None:
277        self._reschedule_timeout()
278
279    @property
280    def read_timeout(self) -> Optional[float]:
281        return self._read_timeout
282
283    @read_timeout.setter
284    def read_timeout(self, read_timeout: Optional[float]) -> None:
285        self._read_timeout = read_timeout
286
287    def _on_read_timeout(self) -> None:
288        exc = SocketTimeoutError("Timeout on reading data from socket")
289        self.set_exception(exc)
290        if self._payload is not None:
291            set_exception(self._payload, exc)
292
293    def data_received(self, data: bytes) -> None:
294        self._reschedule_timeout()
295
296        if not data:
297            return
298
299        # custom payload parser - currently always WebSocketReader
300        if self._payload_parser is not None:
301            eof, tail = self._payload_parser.feed_data(data)
302            if eof:
303                self._payload = None
304                self._payload_parser = None
305
306                if tail:
307                    self.data_received(tail)
308            return
309
310        if self._upgraded or self._parser is None:
311            # i.e. websocket connection, websocket parser is not set yet
312            self._tail += data
313            return
314
315        # parse http messages
316        try:
317            messages, upgraded, tail = self._parser.feed_data(data)
318        except BaseException as underlying_exc:
319            if self.transport is not None:
320                # connection.release() could be called BEFORE
321                # data_received(), the transport is already
322                # closed in this case
323                self.transport.close()
324            # should_close is True after the call
325            if isinstance(underlying_exc, HttpProcessingError):
326                exc = HttpProcessingError(
327                    code=underlying_exc.code,
328                    message=underlying_exc.message,
329                    headers=underlying_exc.headers,
330                )
331            else:
332                exc = HttpProcessingError()
333            self.set_exception(exc, underlying_exc)
334            return
335
336        self._upgraded = upgraded
337
338        payload: Optional[StreamReader] = None
339        for message, payload in messages:
340            if message.should_close:
341                self._should_close = True
342
343            self._payload = payload
344
345            if self._skip_payload or message.code in EMPTY_BODY_STATUS_CODES:
346                self.feed_data((message, EMPTY_PAYLOAD), 0)
347            else:
348                self.feed_data((message, payload), 0)
349
350        if payload is not None:
351            # new message(s) was processed
352            # register timeout handler unsubscribing
353            # either on end-of-stream or immediately for
354            # EMPTY_PAYLOAD
355            if payload is not EMPTY_PAYLOAD:
356                payload.on_eof(self._drop_timeout)
357            else:
358                self._drop_timeout()
359
360        if upgraded and tail:
361            self.data_received(tail)
362