codekingpro/portable-devtools
115k
1import asyncio
2import base64
3import binascii
4import hashlib
5import json
6import sys
7from typing import Any, Final, Iterable, Optional, Tuple, Union, cast
8
9import attr
10from multidict import CIMultiDict
11
12from . import hdrs
13from ._websocket.reader import WebSocketDataQueue
14from ._websocket.writer import DEFAULT_LIMIT
15from .abc import AbstractStreamWriter
16from .client_exceptions import WSMessageTypeError
17from .helpers import calculate_timeout_when, set_exception, set_result
18from .http import (
19 WS_CLOSED_MESSAGE,
20 WS_CLOSING_MESSAGE,
21 WS_KEY,
22 WebSocketError,
23 WebSocketReader,
24 WebSocketWriter,
25 WSCloseCode,
26 WSMessage,
27 WSMsgType as WSMsgType,
28 ws_ext_gen,
29 ws_ext_parse,
30)
31from .http_websocket import _INTERNAL_RECEIVE_TYPES
32from .log import ws_logger
33from .streams import EofStream
34from .typedefs import JSONDecoder, JSONEncoder
35from .web_exceptions import HTTPBadRequest, HTTPException
36from .web_request import BaseRequest
37from .web_response import StreamResponse
38
39if sys.version_info >= (3, 11):
40 import asyncio as async_timeout
41else:
42 import async_timeout
43
44__all__ = (
45 "WebSocketResponse",
46 "WebSocketReady",
47 "WSMsgType",
48)
49
50THRESHOLD_CONNLOST_ACCESS: Final[int] = 5
51
52
53@attr.s(auto_attribs=True, frozen=True, slots=True)
54class WebSocketReady:
55 ok: bool
56 protocol: Optional[str]
57
58 def __bool__(self) -> bool:
59 return self.ok
60
61
62class WebSocketResponse(StreamResponse):
63
64 _length_check: bool = False
65 _ws_protocol: Optional[str] = None
66 _writer: Optional[WebSocketWriter] = None
67 _reader: Optional[WebSocketDataQueue] = None
68 _closed: bool = False
69 _closing: bool = False
70 _conn_lost: int = 0
71 _close_code: Optional[int] = None
72 _loop: Optional[asyncio.AbstractEventLoop] = None
73 _waiting: bool = False
74 _close_wait: Optional[asyncio.Future[None]] = None
75 _exception: Optional[BaseException] = None
76 _heartbeat_when: float = 0.0
77 _heartbeat_cb: Optional[asyncio.TimerHandle] = None
78 _pong_response_cb: Optional[asyncio.TimerHandle] = None
79 _ping_task: Optional[asyncio.Task[None]] = None
80
81 def __init__(
82 self,
83 *,
84 timeout: float = 10.0,
85 receive_timeout: Optional[float] = None,
86 autoclose: bool = True,
87 autoping: bool = True,
88 heartbeat: Optional[float] = None,
89 protocols: Iterable[str] = (),
90 compress: bool = True,
91 max_msg_size: int = 4 * 1024 * 1024,
92 writer_limit: int = DEFAULT_LIMIT,
93 ) -> None:
94 super().__init__(status=101)
95 self._protocols = protocols
96 self._timeout = timeout
97 self._receive_timeout = receive_timeout
98 self._autoclose = autoclose
99 self._autoping = autoping
100 self._heartbeat = heartbeat
101 if heartbeat is not None:
102 self._pong_heartbeat = heartbeat / 2.0
103 self._compress: Union[bool, int] = compress
104 self._max_msg_size = max_msg_size
105 self._writer_limit = writer_limit
106
107 def _cancel_heartbeat(self) -> None:
108 self._cancel_pong_response_cb()
109 if self._heartbeat_cb is not None:
110 self._heartbeat_cb.cancel()
111 self._heartbeat_cb = None
112 if self._ping_task is not None:
113 self._ping_task.cancel()
114 self._ping_task = None
115
116 def _cancel_pong_response_cb(self) -> None:
117 if self._pong_response_cb is not None:
118 self._pong_response_cb.cancel()
119 self._pong_response_cb = None
120
121 def _reset_heartbeat(self) -> None:
122 if self._heartbeat is None:
123 return
124 self._cancel_pong_response_cb()
125 req = self._req
126 timeout_ceil_threshold = (
127 req._protocol._timeout_ceil_threshold if req is not None else 5
128 )
129 loop = self._loop
130 assert loop is not None
131 now = loop.time()
132 when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold)
133 self._heartbeat_when = when
134 if self._heartbeat_cb is None:
135 # We do not cancel the previous heartbeat_cb here because
136 # it generates a significant amount of TimerHandle churn
137 # which causes asyncio to rebuild the heap frequently.
138 # Instead _send_heartbeat() will reschedule the next
139 # heartbeat if it fires too early.
140 self._heartbeat_cb = loop.call_at(when, self._send_heartbeat)
141
142 def _send_heartbeat(self) -> None:
143 self._heartbeat_cb = None
144 loop = self._loop
145 assert loop is not None and self._writer is not None
146 now = loop.time()
147 if now < self._heartbeat_when:
148 # Heartbeat fired too early, reschedule
149 self._heartbeat_cb = loop.call_at(
150 self._heartbeat_when, self._send_heartbeat
151 )
152 return
153
154 req = self._req
155 timeout_ceil_threshold = (
156 req._protocol._timeout_ceil_threshold if req is not None else 5
157 )
158 when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold)
159 self._cancel_pong_response_cb()
160 self._pong_response_cb = loop.call_at(when, self._pong_not_received)
161
162 coro = self._writer.send_frame(b"", WSMsgType.PING)
163 if sys.version_info >= (3, 12):
164 # Optimization for Python 3.12, try to send the ping
165 # immediately to avoid having to schedule
166 # the task on the event loop.
167 ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
168 else:
169 ping_task = loop.create_task(coro)
170
171 if not ping_task.done():
172 self._ping_task = ping_task
173 ping_task.add_done_callback(self._ping_task_done)
174 else:
175 self._ping_task_done(ping_task)
176
177 def _ping_task_done(self, task: "asyncio.Task[None]") -> None:
178 """Callback for when the ping task completes."""
179 if not task.cancelled() and (exc := task.exception()):
180 self._handle_ping_pong_exception(exc)
181 self._ping_task = None
182
183 def _pong_not_received(self) -> None:
184 if self._req is not None and self._req.transport is not None:
185 self._handle_ping_pong_exception(
186 asyncio.TimeoutError(
187 f"No PONG received after {self._pong_heartbeat} seconds"
188 )
189 )
190
191 def _handle_ping_pong_exception(self, exc: BaseException) -> None:
192 """Handle exceptions raised during ping/pong processing."""
193 if self._closed:
194 return
195 self._set_closed()
196 self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
197 self._exception = exc
198 if self._waiting and not self._closing and self._reader is not None:
199 self._reader.feed_data(WSMessage(WSMsgType.ERROR, exc, None), 0)
200
201 def _set_closed(self) -> None:
202 """Set the connection to closed.
203
204 Cancel any heartbeat timers and set the closed flag.
205 """
206 self._closed = True
207 self._cancel_heartbeat()
208
209 async def prepare(self, request: BaseRequest) -> AbstractStreamWriter:
210 # make pre-check to don't hide it by do_handshake() exceptions
211 if self._payload_writer is not None:
212 return self._payload_writer
213
214 protocol, writer = self._pre_start(request)
215 payload_writer = await super().prepare(request)
216 assert payload_writer is not None
217 self._post_start(request, protocol, writer)
218 await payload_writer.drain()
219 return payload_writer
220
221 def _handshake(
222 self, request: BaseRequest
223 ) -> Tuple["CIMultiDict[str]", Optional[str], int, bool]:
224 headers = request.headers
225 if "websocket" != headers.get(hdrs.UPGRADE, "").lower().strip():
226 raise HTTPBadRequest(
227 text=(
228 "No WebSocket UPGRADE hdr: {}\n Can "
229 '"Upgrade" only to "WebSocket".'
230 ).format(headers.get(hdrs.UPGRADE))
231 )
232
233 if "upgrade" not in headers.get(hdrs.CONNECTION, "").lower():
234 raise HTTPBadRequest(
235 text="No CONNECTION upgrade hdr: {}".format(
236 headers.get(hdrs.CONNECTION)
237 )
238 )
239
240 # find common sub-protocol between client and server
241 protocol: Optional[str] = None
242 if hdrs.SEC_WEBSOCKET_PROTOCOL in headers:
243 req_protocols = [
244 str(proto.strip())
245 for proto in headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(",")
246 ]
247
248 for proto in req_protocols:
249 if proto in self._protocols:
250 protocol = proto
251 break
252 else:
253 # No overlap found: Return no protocol as per spec
254 ws_logger.warning(
255 "%s: Client protocols %r don’t overlap server-known ones %r",
256 request.remote,
257 req_protocols,
258 self._protocols,
259 )
260
261 # check supported version
262 version = headers.get(hdrs.SEC_WEBSOCKET_VERSION, "")
263 if version not in ("13", "8", "7"):
264 raise HTTPBadRequest(text=f"Unsupported version: {version}")
265
266 # check client handshake for validity
267 key = headers.get(hdrs.SEC_WEBSOCKET_KEY)
268 try:
269 if not key or len(base64.b64decode(key)) != 16:
270 raise HTTPBadRequest(text=f"Handshake error: {key!r}")
271 except binascii.Error:
272 raise HTTPBadRequest(text=f"Handshake error: {key!r}") from None
273
274 accept_val = base64.b64encode(
275 hashlib.sha1(key.encode() + WS_KEY).digest()
276 ).decode()
277 response_headers = CIMultiDict(
278 {
279 hdrs.UPGRADE: "websocket",
280 hdrs.CONNECTION: "upgrade",
281 hdrs.SEC_WEBSOCKET_ACCEPT: accept_val,
282 }
283 )
284
285 notakeover = False
286 compress = 0
287 if self._compress:
288 extensions = headers.get(hdrs.SEC_WEBSOCKET_EXTENSIONS)
289 # Server side always get return with no exception.
290 # If something happened, just drop compress extension
291 compress, notakeover = ws_ext_parse(extensions, isserver=True)
292 if compress:
293 enabledext = ws_ext_gen(
294 compress=compress, isserver=True, server_notakeover=notakeover
295 )
296 response_headers[hdrs.SEC_WEBSOCKET_EXTENSIONS] = enabledext
297
298 if protocol:
299 response_headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = protocol
300 return (
301 response_headers,
302 protocol,
303 compress,
304 notakeover,
305 )
306
307 def _pre_start(self, request: BaseRequest) -> Tuple[Optional[str], WebSocketWriter]:
308 self._loop = request._loop
309
310 headers, protocol, compress, notakeover = self._handshake(request)
311
312 self.set_status(101)
313 self.headers.update(headers)
314 self.force_close()
315 self._compress = compress
316 transport = request._protocol.transport
317 assert transport is not None
318 writer = WebSocketWriter(
319 request._protocol,
320 transport,
321 compress=compress,
322 notakeover=notakeover,
323 limit=self._writer_limit,
324 )
325
326 return protocol, writer
327
328 def _post_start(
329 self, request: BaseRequest, protocol: Optional[str], writer: WebSocketWriter
330 ) -> None:
331 self._ws_protocol = protocol
332 self._writer = writer
333
334 self._reset_heartbeat()
335
336 loop = self._loop
337 assert loop is not None
338 self._reader = WebSocketDataQueue(request._protocol, 2**16, loop=loop)
339 request.protocol.set_parser(
340 WebSocketReader(
341 self._reader, self._max_msg_size, compress=bool(self._compress)
342 )
343 )
344 # disable HTTP keepalive for WebSocket
345 request.protocol.keep_alive(False)
346
347 def can_prepare(self, request: BaseRequest) -> WebSocketReady:
348 if self._writer is not None:
349 raise RuntimeError("Already started")
350 try:
351 _, protocol, _, _ = self._handshake(request)
352 except HTTPException:
353 return WebSocketReady(False, None)
354 else:
355 return WebSocketReady(True, protocol)
356
357 @property
358 def prepared(self) -> bool:
359 return self._writer is not None
360
361 @property
362 def closed(self) -> bool:
363 return self._closed
364
365 @property
366 def close_code(self) -> Optional[int]:
367 return self._close_code
368
369 @property
370 def ws_protocol(self) -> Optional[str]:
371 return self._ws_protocol
372
373 @property
374 def compress(self) -> Union[int, bool]:
375 return self._compress
376
377 def get_extra_info(self, name: str, default: Any = None) -> Any:
378 """Get optional transport information.
379
380 If no value associated with ``name`` is found, ``default`` is returned.
381 """
382 writer = self._writer
383 if writer is None:
384 return default
385 transport = writer.transport
386 if transport is None:
387 return default
388 return transport.get_extra_info(name, default)
389
390 def exception(self) -> Optional[BaseException]:
391 return self._exception
392
393 async def ping(self, message: bytes = b"") -> None:
394 if self._writer is None:
395 raise RuntimeError("Call .prepare() first")
396 await self._writer.send_frame(message, WSMsgType.PING)
397
398 async def pong(self, message: bytes = b"") -> None:
399 # unsolicited pong
400 if self._writer is None:
401 raise RuntimeError("Call .prepare() first")
402 await self._writer.send_frame(message, WSMsgType.PONG)
403
404 async def send_frame(
405 self, message: bytes, opcode: WSMsgType, compress: Optional[int] = None
406 ) -> None:
407 """Send a frame over the websocket."""
408 if self._writer is None:
409 raise RuntimeError("Call .prepare() first")
410 await self._writer.send_frame(message, opcode, compress)
411
412 async def send_str(self, data: str, compress: Optional[int] = None) -> None:
413 if self._writer is None:
414 raise RuntimeError("Call .prepare() first")
415 if not isinstance(data, str):
416 raise TypeError("data argument must be str (%r)" % type(data))
417 await self._writer.send_frame(
418 data.encode("utf-8"), WSMsgType.TEXT, compress=compress
419 )
420
421 async def send_bytes(self, data: bytes, compress: Optional[int] = None) -> None:
422 if self._writer is None:
423 raise RuntimeError("Call .prepare() first")
424 if not isinstance(data, (bytes, bytearray, memoryview)):
425 raise TypeError("data argument must be byte-ish (%r)" % type(data))
426 await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress)
427
428 async def send_json(
429 self,
430 data: Any,
431 compress: Optional[int] = None,
432 *,
433 dumps: JSONEncoder = json.dumps,
434 ) -> None:
435 await self.send_str(dumps(data), compress=compress)
436
437 async def write_eof(self) -> None: # type: ignore[override]
438 if self._eof_sent:
439 return
440 if self._payload_writer is None:
441 raise RuntimeError("Response has not been started")
442
443 await self.close()
444 self._eof_sent = True
445
446 async def close(
447 self, *, code: int = WSCloseCode.OK, message: bytes = b"", drain: bool = True
448 ) -> bool:
449 """Close websocket connection."""
450 if self._writer is None:
451 raise RuntimeError("Call .prepare() first")
452
453 if self._closed:
454 return False
455 self._set_closed()
456
457 try:
458 await self._writer.close(code, message)
459 writer = self._payload_writer
460 assert writer is not None
461 if drain:
462 await writer.drain()
463 except (asyncio.CancelledError, asyncio.TimeoutError):
464 self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
465 raise
466 except Exception as exc:
467 self._exception = exc
468 self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
469 return True
470
471 reader = self._reader
472 assert reader is not None
473 # we need to break `receive()` cycle before we can call
474 # `reader.read()` as `close()` may be called from different task
475 if self._waiting:
476 assert self._loop is not None
477 assert self._close_wait is None
478 self._close_wait = self._loop.create_future()
479 reader.feed_data(WS_CLOSING_MESSAGE, 0)
480 await self._close_wait
481
482 if self._closing:
483 self._close_transport()
484 return True
485
486 try:
487 async with async_timeout.timeout(self._timeout):
488 while True:
489 msg = await reader.read()
490 if msg.type is WSMsgType.CLOSE:
491 self._set_code_close_transport(msg.data)
492 return True
493 except asyncio.CancelledError:
494 self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
495 raise
496 except Exception as exc:
497 self._exception = exc
498 self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
499 return True
500
501 def _set_closing(self, code: WSCloseCode) -> None:
502 """Set the close code and mark the connection as closing."""
503 self._closing = True
504 self._close_code = code
505 self._cancel_heartbeat()
506
507 def _set_code_close_transport(self, code: WSCloseCode) -> None:
508 """Set the close code and close the transport."""
509 self._close_code = code
510 self._close_transport()
511
512 def _close_transport(self) -> None:
513 """Close the transport."""
514 if self._req is not None and self._req.transport is not None:
515 self._req.transport.close()
516
517 async def receive(self, timeout: Optional[float] = None) -> WSMessage:
518 if self._reader is None:
519 raise RuntimeError("Call .prepare() first")
520
521 receive_timeout = timeout or self._receive_timeout
522 while True:
523 if self._waiting:
524 raise RuntimeError("Concurrent call to receive() is not allowed")
525
526 if self._closed:
527 self._conn_lost += 1
528 if self._conn_lost >= THRESHOLD_CONNLOST_ACCESS:
529 raise RuntimeError("WebSocket connection is closed.")
530 return WS_CLOSED_MESSAGE
531 elif self._closing:
532 return WS_CLOSING_MESSAGE
533
534 try:
535 self._waiting = True
536 try:
537 if receive_timeout:
538 # Entering the context manager and creating
539 # Timeout() object can take almost 50% of the
540 # run time in this loop so we avoid it if
541 # there is no read timeout.
542 async with async_timeout.timeout(receive_timeout):
543 msg = await self._reader.read()
544 else:
545 msg = await self._reader.read()
546 self._reset_heartbeat()
547 finally:
548 self._waiting = False
549 if self._close_wait:
550 set_result(self._close_wait, None)
551 except asyncio.TimeoutError:
552 raise
553 except EofStream:
554 self._close_code = WSCloseCode.OK
555 await self.close()
556 return WSMessage(WSMsgType.CLOSED, None, None)
557 except WebSocketError as exc:
558 self._close_code = exc.code
559 await self.close(code=exc.code)
560 return WSMessage(WSMsgType.ERROR, exc, None)
561 except Exception as exc:
562 self._exception = exc
563 self._set_closing(WSCloseCode.ABNORMAL_CLOSURE)
564 await self.close()
565 return WSMessage(WSMsgType.ERROR, exc, None)
566
567 if msg.type not in _INTERNAL_RECEIVE_TYPES:
568 # If its not a close/closing/ping/pong message
569 # we can return it immediately
570 return msg
571
572 if msg.type is WSMsgType.CLOSE:
573 self._set_closing(msg.data)
574 # Could be closed while awaiting reader.
575 if not self._closed and self._autoclose:
576 # The client is likely going to close the
577 # connection out from under us so we do not
578 # want to drain any pending writes as it will
579 # likely result writing to a broken pipe.
580 await self.close(drain=False)
581 elif msg.type is WSMsgType.CLOSING:
582 self._set_closing(WSCloseCode.OK)
583 elif msg.type is WSMsgType.PING and self._autoping:
584 await self.pong(msg.data)
585 continue
586 elif msg.type is WSMsgType.PONG and self._autoping:
587 continue
588
589 return msg
590
591 async def receive_str(self, *, timeout: Optional[float] = None) -> str:
592 msg = await self.receive(timeout)
593 if msg.type is not WSMsgType.TEXT:
594 raise WSMessageTypeError(
595 f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT"
596 )
597 return cast(str, msg.data)
598
599 async def receive_bytes(self, *, timeout: Optional[float] = None) -> bytes:
600 msg = await self.receive(timeout)
601 if msg.type is not WSMsgType.BINARY:
602 raise WSMessageTypeError(
603 f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY"
604 )
605 return cast(bytes, msg.data)
606
607 async def receive_json(
608 self, *, loads: JSONDecoder = json.loads, timeout: Optional[float] = None
609 ) -> Any:
610 data = await self.receive_str(timeout=timeout)
611 return loads(data)
612
613 async def write(self, data: bytes) -> None:
614 raise RuntimeError("Cannot call .write() for websocket")
615
616 def __aiter__(self) -> "WebSocketResponse":
617 return self
618
619 async def __anext__(self) -> WSMessage:
620 msg = await self.receive()
621 if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
622 raise StopAsyncIteration
623 return msg
624
625 def _cancel(self, exc: BaseException) -> None:
626 # web_protocol calls this from connection_lost
627 # or when the server is shutting down.
628 self._closing = True
629 self._cancel_heartbeat()
630 if self._reader is not None:
631 set_exception(self._reader, exc)
632