chore: rename project coven → hack-house ⛧
Rebrand the Rust client crate (coven/ → hh/, package+binary "hack-house"), README, CLI strings, and branch (coven → hack-house). Gitea repo renamed cmd-chat → hack-house to match. Crypto/server logic unchanged; selftest + golden-vector test still green, binary is now `hack-house`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
leetcrypt
@@ -0,0 +1,15 @@
|
||||
from sanic.models.server_types import ConnInfo, Signal
|
||||
from sanic.server.async_server import AsyncioServer
|
||||
from sanic.server.loop import try_use_uvloop
|
||||
from sanic.server.protocols.http_protocol import HttpProtocol
|
||||
from sanic.server.runners import serve
|
||||
|
||||
|
||||
__all__ = (
|
||||
"AsyncioServer",
|
||||
"ConnInfo",
|
||||
"HttpProtocol",
|
||||
"Signal",
|
||||
"serve",
|
||||
"try_use_uvloop",
|
||||
)
|
||||
@@ -0,0 +1,112 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sanic.exceptions import SanicException
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sanic import Sanic
|
||||
|
||||
|
||||
class AsyncioServer:
|
||||
"""Wraps an asyncio server with functionality that might be useful to a user who needs to manage the server lifecycle manually.""" # noqa: E501
|
||||
|
||||
__slots__ = ("app", "connections", "loop", "serve_coro", "server")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app: Sanic,
|
||||
loop,
|
||||
serve_coro,
|
||||
connections,
|
||||
):
|
||||
# Note, Sanic already called "before_server_start" events
|
||||
# before this helper was even created. So we don't need it here.
|
||||
self.app = app
|
||||
self.connections = connections
|
||||
self.loop = loop
|
||||
self.serve_coro = serve_coro
|
||||
self.server = None
|
||||
|
||||
def startup(self):
|
||||
"""Trigger "startup" operations on the app"""
|
||||
return self.app._startup()
|
||||
|
||||
def before_start(self):
|
||||
"""Trigger "before_server_start" events"""
|
||||
return self._server_event("init", "before")
|
||||
|
||||
def after_start(self):
|
||||
"""Trigger "after_server_start" events"""
|
||||
return self._server_event("init", "after")
|
||||
|
||||
def before_stop(self):
|
||||
"""Trigger "before_server_stop" events"""
|
||||
return self._server_event("shutdown", "before")
|
||||
|
||||
def after_stop(self):
|
||||
"""Trigger "after_server_stop" events"""
|
||||
return self._server_event("shutdown", "after")
|
||||
|
||||
def is_serving(self) -> bool:
|
||||
"""Returns True if the server is running, False otherwise"""
|
||||
if self.server:
|
||||
return self.server.is_serving()
|
||||
return False
|
||||
|
||||
def wait_closed(self):
|
||||
"""Wait until the server is closed"""
|
||||
if self.server:
|
||||
return self.server.wait_closed()
|
||||
|
||||
def close(self):
|
||||
"""Close the server"""
|
||||
if self.server:
|
||||
self.server.close()
|
||||
coro = self.wait_closed()
|
||||
task = self.loop.create_task(coro)
|
||||
return task
|
||||
|
||||
def start_serving(self):
|
||||
"""Start serving requests"""
|
||||
return self._serve(self.server.start_serving)
|
||||
|
||||
def serve_forever(self):
|
||||
"""Serve requests until the server is stopped"""
|
||||
return self._serve(self.server.serve_forever)
|
||||
|
||||
def _serve(self, serve_func):
|
||||
if self.server:
|
||||
if not self.app.state.is_started:
|
||||
raise SanicException(
|
||||
"Cannot run Sanic server without first running "
|
||||
"await server.startup()"
|
||||
)
|
||||
|
||||
try:
|
||||
return serve_func()
|
||||
except AttributeError:
|
||||
name = serve_func.__name__
|
||||
raise NotImplementedError(
|
||||
f"server.{name} not available in this version "
|
||||
"of asyncio or uvloop."
|
||||
)
|
||||
|
||||
def _server_event(self, concern: str, action: str):
|
||||
if not self.app.state.is_started:
|
||||
raise SanicException(
|
||||
"Cannot dispatch server event without "
|
||||
"first running await server.startup()"
|
||||
)
|
||||
return self.app._server_event(concern, action, loop=self.loop)
|
||||
|
||||
def __await__(self):
|
||||
"""
|
||||
Starts the asyncio server, returns AsyncServerCoro
|
||||
"""
|
||||
task = self.loop.create_task(self.serve_coro)
|
||||
while not task.done():
|
||||
yield
|
||||
self.server = task.result()
|
||||
return self
|
||||
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from inspect import isawaitable
|
||||
from typing import TYPE_CHECKING, Any, Callable
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sanic import Sanic
|
||||
|
||||
|
||||
def trigger_events(
|
||||
events: Iterable[Callable[..., Any]] | None,
|
||||
loop,
|
||||
app: Sanic | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Trigger event callbacks (functions or async)
|
||||
|
||||
Args:
|
||||
events (Optional[Iterable[Callable[..., Any]]]): [description]
|
||||
loop ([type]): [description]
|
||||
app (Optional[Sanic], optional): [description]. Defaults to None.
|
||||
"""
|
||||
if events:
|
||||
for event in events:
|
||||
try:
|
||||
result = event(**kwargs) if not app else event(app, **kwargs)
|
||||
except TypeError:
|
||||
result = (
|
||||
event(loop, **kwargs)
|
||||
if not app
|
||||
else event(app, loop, **kwargs)
|
||||
)
|
||||
if isawaitable(result):
|
||||
loop.run_until_complete(result)
|
||||
@@ -0,0 +1,31 @@
|
||||
# flake8: noqa: E501
|
||||
|
||||
import random
|
||||
import sys
|
||||
|
||||
|
||||
# fmt: off
|
||||
ascii_phrases = {
|
||||
'Farewell', 'Bye', 'See you later', 'Take care', 'So long', 'Adieu', 'Cheerio',
|
||||
'Goodbye', 'Adios', 'Au revoir', 'Arrivederci', 'Sayonara', 'Auf Wiedersehen',
|
||||
'Do svidaniya', 'Annyeong', 'Tot ziens', 'Ha det', 'Selamat tinggal',
|
||||
'Hasta luego', 'Nos vemos', 'Salut', 'Ciao', 'A presto',
|
||||
'Dag', 'Tot later', 'Vi ses', 'Sampai jumpa',
|
||||
}
|
||||
|
||||
non_ascii_phrases = {
|
||||
'Tschüss', 'Zài jiàn', 'Bāi bāi', 'Míngtiān jiàn', 'Adeus', 'Tchau', 'Até logo',
|
||||
'Hejdå', 'À bientôt', 'Bis später', 'Adjø',
|
||||
'じゃね', 'またね', '안녕히 계세요', '잘 가', 'שלום',
|
||||
'להתראות', 'مع السلامة', 'إلى اللقاء', 'وداعاً', 'अलविदा',
|
||||
'फिर मिलेंगे',
|
||||
}
|
||||
|
||||
all_phrases = ascii_phrases | non_ascii_phrases
|
||||
# fmt: on
|
||||
|
||||
|
||||
def get_goodbye() -> str: # pragma: no cover
|
||||
is_utf8 = sys.stdout.encoding.lower() == "utf-8"
|
||||
phrases = all_phrases if is_utf8 else ascii_phrases
|
||||
return random.choice(list(phrases)) # nosec: B311
|
||||
@@ -0,0 +1,64 @@
|
||||
import asyncio
|
||||
|
||||
from os import getenv
|
||||
|
||||
from sanic.compat import OS_IS_WINDOWS
|
||||
from sanic.log import error_logger
|
||||
from sanic.utils import str_to_bool
|
||||
|
||||
|
||||
def try_use_uvloop() -> None:
|
||||
"""Use uvloop instead of the default asyncio loop."""
|
||||
if OS_IS_WINDOWS:
|
||||
error_logger.warning(
|
||||
"You are trying to use uvloop, but uvloop is not compatible "
|
||||
"with your system. You can disable uvloop completely by setting "
|
||||
"the 'USE_UVLOOP' configuration value to false, or simply not "
|
||||
"defining it and letting Sanic handle it for you. Sanic will now "
|
||||
"continue to run using the default event loop."
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
import uvloop # type: ignore
|
||||
except ImportError:
|
||||
error_logger.warning(
|
||||
"You are trying to use uvloop, but uvloop is not "
|
||||
"installed in your system. In order to use uvloop "
|
||||
"you must first install it. Otherwise, you can disable "
|
||||
"uvloop completely by setting the 'USE_UVLOOP' "
|
||||
"configuration value to false. Sanic will now continue "
|
||||
"to run with the default event loop."
|
||||
)
|
||||
return
|
||||
|
||||
uvloop_install_removed = str_to_bool(getenv("SANIC_NO_UVLOOP", "no"))
|
||||
if uvloop_install_removed:
|
||||
error_logger.info(
|
||||
"You are requesting to run Sanic using uvloop, but the "
|
||||
"install-time 'SANIC_NO_UVLOOP' environment variable (used to "
|
||||
"opt-out of installing uvloop with Sanic) is set to true. If "
|
||||
"you want to prevent Sanic from overriding the event loop policy "
|
||||
"during runtime, set the 'USE_UVLOOP' configuration value to "
|
||||
"false."
|
||||
)
|
||||
|
||||
if not isinstance(asyncio.get_event_loop_policy(), uvloop.EventLoopPolicy):
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
|
||||
|
||||
def try_windows_loop():
|
||||
"""Try to use the WindowsSelectorEventLoopPolicy instead of the default"""
|
||||
if not OS_IS_WINDOWS:
|
||||
error_logger.warning(
|
||||
"You are trying to use an event loop policy that is not "
|
||||
"compatible with your system. You can simply let Sanic handle "
|
||||
"selecting the best loop for you. Sanic will now continue to run "
|
||||
"using the default event loop."
|
||||
)
|
||||
return
|
||||
|
||||
if not isinstance(
|
||||
asyncio.get_event_loop_policy(), asyncio.WindowsSelectorEventLoopPolicy
|
||||
):
|
||||
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
||||
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sanic.exceptions import RequestCancelled
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sanic.app import Sanic
|
||||
|
||||
import asyncio
|
||||
|
||||
from asyncio.transports import Transport
|
||||
from time import monotonic as current_time
|
||||
|
||||
from sanic.log import error_logger
|
||||
from sanic.models.server_types import ConnInfo, Signal
|
||||
|
||||
|
||||
class SanicProtocol(asyncio.Protocol):
|
||||
__slots__ = (
|
||||
"app",
|
||||
# event loop, connection
|
||||
"loop",
|
||||
"transport",
|
||||
"connections",
|
||||
"conn_info",
|
||||
"signal",
|
||||
"_can_write",
|
||||
"_time",
|
||||
"_task",
|
||||
"_unix",
|
||||
"_data_received",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
loop,
|
||||
app: Sanic,
|
||||
signal=None,
|
||||
connections=None,
|
||||
unix=None,
|
||||
**kwargs,
|
||||
):
|
||||
asyncio.set_event_loop(loop)
|
||||
self.loop = loop
|
||||
self.app: Sanic = app
|
||||
self.signal = signal or Signal()
|
||||
self.transport: Transport | None = None
|
||||
self.connections = connections if connections is not None else set()
|
||||
self.conn_info: ConnInfo | None = None
|
||||
self._can_write = asyncio.Event()
|
||||
self._can_write.set()
|
||||
self._unix = unix
|
||||
self._time = 0.0 # type: float
|
||||
self._task: asyncio.Task | None = None
|
||||
self._data_received = asyncio.Event()
|
||||
|
||||
@property
|
||||
def ctx(self):
|
||||
if self.conn_info is not None:
|
||||
return self.conn_info.ctx
|
||||
else:
|
||||
return None
|
||||
|
||||
async def send(self, data):
|
||||
"""
|
||||
Generic data write implementation with backpressure control.
|
||||
"""
|
||||
await self._can_write.wait()
|
||||
if self.transport.is_closing():
|
||||
raise RequestCancelled
|
||||
self.transport.write(data)
|
||||
self._time = current_time()
|
||||
|
||||
async def receive_more(self):
|
||||
"""
|
||||
Wait until more data is received into the Server protocol's buffer
|
||||
"""
|
||||
self.transport.resume_reading()
|
||||
self._data_received.clear()
|
||||
await self._data_received.wait()
|
||||
|
||||
def close(self, timeout: float | None = None):
|
||||
"""
|
||||
Attempt close the connection.
|
||||
"""
|
||||
if self.transport is None or self.transport.is_closing():
|
||||
# do not attempt to close again, already aborted or closing
|
||||
return
|
||||
|
||||
# Check if write is already paused _before_ close() is called.
|
||||
write_was_paused = not self._can_write.is_set()
|
||||
# Trigger the UVLoop Stream Transport Close routine
|
||||
# Causes a call to connection_lost where further cleanup occurs
|
||||
# Close may fully close the connection now, but if there is still
|
||||
# data in the libuv buffer, then close becomes an async operation
|
||||
self.transport.close()
|
||||
try:
|
||||
# Check write-buffer data left _after_ close is called.
|
||||
# in UVLoop, get the data in the libuv transport write-buffer
|
||||
data_left = self.transport.get_write_buffer_size()
|
||||
# Some asyncio implementations don't support get_write_buffer_size
|
||||
except (AttributeError, NotImplementedError):
|
||||
data_left = 0
|
||||
if write_was_paused or data_left > 0:
|
||||
# don't call resume_writing here, it gets called by the transport
|
||||
# to unpause the protocol when it is ready for more data
|
||||
|
||||
# Schedule the async close checker, to close the connection
|
||||
# after the transport is done, and clean everything up.
|
||||
if timeout is None:
|
||||
# This close timeout needs to be less than the graceful
|
||||
# shutdown timeout. The graceful shutdown _could_ be waiting
|
||||
# for this transport to close before shutting down the app.
|
||||
timeout = self.app.config.GRACEFUL_TCP_CLOSE_TIMEOUT
|
||||
# This is 5s by default.
|
||||
else:
|
||||
# Schedule the async close checker but with no timeout,
|
||||
# this will ensure abort() is called if required.
|
||||
if timeout is None:
|
||||
timeout = 0
|
||||
self.loop.call_soon(
|
||||
_async_protocol_transport_close,
|
||||
self,
|
||||
self.loop,
|
||||
timeout,
|
||||
)
|
||||
|
||||
def abort(self):
|
||||
"""
|
||||
Force close the connection.
|
||||
"""
|
||||
# Cause a call to connection_lost where further cleanup occurs
|
||||
if self.transport:
|
||||
self.transport.abort()
|
||||
self.transport = None
|
||||
|
||||
# asyncio.Protocol API Callbacks #
|
||||
# ------------------------------ #
|
||||
def connection_made(self, transport):
|
||||
"""
|
||||
Generic connection-made, with no connection_task, and no recv_buffer.
|
||||
Override this for protocol-specific connection implementations.
|
||||
"""
|
||||
try:
|
||||
transport.set_write_buffer_limits(low=16384, high=65536)
|
||||
self.connections.add(self)
|
||||
self.transport = transport
|
||||
self.conn_info = ConnInfo(self.transport, unix=self._unix)
|
||||
except Exception:
|
||||
error_logger.exception("protocol.connect_made")
|
||||
|
||||
def connection_lost(self, exc):
|
||||
"""
|
||||
This is a callback handler that is called from the asyncio
|
||||
transport layer implementation (eg, UVLoop's UVStreamTransport).
|
||||
It is scheduled to be called async after the transport has closed.
|
||||
When data is still in the send buffer, this call to connection_lost
|
||||
will be delayed until _after_ the buffer is finished being sent.
|
||||
|
||||
So we can use this callback as a confirmation callback
|
||||
that the async write-buffer transfer is finished.
|
||||
"""
|
||||
try:
|
||||
self.connections.discard(self)
|
||||
# unblock the send queue if it is paused,
|
||||
# this allows the route handler to see
|
||||
# the CancelledError exception
|
||||
self.resume_writing()
|
||||
self.conn_info.lost = True
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
except BaseException:
|
||||
error_logger.exception("protocol.connection_lost")
|
||||
|
||||
def pause_writing(self):
|
||||
self._can_write.clear()
|
||||
|
||||
def resume_writing(self):
|
||||
self._can_write.set()
|
||||
|
||||
def data_received(self, data: bytes):
|
||||
try:
|
||||
self._time = current_time()
|
||||
if not data:
|
||||
return self.close()
|
||||
|
||||
if self._data_received:
|
||||
self._data_received.set()
|
||||
except BaseException:
|
||||
error_logger.exception("protocol.data_received")
|
||||
|
||||
|
||||
def _async_protocol_transport_close(
|
||||
protocol: SanicProtocol,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
timeout: float,
|
||||
):
|
||||
"""
|
||||
This function is scheduled to be called after close() is called.
|
||||
It checks that the transport has shut down properly, or waits
|
||||
for any remaining data to be sent, and aborts after a timeout.
|
||||
This is required if the transport is closed while there is an async
|
||||
large async transport write operation in progress.
|
||||
This is observed when NGINX reverse-proxy is the client.
|
||||
"""
|
||||
if protocol.transport is None:
|
||||
# abort() is the only method that can make
|
||||
# protocol.transport be None, so abort was already called
|
||||
return
|
||||
# protocol.connection_lost does not set protocol.transport to None
|
||||
# so to detect it a different way with conninfo.lost
|
||||
elif protocol.conn_info is not None and protocol.conn_info.lost:
|
||||
# Terminus. Most connections finish the protocol here!
|
||||
# Connection_lost callback was executed already,
|
||||
# so transport did complete and close itself properly.
|
||||
# No need to call abort().
|
||||
|
||||
# This is the last part of cleanup to do
|
||||
# that is not done by connection_lost handler.
|
||||
# Ensure transport is cleaned up by GC.
|
||||
protocol.transport = None
|
||||
return
|
||||
elif not protocol.transport.is_closing():
|
||||
raise RuntimeError(
|
||||
"You must call transport.close() before "
|
||||
"protocol._async_transport_close() runs."
|
||||
)
|
||||
|
||||
write_is_paused = not protocol._can_write.is_set()
|
||||
try:
|
||||
# in UVLoop, get the data in the libuv write-buffer
|
||||
data_left = protocol.transport.get_write_buffer_size()
|
||||
# Some asyncio implementations don't support get_write_buffer_size
|
||||
except (AttributeError, NotImplementedError):
|
||||
data_left = 0
|
||||
if write_is_paused or data_left > 0:
|
||||
# don't need to call resume_writing here to unpause
|
||||
if timeout <= 0:
|
||||
# timeout is 0 or less, so we can simply abort now
|
||||
loop.call_soon(SanicProtocol.abort, protocol)
|
||||
else:
|
||||
next_check_interval = min(timeout, 0.1)
|
||||
next_check_timeout = timeout - next_check_interval
|
||||
loop.call_later(
|
||||
# Recurse back in after the timeout, to check again
|
||||
next_check_interval,
|
||||
# this next time with reduced timeout.
|
||||
_async_protocol_transport_close,
|
||||
protocol,
|
||||
loop,
|
||||
next_check_timeout,
|
||||
)
|
||||
else:
|
||||
# Not paused, and no data left in the buffer, but transport
|
||||
# is still open, connection_lost has not been called yet.
|
||||
# We can call abort() to fix that.
|
||||
loop.call_soon(SanicProtocol.abort, protocol)
|
||||
@@ -0,0 +1,360 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sanic.http.constants import HTTP
|
||||
from sanic.http.http3 import Http3
|
||||
from sanic.touchup.meta import TouchUpMeta
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sanic.app import Sanic
|
||||
|
||||
|
||||
from asyncio import CancelledError
|
||||
from time import monotonic as current_time
|
||||
|
||||
from sanic.exceptions import (
|
||||
RequestCancelled,
|
||||
RequestTimeout,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
from sanic.http import Http, Stage
|
||||
from sanic.log import (
|
||||
Colors,
|
||||
access_logger,
|
||||
error_logger,
|
||||
logger,
|
||||
websockets_logger,
|
||||
)
|
||||
from sanic.models.server_types import ConnInfo
|
||||
from sanic.request import Request
|
||||
from sanic.server.protocols.base_protocol import SanicProtocol
|
||||
|
||||
|
||||
ConnectionProtocol = type("ConnectionProtocol", (), {})
|
||||
try:
|
||||
from aioquic.asyncio import QuicConnectionProtocol
|
||||
from aioquic.h3.connection import H3_ALPN, H3Connection
|
||||
from aioquic.quic.events import (
|
||||
DatagramFrameReceived,
|
||||
ProtocolNegotiated,
|
||||
QuicEvent,
|
||||
)
|
||||
|
||||
ConnectionProtocol = QuicConnectionProtocol
|
||||
except ModuleNotFoundError: # no cov
|
||||
...
|
||||
|
||||
|
||||
class HttpProtocolMixin:
|
||||
__slots__ = ()
|
||||
__version__: HTTP
|
||||
|
||||
def _setup_connection(self, *args, **kwargs):
|
||||
self._http = self.HTTP_CLASS(self, *args, **kwargs)
|
||||
self._time = current_time()
|
||||
try:
|
||||
self.check_timeouts()
|
||||
except AttributeError:
|
||||
...
|
||||
|
||||
def _setup(self):
|
||||
self.request: Request | None = None
|
||||
self.access_log = self.app.config.ACCESS_LOG
|
||||
self.request_handler = self.app.handle_request
|
||||
self.error_handler = self.app.error_handler
|
||||
self.request_timeout = self.app.config.REQUEST_TIMEOUT
|
||||
self.response_timeout = self.app.config.RESPONSE_TIMEOUT
|
||||
self.keep_alive_timeout = self.app.config.KEEP_ALIVE_TIMEOUT
|
||||
self.request_max_size = self.app.config.REQUEST_MAX_SIZE
|
||||
self.request_class = self.app.request_class or Request
|
||||
|
||||
@property
|
||||
def http(self):
|
||||
if not hasattr(self, "_http"):
|
||||
return None
|
||||
return self._http
|
||||
|
||||
@property
|
||||
def version(self):
|
||||
return self.__class__.__version__
|
||||
|
||||
|
||||
class HttpProtocol(HttpProtocolMixin, SanicProtocol, metaclass=TouchUpMeta):
|
||||
"""
|
||||
This class provides implements the HTTP 1.1 protocol on top of our
|
||||
Sanic Server transport
|
||||
"""
|
||||
|
||||
HTTP_CLASS = Http
|
||||
|
||||
__touchup__ = (
|
||||
"send",
|
||||
"connection_task",
|
||||
)
|
||||
__version__ = HTTP.VERSION_1
|
||||
__slots__ = (
|
||||
# request params
|
||||
"request",
|
||||
# request config
|
||||
"request_handler",
|
||||
"request_timeout",
|
||||
"response_timeout",
|
||||
"keep_alive_timeout",
|
||||
"request_max_size",
|
||||
"request_class",
|
||||
"error_handler",
|
||||
# enable or disable access log purpose
|
||||
"access_log",
|
||||
# connection management
|
||||
"state",
|
||||
"url",
|
||||
"_handler_task",
|
||||
"_http",
|
||||
"_exception",
|
||||
"recv_buffer",
|
||||
"_callback_check_timeouts",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
loop,
|
||||
app: Sanic,
|
||||
signal=None,
|
||||
connections=None,
|
||||
state=None,
|
||||
unix=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
loop=loop,
|
||||
app=app,
|
||||
signal=signal,
|
||||
connections=connections,
|
||||
unix=unix,
|
||||
)
|
||||
self.url = None
|
||||
self.state = state if state else {}
|
||||
self._setup()
|
||||
if "requests_count" not in self.state:
|
||||
self.state["requests_count"] = 0
|
||||
self._exception = None
|
||||
self._callback_check_timeouts = None
|
||||
|
||||
async def connection_task(self): # no cov
|
||||
"""
|
||||
Run a HTTP connection.
|
||||
|
||||
Timeouts and some additional error handling occur here, while most of
|
||||
everything else happens in class Http or in code called from there.
|
||||
"""
|
||||
try:
|
||||
self._setup_connection()
|
||||
await self.app.dispatch(
|
||||
"http.lifecycle.begin",
|
||||
inline=True,
|
||||
context={"conn_info": self.conn_info},
|
||||
)
|
||||
await self._http.http1()
|
||||
except CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
error_logger.exception("protocol.connection_task uncaught")
|
||||
finally:
|
||||
if (
|
||||
self.access_log
|
||||
and self._http
|
||||
and self.transport
|
||||
and not self._http.upgrade_websocket
|
||||
):
|
||||
self.log_disconnect()
|
||||
self._http = None
|
||||
self._task = None
|
||||
try:
|
||||
self.close()
|
||||
except BaseException:
|
||||
error_logger.exception("Closing failed")
|
||||
finally:
|
||||
await self.app.dispatch(
|
||||
"http.lifecycle.complete",
|
||||
inline=True,
|
||||
context={"conn_info": self.conn_info},
|
||||
)
|
||||
# Important to keep this Ellipsis here for the TouchUp module
|
||||
...
|
||||
|
||||
def log_disconnect(self):
|
||||
ip = self.transport.get_extra_info("peername")
|
||||
|
||||
req = self._http.request
|
||||
res = self._http.response
|
||||
extra = {
|
||||
"status": res.status if res else str(self._http.stage),
|
||||
"byte": "DISCONNECTED",
|
||||
"host": f"{id(self):X}"[-5:-1] + "unx",
|
||||
"request": "nil",
|
||||
"duration": "",
|
||||
}
|
||||
if req is not None:
|
||||
if ip := req.client_ip:
|
||||
extra["host"] = f"{ip}:{req.port}"
|
||||
extra["request"] = f"{req.method} {req.url}"
|
||||
access_logger.info("", extra=extra)
|
||||
|
||||
def check_timeouts(self):
|
||||
"""
|
||||
Runs itself periodically to enforce any expired timeouts.
|
||||
"""
|
||||
try:
|
||||
if not self._task:
|
||||
return
|
||||
duration = current_time() - self._time
|
||||
stage = self._http.stage
|
||||
if stage is Stage.IDLE and duration > self.keep_alive_timeout:
|
||||
logger.debug("KeepAlive Timeout. Closing connection.")
|
||||
elif stage is Stage.REQUEST and duration > self.request_timeout:
|
||||
logger.debug("Request Timeout. Closing connection.")
|
||||
self._http.exception = RequestTimeout("Request Timeout")
|
||||
elif stage is Stage.HANDLER and self._http.upgrade_websocket:
|
||||
websockets_logger.debug(
|
||||
"Handling websocket. Timeouts disabled."
|
||||
)
|
||||
return
|
||||
elif (
|
||||
stage in (Stage.HANDLER, Stage.RESPONSE, Stage.FAILED)
|
||||
and duration > self.response_timeout
|
||||
):
|
||||
logger.debug("Response Timeout. Closing connection.")
|
||||
self._http.exception = ServiceUnavailable("Response Timeout")
|
||||
else:
|
||||
interval = (
|
||||
min(
|
||||
self.keep_alive_timeout,
|
||||
self.request_timeout,
|
||||
self.response_timeout,
|
||||
)
|
||||
/ 2
|
||||
)
|
||||
_interval = max(0.1, interval)
|
||||
self._callback_check_timeouts = self.loop.call_later(
|
||||
_interval, self.check_timeouts
|
||||
)
|
||||
return
|
||||
if sys.version_info < (3, 14):
|
||||
self._task.cancel("Cancel connection task with a timeout")
|
||||
else:
|
||||
self._task.cancel()
|
||||
except Exception:
|
||||
error_logger.exception("protocol.check_timeouts")
|
||||
|
||||
def close(self, timeout: float | None = None):
|
||||
"""
|
||||
Requires to prevent checking timeouts for closed connections
|
||||
"""
|
||||
|
||||
if self._callback_check_timeouts:
|
||||
self._callback_check_timeouts.cancel()
|
||||
return super().close(timeout=timeout)
|
||||
|
||||
async def send(self, data): # no cov
|
||||
"""
|
||||
Writes HTTP data with backpressure control.
|
||||
"""
|
||||
await self._can_write.wait()
|
||||
if self.transport.is_closing():
|
||||
raise RequestCancelled
|
||||
await self.app.dispatch(
|
||||
"http.lifecycle.send",
|
||||
inline=True,
|
||||
context={"data": data},
|
||||
)
|
||||
self.transport.write(data)
|
||||
self._time = current_time()
|
||||
|
||||
def close_if_idle(self) -> bool:
|
||||
"""
|
||||
Close the connection if a request is not being sent or received
|
||||
|
||||
:return: boolean - True if closed, false if staying open
|
||||
"""
|
||||
if self.http is None or self.http.stage is Stage.IDLE:
|
||||
self.close()
|
||||
return True
|
||||
return False
|
||||
|
||||
# -------------------------------------------- #
|
||||
# Only asyncio.Protocol callbacks below this
|
||||
# -------------------------------------------- #
|
||||
|
||||
def connection_made(self, transport):
|
||||
"""
|
||||
HTTP-protocol-specific new connection handler
|
||||
"""
|
||||
try:
|
||||
# TODO: Benchmark to find suitable write buffer limits
|
||||
transport.set_write_buffer_limits(low=16384, high=65536)
|
||||
self.connections.add(self)
|
||||
self.transport = transport
|
||||
self._task = self.loop.create_task(self.connection_task())
|
||||
self.recv_buffer = bytearray()
|
||||
self.conn_info = ConnInfo(self.transport, unix=self._unix)
|
||||
except Exception:
|
||||
error_logger.exception("protocol.connect_made")
|
||||
|
||||
def data_received(self, data: bytes):
|
||||
try:
|
||||
self._time = current_time()
|
||||
if not data:
|
||||
return self.close()
|
||||
self.recv_buffer += data
|
||||
|
||||
if (
|
||||
len(self.recv_buffer) >= self.app.config.REQUEST_BUFFER_SIZE
|
||||
and self.transport
|
||||
):
|
||||
self.transport.pause_reading()
|
||||
|
||||
if self._data_received:
|
||||
self._data_received.set()
|
||||
except Exception:
|
||||
error_logger.exception("protocol.data_received")
|
||||
|
||||
|
||||
class Http3Protocol(HttpProtocolMixin, ConnectionProtocol): # type: ignore
|
||||
HTTP_CLASS = Http3
|
||||
__version__ = HTTP.VERSION_3
|
||||
|
||||
def __init__(self, *args, app: Sanic, **kwargs) -> None:
|
||||
self.app = app
|
||||
super().__init__(*args, **kwargs)
|
||||
self._setup()
|
||||
self._connection: H3Connection | None = None
|
||||
|
||||
def quic_event_received(self, event: QuicEvent) -> None:
|
||||
logger.debug(
|
||||
f"{Colors.BLUE}[quic_event_received]: "
|
||||
f"{Colors.PURPLE}{event}{Colors.END}",
|
||||
extra={"verbosity": 2},
|
||||
)
|
||||
if isinstance(event, ProtocolNegotiated):
|
||||
self._setup_connection(transmit=self.transmit)
|
||||
if event.alpn_protocol in H3_ALPN:
|
||||
self._connection = H3Connection(
|
||||
self._quic, enable_webtransport=True
|
||||
)
|
||||
elif isinstance(event, DatagramFrameReceived):
|
||||
if event.data == b"quack":
|
||||
self._quic.send_datagram_frame(b"quack-ack")
|
||||
|
||||
# pass event to the HTTP layer
|
||||
if self._connection is not None:
|
||||
for http_event in self._connection.handle_event(event):
|
||||
self._http.http_event_received(http_event)
|
||||
|
||||
@property
|
||||
def connection(self) -> H3Connection | None:
|
||||
return self._connection
|
||||
@@ -0,0 +1,230 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import cast
|
||||
|
||||
|
||||
try: # websockets >= 11.0
|
||||
from websockets.protocol import State # type: ignore
|
||||
from websockets.server import ServerProtocol # type: ignore
|
||||
except ImportError: # websockets < 11.0
|
||||
from websockets.connection import State
|
||||
from websockets.server import ServerConnection as ServerProtocol
|
||||
|
||||
from websockets import http11
|
||||
from websockets.datastructures import Headers as WSHeaders
|
||||
from websockets.typing import Subprotocol
|
||||
|
||||
from sanic.exceptions import SanicException
|
||||
from sanic.log import access_logger, websockets_logger
|
||||
from sanic.request import Request
|
||||
from sanic.server import HttpProtocol
|
||||
|
||||
from ..websockets.impl import WebsocketImplProtocol
|
||||
|
||||
|
||||
OPEN = State.OPEN
|
||||
CLOSING = State.CLOSING
|
||||
CLOSED = State.CLOSED
|
||||
|
||||
|
||||
class WebSocketProtocol(HttpProtocol):
|
||||
__slots__ = (
|
||||
"websocket",
|
||||
"websocket_timeout",
|
||||
"websocket_max_size",
|
||||
"websocket_ping_interval",
|
||||
"websocket_ping_timeout",
|
||||
"websocket_url",
|
||||
"websocket_peer",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
websocket_timeout: float = 10.0,
|
||||
websocket_max_size: int | None = None,
|
||||
websocket_ping_interval: float | None = 20.0,
|
||||
websocket_ping_timeout: float | None = 20.0,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.websocket: WebsocketImplProtocol | None = None
|
||||
self.websocket_timeout = websocket_timeout
|
||||
self.websocket_max_size = websocket_max_size
|
||||
self.websocket_ping_interval = websocket_ping_interval
|
||||
self.websocket_ping_timeout = websocket_ping_timeout
|
||||
self.websocket_url: str | None = None
|
||||
self.websocket_peer: str | None = None
|
||||
|
||||
def connection_lost(self, exc):
|
||||
if self.websocket is not None:
|
||||
self.websocket.connection_lost(exc)
|
||||
super().connection_lost(exc)
|
||||
self.log_websocket("CLOSE")
|
||||
self.websocket_url = None
|
||||
self.websocket_peer = None
|
||||
|
||||
def data_received(self, data):
|
||||
if self.websocket is not None:
|
||||
self.websocket.data_received(data)
|
||||
else:
|
||||
# Pass it to HttpProtocol handler first
|
||||
# That will (hopefully) upgrade it to a websocket.
|
||||
super().data_received(data)
|
||||
|
||||
def eof_received(self) -> bool | None:
|
||||
if self.websocket is not None:
|
||||
return self.websocket.eof_received()
|
||||
else:
|
||||
return False
|
||||
|
||||
def close(self, timeout: float | None = None):
|
||||
# Called by HttpProtocol at the end of connection_task
|
||||
# If we've upgraded to websocket, we do our own closing
|
||||
if self.websocket is not None:
|
||||
# Note, we don't want to use websocket.close()
|
||||
# That is used for user's application code to send a
|
||||
# websocket close packet. This is different.
|
||||
self.websocket.end_connection(1001)
|
||||
else:
|
||||
super().close()
|
||||
|
||||
def close_if_idle(self):
|
||||
# Called by Sanic Server when shutting down
|
||||
# If we've upgraded to websocket, shut it down
|
||||
if self.websocket is not None:
|
||||
if self.websocket.ws_proto.state in (CLOSING, CLOSED):
|
||||
return True
|
||||
elif self.websocket.loop is not None:
|
||||
self.websocket.loop.create_task(self.websocket.close(1001))
|
||||
else:
|
||||
self.websocket.end_connection(1001)
|
||||
else:
|
||||
return super().close_if_idle()
|
||||
|
||||
@staticmethod
|
||||
def sanic_request_to_ws_request(request: Request):
|
||||
return http11.Request(
|
||||
path=request.path,
|
||||
headers=WSHeaders(request.headers),
|
||||
)
|
||||
|
||||
async def websocket_handshake(
|
||||
self, request, subprotocols: Sequence[str] | None = None
|
||||
):
|
||||
# let the websockets package do the handshake with the client
|
||||
try:
|
||||
if subprotocols is not None:
|
||||
# subprotocols can be a set or frozenset,
|
||||
# but ServerProtocol needs a list
|
||||
subprotocols = cast(
|
||||
Sequence[Subprotocol] | None,
|
||||
list(
|
||||
[
|
||||
Subprotocol(subprotocol)
|
||||
for subprotocol in subprotocols
|
||||
]
|
||||
),
|
||||
)
|
||||
ws_proto = ServerProtocol(
|
||||
max_size=self.websocket_max_size,
|
||||
subprotocols=subprotocols,
|
||||
state=OPEN,
|
||||
logger=websockets_logger,
|
||||
)
|
||||
resp = ws_proto.accept(self.sanic_request_to_ws_request(request))
|
||||
except Exception:
|
||||
msg = (
|
||||
"Failed to open a WebSocket connection.\n"
|
||||
"See server log for more information.\n"
|
||||
)
|
||||
raise SanicException(msg, status_code=500)
|
||||
if 100 <= resp.status_code <= 299:
|
||||
first_line = (
|
||||
f"HTTP/1.1 {resp.status_code} {resp.reason_phrase}\r\n"
|
||||
).encode()
|
||||
rbody = bytearray(first_line)
|
||||
rbody += (
|
||||
"".join([f"{k}: {v}\r\n" for k, v in resp.headers.items()])
|
||||
).encode()
|
||||
rbody += b"\r\n"
|
||||
if resp.body:
|
||||
rbody += resp.body
|
||||
rbody += b"\r\n\r\n"
|
||||
await super().send(rbody)
|
||||
else:
|
||||
raise SanicException(resp.body, resp.status_code)
|
||||
self.websocket = WebsocketImplProtocol(
|
||||
ws_proto,
|
||||
ping_interval=self.websocket_ping_interval,
|
||||
ping_timeout=self.websocket_ping_timeout,
|
||||
close_timeout=self.websocket_timeout,
|
||||
)
|
||||
loop = (
|
||||
request.transport.loop
|
||||
if hasattr(request, "transport")
|
||||
and hasattr(request.transport, "loop")
|
||||
else None
|
||||
)
|
||||
await self.websocket.connection_made(self, loop=loop)
|
||||
self.websocket_url = self._http.request.url
|
||||
self.websocket_peer = f"{id(self):X}"[-5:-1] + "unx"
|
||||
if ip := self._http.request.client_ip:
|
||||
self.websocket_peer = f"{ip}:{self._http.request.port}"
|
||||
self.log_websocket("OPEN")
|
||||
return self.websocket
|
||||
|
||||
def log_websocket(self, message):
|
||||
if not self.access_log or not self.websocket_url:
|
||||
return
|
||||
status = ""
|
||||
close = ""
|
||||
try:
|
||||
# Can we get some useful statistics?
|
||||
p = self.websocket.ws_proto
|
||||
state = p.state
|
||||
if state == CLOSED:
|
||||
codes = {
|
||||
1000: "NORMAL",
|
||||
1001: "GOING AWAY",
|
||||
1005: "NO STATUS",
|
||||
1006: "ABNORMAL",
|
||||
1011: "SERVER ERR",
|
||||
}
|
||||
if p.close_code == 1006:
|
||||
message = "CLOSE_ABN"
|
||||
scode = rcode = 1006 # Abnormal closure (disconnection)
|
||||
sdesc = rdesc = ""
|
||||
if p.close_sent:
|
||||
scode = p.close_sent.code
|
||||
sdesc = p.close_sent.reason
|
||||
if p.close_rcvd:
|
||||
rcode = p.close_rcvd.code
|
||||
rdesc = p.close_rcvd.reason
|
||||
# Use repr() to escape any control characters
|
||||
sdesc = repr(sdesc[:256]) if sdesc else codes.get(scode, "")
|
||||
rdesc = repr(rdesc[:256]) if rdesc else codes.get(rcode, "")
|
||||
if p.close_rcvd_then_sent or scode == 1006:
|
||||
status = rcode
|
||||
close = (
|
||||
f"{rdesc} from client"
|
||||
if scode in (rcode, 1006)
|
||||
else f"{rdesc} ▼▲ {scode} {sdesc}"
|
||||
)
|
||||
else:
|
||||
status = scode
|
||||
close = (
|
||||
f"{sdesc} from server"
|
||||
if rcode in (scode, 1006)
|
||||
else f"{sdesc} ▲▼ {rcode} {rdesc}"
|
||||
)
|
||||
|
||||
except AttributeError:
|
||||
...
|
||||
extra = {
|
||||
"status": status,
|
||||
"byte": close,
|
||||
"host": self.websocket_peer,
|
||||
"request": f" 🔌 {self.websocket_url}",
|
||||
"duration": "",
|
||||
}
|
||||
access_logger.info(message, extra=extra)
|
||||
@@ -0,0 +1,404 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from ssl import SSLContext
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sanic.config import Config
|
||||
from sanic.exceptions import ServerError
|
||||
from sanic.http.constants import HTTP
|
||||
from sanic.http.tls import get_ssl_context
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sanic.app import Sanic
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
|
||||
from functools import partial
|
||||
from signal import SIG_IGN, SIGINT, SIGTERM
|
||||
from signal import signal as signal_func
|
||||
|
||||
from sanic.application.ext import setup_ext
|
||||
from sanic.compat import OS_IS_WINDOWS, ctrlc_workaround_for_windows
|
||||
from sanic.http.http3 import SessionTicketStore, get_config
|
||||
from sanic.log import error_logger, server_logger
|
||||
from sanic.logging.setup import setup_logging
|
||||
from sanic.models.server_types import Signal
|
||||
from sanic.server.async_server import AsyncioServer
|
||||
from sanic.server.protocols.http_protocol import Http3Protocol, HttpProtocol
|
||||
from sanic.server.socket import bind_unix_socket, remove_unix_socket
|
||||
|
||||
|
||||
try:
|
||||
from aioquic.asyncio import serve as quic_serve
|
||||
|
||||
HTTP3_AVAILABLE = True
|
||||
except ModuleNotFoundError: # no cov
|
||||
HTTP3_AVAILABLE = False
|
||||
|
||||
|
||||
def serve(
|
||||
host,
|
||||
port,
|
||||
app: Sanic,
|
||||
ssl: SSLContext | None = None,
|
||||
sock: socket.socket | None = None,
|
||||
unix: str | None = None,
|
||||
reuse_port: bool = False,
|
||||
loop=None,
|
||||
protocol: type[asyncio.Protocol] = HttpProtocol,
|
||||
backlog: int = 100,
|
||||
register_sys_signals: bool = True,
|
||||
run_multiple: bool = False,
|
||||
run_async: bool = False,
|
||||
connections=None,
|
||||
signal=Signal(),
|
||||
state=None,
|
||||
asyncio_server_kwargs=None,
|
||||
version=HTTP.VERSION_1,
|
||||
):
|
||||
"""Start asynchronous HTTP Server on an individual process.
|
||||
|
||||
:param host: Address to host on
|
||||
:param port: Port to host on
|
||||
:param before_start: function to be executed before the server starts
|
||||
listening. Takes arguments `app` instance and `loop`
|
||||
:param after_start: function to be executed after the server starts
|
||||
listening. Takes arguments `app` instance and `loop`
|
||||
:param before_stop: function to be executed when a stop signal is
|
||||
received before it is respected. Takes arguments
|
||||
`app` instance and `loop`
|
||||
:param after_stop: function to be executed when a stop signal is
|
||||
received after it is respected. Takes arguments
|
||||
`app` instance and `loop`
|
||||
:param ssl: SSLContext
|
||||
:param sock: Socket for the server to accept connections from
|
||||
:param unix: Unix socket to listen on instead of TCP port
|
||||
:param reuse_port: `True` for multiple workers
|
||||
:param loop: asyncio compatible event loop
|
||||
:param run_async: bool: Do not create a new event loop for the server,
|
||||
and return an AsyncServer object rather than running it
|
||||
:param asyncio_server_kwargs: key-value args for asyncio/uvloop
|
||||
create_server method
|
||||
:return: Nothing
|
||||
|
||||
Args:
|
||||
host (str): Address to host on
|
||||
port (int): Port to host on
|
||||
app (Sanic): Sanic app instance
|
||||
ssl (Optional[SSLContext], optional): SSLContext. Defaults to `None`.
|
||||
sock (Optional[socket.socket], optional): Socket for the server to
|
||||
accept connections from. Defaults to `None`.
|
||||
unix (Optional[str], optional): Unix socket to listen on instead of
|
||||
TCP port. Defaults to `None`.
|
||||
reuse_port (bool, optional): `True` for multiple workers. Defaults
|
||||
to `False`.
|
||||
loop: asyncio compatible event loop. Defaults
|
||||
to `None`.
|
||||
protocol (Type[asyncio.Protocol], optional): Protocol to use. Defaults
|
||||
to `HttpProtocol`.
|
||||
backlog (int, optional): The maximum number of queued connections
|
||||
passed to socket.listen(). Defaults to `100`.
|
||||
register_sys_signals (bool, optional): Register SIGINT and SIGTERM.
|
||||
Defaults to `True`.
|
||||
run_multiple (bool, optional): Run multiple workers. Defaults
|
||||
to `False`.
|
||||
run_async (bool, optional): Return an AsyncServer object.
|
||||
Defaults to `False`.
|
||||
connections: Connections. Defaults to `None`.
|
||||
signal (Signal, optional): Signal. Defaults to `Signal()`.
|
||||
state: State. Defaults to `None`.
|
||||
asyncio_server_kwargs (Optional[Dict[str, Union[int, float]]], optional):
|
||||
key-value args for asyncio/uvloop create_server method. Defaults
|
||||
to `None`.
|
||||
version (str, optional): HTTP version. Defaults to `HTTP.VERSION_1`.
|
||||
|
||||
Raises:
|
||||
ServerError: Cannot run HTTP/3 server without aioquic installed.
|
||||
|
||||
Returns:
|
||||
AsyncioServer: AsyncioServer object if `run_async` is `True`.
|
||||
""" # noqa: E501
|
||||
if not run_async and not loop:
|
||||
# create new event_loop after fork
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
setup_logging(app.debug, app.config.NO_COLOR, app.config.LOG_EXTRA)
|
||||
|
||||
if app.debug:
|
||||
loop.set_debug(app.debug)
|
||||
|
||||
app.asgi = False
|
||||
|
||||
if version is HTTP.VERSION_3:
|
||||
return _serve_http_3(host, port, app, loop, ssl)
|
||||
return _serve_http_1(
|
||||
host,
|
||||
port,
|
||||
app,
|
||||
ssl,
|
||||
sock,
|
||||
unix,
|
||||
reuse_port,
|
||||
loop,
|
||||
protocol,
|
||||
backlog,
|
||||
register_sys_signals,
|
||||
run_multiple,
|
||||
run_async,
|
||||
connections,
|
||||
signal,
|
||||
state,
|
||||
asyncio_server_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _setup_system_signals(
|
||||
app: Sanic,
|
||||
run_multiple: bool,
|
||||
register_sys_signals: bool,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
) -> None: # no cov
|
||||
signal_func(SIGINT, SIG_IGN)
|
||||
signal_func(SIGTERM, SIG_IGN)
|
||||
os.environ["SANIC_WORKER_PROCESS"] = "true"
|
||||
# Register signals for graceful termination
|
||||
if register_sys_signals:
|
||||
if OS_IS_WINDOWS:
|
||||
ctrlc_workaround_for_windows(app)
|
||||
else:
|
||||
for _signal in [SIGINT, SIGTERM]:
|
||||
loop.add_signal_handler(
|
||||
_signal, partial(app.stop, terminate=False)
|
||||
)
|
||||
|
||||
|
||||
def _run_shutdown_coro(loop, coro):
|
||||
"""Run a shutdown coroutine, handling the case where the loop was stopped.
|
||||
|
||||
When loop.stop() is called during run_forever(), asyncio sets an internal
|
||||
flag that causes the first run_until_complete() to fail. For standard
|
||||
asyncio, we clear the _stopping flag. For uvloop (which doesn't expose
|
||||
this flag), the first attempt may fail but subsequent attempts succeed.
|
||||
"""
|
||||
# Clear asyncio's stopped state if accessible
|
||||
if hasattr(loop, "_stopping"):
|
||||
loop._stopping = False
|
||||
|
||||
try:
|
||||
loop.run_until_complete(coro())
|
||||
except (RuntimeError, KeyboardInterrupt):
|
||||
# RuntimeError: loop was stopped (uvloop behavior)
|
||||
# KeyboardInterrupt: signal arrived during select (asyncio behavior)
|
||||
# Try once more - this handles uvloop's behavior where the first
|
||||
# run_until_complete after stop() fails but subsequent calls succeed.
|
||||
if hasattr(loop, "_stopping"):
|
||||
loop._stopping = False
|
||||
try:
|
||||
loop.run_until_complete(coro())
|
||||
except (RuntimeError, KeyboardInterrupt):
|
||||
# If it still fails, the loop is truly unusable
|
||||
pass
|
||||
|
||||
|
||||
def _run_server_forever(loop, before_stop, after_stop, cleanup, unix, pid):
|
||||
try:
|
||||
server_logger.info("Worker ready [%s]", pid)
|
||||
loop.run_forever()
|
||||
finally:
|
||||
server_logger.info("Stopping worker [%s]", pid)
|
||||
|
||||
for _signal in [SIGINT, SIGTERM]:
|
||||
try:
|
||||
loop.remove_signal_handler(_signal)
|
||||
except (NotImplementedError, OSError):
|
||||
pass
|
||||
|
||||
_run_shutdown_coro(loop, before_stop)
|
||||
|
||||
if cleanup:
|
||||
cleanup()
|
||||
|
||||
_run_shutdown_coro(loop, after_stop)
|
||||
|
||||
remove_unix_socket(unix)
|
||||
loop.close()
|
||||
server_logger.info("Worker complete [%s]", pid)
|
||||
|
||||
|
||||
def _serve_http_1(
|
||||
host,
|
||||
port,
|
||||
app,
|
||||
ssl,
|
||||
sock,
|
||||
unix,
|
||||
reuse_port,
|
||||
loop,
|
||||
protocol,
|
||||
backlog,
|
||||
register_sys_signals,
|
||||
run_multiple,
|
||||
run_async,
|
||||
connections,
|
||||
signal,
|
||||
state,
|
||||
asyncio_server_kwargs,
|
||||
):
|
||||
connections = connections if connections is not None else set()
|
||||
protocol_kwargs = _build_protocol_kwargs(protocol, app.config)
|
||||
server = partial(
|
||||
protocol,
|
||||
loop=loop,
|
||||
connections=connections,
|
||||
signal=signal,
|
||||
app=app,
|
||||
state=state,
|
||||
unix=unix,
|
||||
**protocol_kwargs,
|
||||
)
|
||||
asyncio_server_kwargs = (
|
||||
asyncio_server_kwargs if asyncio_server_kwargs else {}
|
||||
)
|
||||
if OS_IS_WINDOWS and sock:
|
||||
pid = os.getpid()
|
||||
sock = sock.share(pid)
|
||||
sock = socket.fromshare(sock)
|
||||
# UNIX sockets are always bound by us (to preserve semantics between modes)
|
||||
elif unix:
|
||||
sock = bind_unix_socket(unix, backlog=backlog)
|
||||
server_coroutine = loop.create_server(
|
||||
server,
|
||||
None if sock else host,
|
||||
None if sock else port,
|
||||
ssl=ssl,
|
||||
reuse_port=reuse_port,
|
||||
sock=sock,
|
||||
backlog=backlog,
|
||||
**asyncio_server_kwargs,
|
||||
)
|
||||
|
||||
setup_ext(app)
|
||||
if run_async:
|
||||
return AsyncioServer(
|
||||
app=app,
|
||||
loop=loop,
|
||||
serve_coro=server_coroutine,
|
||||
connections=connections,
|
||||
)
|
||||
|
||||
pid = os.getpid()
|
||||
server_logger.info("Starting worker [%s]", pid)
|
||||
loop.run_until_complete(app._startup())
|
||||
loop.run_until_complete(app._server_event("init", "before"))
|
||||
app.ack()
|
||||
|
||||
try:
|
||||
http_server = loop.run_until_complete(server_coroutine)
|
||||
except BaseException:
|
||||
error_logger.exception("Unable to start server", exc_info=True)
|
||||
return
|
||||
|
||||
def _cleanup():
|
||||
# Wait for event loop to finish and all connections to drain
|
||||
http_server.close()
|
||||
loop.run_until_complete(http_server.wait_closed())
|
||||
|
||||
# Complete all tasks on the loop
|
||||
signal.stopped = True
|
||||
for connection in connections:
|
||||
connection.close_if_idle()
|
||||
|
||||
# Gracefully shutdown timeout.
|
||||
# We should provide graceful_shutdown_timeout,
|
||||
# instead of letting connection hangs forever.
|
||||
# Let's roughly calcucate time.
|
||||
graceful = app.config.GRACEFUL_SHUTDOWN_TIMEOUT
|
||||
start_shutdown: float = 0
|
||||
while connections and (start_shutdown < graceful):
|
||||
loop.run_until_complete(asyncio.sleep(0.1))
|
||||
start_shutdown = start_shutdown + 0.1
|
||||
|
||||
app.shutdown_tasks(graceful - start_shutdown)
|
||||
|
||||
# Force close non-idle connection after waiting for
|
||||
# graceful_shutdown_timeout
|
||||
for conn in connections:
|
||||
if hasattr(conn, "websocket") and conn.websocket:
|
||||
conn.websocket.fail_connection(code=1001)
|
||||
else:
|
||||
conn.abort()
|
||||
|
||||
try:
|
||||
app.set_serving(False)
|
||||
except (BrokenPipeError, ConnectionResetError, EOFError):
|
||||
pass
|
||||
|
||||
_setup_system_signals(app, run_multiple, register_sys_signals, loop)
|
||||
loop.run_until_complete(app._server_event("init", "after"))
|
||||
app.set_serving(True)
|
||||
_run_server_forever(
|
||||
loop,
|
||||
partial(app._server_event, "shutdown", "before"),
|
||||
partial(app._server_event, "shutdown", "after"),
|
||||
_cleanup,
|
||||
unix,
|
||||
pid,
|
||||
)
|
||||
|
||||
|
||||
def _serve_http_3(
|
||||
host,
|
||||
port,
|
||||
app,
|
||||
loop,
|
||||
ssl,
|
||||
register_sys_signals: bool = True,
|
||||
run_multiple: bool = False,
|
||||
):
|
||||
if not HTTP3_AVAILABLE:
|
||||
raise ServerError(
|
||||
"Cannot run HTTP/3 server without aioquic installed. "
|
||||
)
|
||||
pid = os.getpid()
|
||||
server_logger.info("Starting worker [%s]", pid)
|
||||
protocol = partial(Http3Protocol, app=app)
|
||||
ticket_store = SessionTicketStore()
|
||||
ssl_context = get_ssl_context(app, ssl)
|
||||
config = get_config(app, ssl_context)
|
||||
coro = quic_serve(
|
||||
host,
|
||||
port,
|
||||
configuration=config,
|
||||
create_protocol=protocol,
|
||||
session_ticket_fetcher=ticket_store.pop,
|
||||
session_ticket_handler=ticket_store.add,
|
||||
)
|
||||
server = AsyncioServer(app, loop, coro, [])
|
||||
loop.run_until_complete(server.startup())
|
||||
loop.run_until_complete(server.before_start())
|
||||
app.ack()
|
||||
loop.run_until_complete(server)
|
||||
_setup_system_signals(app, run_multiple, register_sys_signals, loop)
|
||||
loop.run_until_complete(server.after_start())
|
||||
|
||||
# TODO: Create connection cleanup and graceful shutdown
|
||||
cleanup = None
|
||||
_run_server_forever(
|
||||
loop, server.before_stop, server.after_stop, cleanup, None, pid
|
||||
)
|
||||
|
||||
|
||||
def _build_protocol_kwargs(
|
||||
protocol: type[asyncio.Protocol], config: Config
|
||||
) -> dict[str, int | float]:
|
||||
if hasattr(protocol, "websocket_handshake"):
|
||||
return {
|
||||
"websocket_max_size": config.WEBSOCKET_MAX_SIZE,
|
||||
"websocket_ping_timeout": config.WEBSOCKET_PING_TIMEOUT,
|
||||
"websocket_ping_interval": config.WEBSOCKET_PING_INTERVAL,
|
||||
}
|
||||
return {}
|
||||
@@ -0,0 +1,121 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import secrets
|
||||
import socket
|
||||
import stat
|
||||
|
||||
from ipaddress import ip_address
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from sanic.http.constants import HTTP
|
||||
|
||||
|
||||
def bind_socket(host: str, port: int, *, backlog=100) -> socket.socket:
|
||||
"""Create TCP server socket.
|
||||
:param host: IPv4, IPv6 or hostname may be specified
|
||||
:param port: TCP port number
|
||||
:param backlog: Maximum number of connections to queue
|
||||
:return: socket.socket object
|
||||
"""
|
||||
location = (host, port)
|
||||
# socket.share, socket.fromshare
|
||||
try: # IP address: family must be specified for IPv6 at least
|
||||
ip = ip_address(host)
|
||||
host = str(ip)
|
||||
sock = socket.socket(
|
||||
socket.AF_INET6 if ip.version == 6 else socket.AF_INET
|
||||
)
|
||||
except ValueError: # Hostname, may become AF_INET or AF_INET6
|
||||
sock = socket.socket()
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
sock.bind(location)
|
||||
sock.listen(backlog)
|
||||
sock.set_inheritable(True)
|
||||
return sock
|
||||
|
||||
|
||||
def bind_unix_socket(
|
||||
path: Path | str, *, mode=0o666, backlog=100
|
||||
) -> socket.socket:
|
||||
"""Create unix socket.
|
||||
:param path: filesystem path
|
||||
:param backlog: Maximum number of connections to queue
|
||||
:return: socket.socket object
|
||||
"""
|
||||
|
||||
# Sanitise and pre-verify socket path
|
||||
path = Path(path)
|
||||
folder = path.parent
|
||||
if not folder.is_dir():
|
||||
raise FileNotFoundError(f"Socket folder does not exist: {folder}")
|
||||
try:
|
||||
if not stat.S_ISSOCK(path.lstat().st_mode):
|
||||
raise FileExistsError(f"Existing file is not a socket: {path}")
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
# Create new socket with a random temporary name
|
||||
tmp_path = path.with_name(f"{path.name}.{secrets.token_urlsafe()}")
|
||||
sock = socket.socket(socket.AF_UNIX)
|
||||
try:
|
||||
# Critical section begins (filename races)
|
||||
sock.bind(tmp_path.as_posix())
|
||||
try:
|
||||
tmp_path.chmod(mode)
|
||||
# Start listening before rename to avoid connection failures
|
||||
sock.listen(backlog)
|
||||
tmp_path.rename(path)
|
||||
except: # noqa: E722
|
||||
try:
|
||||
tmp_path.unlink()
|
||||
finally:
|
||||
raise
|
||||
except: # noqa: E722
|
||||
try:
|
||||
sock.close()
|
||||
finally:
|
||||
raise
|
||||
return sock
|
||||
|
||||
|
||||
def remove_unix_socket(path: Path | str | None) -> None:
|
||||
"""Remove dead unix socket during server exit."""
|
||||
if not path:
|
||||
return
|
||||
try:
|
||||
path = Path(path)
|
||||
if stat.S_ISSOCK(path.lstat().st_mode):
|
||||
# Is it actually dead (doesn't belong to a new server instance)?
|
||||
with socket.socket(socket.AF_UNIX) as testsock:
|
||||
try:
|
||||
testsock.connect(path.as_posix())
|
||||
except ConnectionRefusedError:
|
||||
path.unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
def configure_socket(
|
||||
server_settings: dict[str, Any],
|
||||
) -> socket.SocketType | None:
|
||||
# Create a listening socket or use the one in settings
|
||||
if server_settings.get("version") is HTTP.VERSION_3:
|
||||
return None
|
||||
sock = server_settings.get("sock")
|
||||
unix = server_settings["unix"]
|
||||
backlog = server_settings["backlog"]
|
||||
if unix:
|
||||
unix = Path(unix).absolute()
|
||||
sock = bind_unix_socket(unix, backlog=backlog)
|
||||
server_settings["unix"] = unix
|
||||
if sock is None:
|
||||
sock = bind_socket(
|
||||
server_settings["host"],
|
||||
server_settings["port"],
|
||||
backlog=backlog,
|
||||
)
|
||||
sock.set_inheritable(True)
|
||||
server_settings["sock"] = sock
|
||||
server_settings["host"] = None
|
||||
server_settings["port"] = None
|
||||
return sock
|
||||
@@ -0,0 +1,82 @@
|
||||
from collections.abc import Awaitable, MutableMapping
|
||||
from typing import Any, Callable
|
||||
|
||||
from sanic.exceptions import InvalidUsage
|
||||
|
||||
|
||||
ASGIMessage = MutableMapping[str, Any]
|
||||
|
||||
|
||||
class WebSocketConnection:
|
||||
"""
|
||||
This is for ASGI Connections.
|
||||
It provides an interface similar to WebsocketProtocol, but
|
||||
sends/receives over an ASGI connection.
|
||||
"""
|
||||
|
||||
# TODO
|
||||
# - Implement ping/pong
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
send: Callable[[ASGIMessage], Awaitable[None]],
|
||||
receive: Callable[[], Awaitable[ASGIMessage]],
|
||||
subprotocols: list[str] | None = None,
|
||||
) -> None:
|
||||
self._send = send
|
||||
self._receive = receive
|
||||
self._subprotocols = subprotocols or []
|
||||
|
||||
async def send(self, data: str | bytes, *args, **kwargs) -> None:
|
||||
message: dict[str, str | bytes] = {"type": "websocket.send"}
|
||||
|
||||
if isinstance(data, bytes):
|
||||
message.update({"bytes": data})
|
||||
else:
|
||||
message.update({"text": str(data)})
|
||||
|
||||
await self._send(message)
|
||||
|
||||
async def recv(self, *args, **kwargs) -> str | bytes | None:
|
||||
message = await self._receive()
|
||||
|
||||
if message["type"] == "websocket.receive":
|
||||
try:
|
||||
return message["text"]
|
||||
except KeyError:
|
||||
try:
|
||||
return message["bytes"]
|
||||
except KeyError:
|
||||
raise InvalidUsage("Bad ASGI message received")
|
||||
elif message["type"] == "websocket.disconnect":
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
receive = recv
|
||||
|
||||
async def accept(self, subprotocols: list[str] | None = None) -> None:
|
||||
subprotocol = None
|
||||
if subprotocols:
|
||||
for subp in subprotocols:
|
||||
if subp in self.subprotocols:
|
||||
subprotocol = subp
|
||||
break
|
||||
|
||||
await self._send(
|
||||
{
|
||||
"type": "websocket.accept",
|
||||
"subprotocol": subprotocol,
|
||||
}
|
||||
)
|
||||
|
||||
async def close(self, code: int = 1000, reason: str = "") -> None:
|
||||
pass
|
||||
|
||||
@property
|
||||
def subprotocols(self):
|
||||
return self._subprotocols
|
||||
|
||||
@subprotocols.setter
|
||||
def subprotocols(self, subprotocols: list[str] | None = None):
|
||||
self._subprotocols = subprotocols or []
|
||||
@@ -0,0 +1,293 @@
|
||||
import asyncio
|
||||
import codecs
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from websockets.frames import Frame, Opcode
|
||||
from websockets.typing import Data
|
||||
|
||||
from sanic.exceptions import ServerError
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .impl import WebsocketImplProtocol
|
||||
|
||||
UTF8Decoder = codecs.getincrementaldecoder("utf-8")
|
||||
|
||||
|
||||
class WebsocketFrameAssembler:
|
||||
"""
|
||||
Assemble a message from frames.
|
||||
Code borrowed from aaugustin/websockets project:
|
||||
https://github.com/aaugustin/websockets/blob/6eb98dd8fa5b2c896b9f6be7e8d117708da82a39/src/websockets/sync/messages.py
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"protocol",
|
||||
"read_mutex",
|
||||
"write_mutex",
|
||||
"message_complete",
|
||||
"message_fetched",
|
||||
"get_in_progress",
|
||||
"decoder",
|
||||
"completed_queue",
|
||||
"chunks",
|
||||
"chunks_queue",
|
||||
"paused",
|
||||
"get_id",
|
||||
"put_id",
|
||||
)
|
||||
if TYPE_CHECKING:
|
||||
protocol: "WebsocketImplProtocol"
|
||||
read_mutex: asyncio.Lock
|
||||
write_mutex: asyncio.Lock
|
||||
message_complete: asyncio.Event
|
||||
message_fetched: asyncio.Event
|
||||
completed_queue: asyncio.Queue
|
||||
get_in_progress: bool
|
||||
decoder: codecs.IncrementalDecoder | None
|
||||
# For streaming chunks rather than messages:
|
||||
chunks: list[Data]
|
||||
chunks_queue: asyncio.Queue[Data | None] | None
|
||||
paused: bool
|
||||
|
||||
def __init__(self, protocol) -> None:
|
||||
self.protocol = protocol
|
||||
|
||||
self.read_mutex = asyncio.Lock()
|
||||
self.write_mutex = asyncio.Lock()
|
||||
|
||||
self.completed_queue = asyncio.Queue(maxsize=1) # type: asyncio.Queue[Data]
|
||||
|
||||
# put() sets this event to tell get() that a message can be fetched.
|
||||
self.message_complete = asyncio.Event()
|
||||
# get() sets this event to let put()
|
||||
self.message_fetched = asyncio.Event()
|
||||
|
||||
# This flag prevents concurrent calls to get() by user code.
|
||||
self.get_in_progress = False
|
||||
|
||||
# Decoder for text frames, None for binary frames.
|
||||
self.decoder = None
|
||||
|
||||
# Buffer data from frames belonging to the same message.
|
||||
self.chunks = []
|
||||
|
||||
# When switching from "buffering" to "streaming", we use a thread-safe
|
||||
# queue for transferring frames from the writing thread (library code)
|
||||
# to the reading thread (user code). We're buffering when chunks_queue
|
||||
# is None and streaming when it's a Queue. None is a sentinel
|
||||
# value marking the end of the stream, superseding message_complete.
|
||||
|
||||
# Stream data from frames belonging to the same message.
|
||||
self.chunks_queue = None
|
||||
|
||||
# Flag to indicate we've paused the protocol
|
||||
self.paused = False
|
||||
|
||||
async def get(self, timeout: float | None = None) -> Data | None:
|
||||
"""
|
||||
Read the next message.
|
||||
:meth:`get` returns a single :class:`str` or :class:`bytes`.
|
||||
If the :message was fragmented, :meth:`get` waits until the last frame
|
||||
is received, then it reassembles the message.
|
||||
If ``timeout`` is set and elapses before a complete message is
|
||||
received, :meth:`get` returns ``None``.
|
||||
"""
|
||||
completed: bool
|
||||
async with self.read_mutex:
|
||||
if timeout is not None and timeout <= 0:
|
||||
if not self.message_complete.is_set():
|
||||
return None
|
||||
if self.get_in_progress:
|
||||
# This should be guarded against with the read_mutex,
|
||||
# exception is only here as a failsafe
|
||||
raise ServerError(
|
||||
"Called get() on Websocket frame assembler "
|
||||
"while asynchronous get is already in progress."
|
||||
)
|
||||
self.get_in_progress = True
|
||||
|
||||
# If the message_complete event isn't set yet, release the lock to
|
||||
# allow put() to run and eventually set it.
|
||||
# Locking with get_in_progress ensures only one task can get here.
|
||||
if timeout is None:
|
||||
completed = await self.message_complete.wait()
|
||||
elif timeout <= 0:
|
||||
completed = self.message_complete.is_set()
|
||||
else:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self.message_complete.wait(), timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
...
|
||||
finally:
|
||||
completed = self.message_complete.is_set()
|
||||
|
||||
# Unpause the transport, if its paused
|
||||
if self.paused:
|
||||
self.protocol.resume_frames()
|
||||
self.paused = False
|
||||
if not self.get_in_progress: # no cov
|
||||
# This should be guarded against with the read_mutex,
|
||||
# exception is here as a failsafe
|
||||
raise ServerError(
|
||||
"State of Websocket frame assembler was modified while an "
|
||||
"asynchronous get was in progress."
|
||||
)
|
||||
self.get_in_progress = False
|
||||
|
||||
# Waiting for a complete message timed out.
|
||||
if not completed:
|
||||
return None
|
||||
if not self.message_complete.is_set():
|
||||
return None
|
||||
|
||||
self.message_complete.clear()
|
||||
|
||||
joiner: Data = b"" if self.decoder is None else ""
|
||||
# mypy cannot figure out that chunks have the proper type.
|
||||
message: Data = joiner.join(self.chunks) # type: ignore
|
||||
if self.message_fetched.is_set():
|
||||
# This should be guarded against with the read_mutex,
|
||||
# and get_in_progress check, this exception is here
|
||||
# as a failsafe
|
||||
raise ServerError(
|
||||
"Websocket get() found a message when "
|
||||
"state was already fetched."
|
||||
)
|
||||
self.message_fetched.set()
|
||||
self.chunks = []
|
||||
# this should already be None, but set it here for safety
|
||||
self.chunks_queue = None
|
||||
return message
|
||||
|
||||
async def get_iter(self) -> AsyncIterator[Data]:
|
||||
"""
|
||||
Stream the next message.
|
||||
Iterating the return value of :meth:`get_iter` yields a :class:`str`
|
||||
or :class:`bytes` for each frame in the message.
|
||||
"""
|
||||
async with self.read_mutex:
|
||||
if self.get_in_progress:
|
||||
# This should be guarded against with the read_mutex,
|
||||
# exception is only here as a failsafe
|
||||
raise ServerError(
|
||||
"Called get_iter on Websocket frame assembler "
|
||||
"while asynchronous get is already in progress."
|
||||
)
|
||||
self.get_in_progress = True
|
||||
|
||||
chunks = self.chunks
|
||||
self.chunks = []
|
||||
self.chunks_queue = asyncio.Queue()
|
||||
|
||||
# Sending None in chunk_queue supersedes setting message_complete
|
||||
# when switching to "streaming". If message is already complete
|
||||
# when the switch happens, put() didn't send None, so we have to.
|
||||
if self.message_complete.is_set():
|
||||
await self.chunks_queue.put(None)
|
||||
|
||||
# Locking with get_in_progress ensures only one task can get here
|
||||
for c in chunks:
|
||||
yield c
|
||||
while True:
|
||||
chunk = await self.chunks_queue.get()
|
||||
if chunk is None:
|
||||
break
|
||||
yield chunk
|
||||
|
||||
# Unpause the transport, if its paused
|
||||
if self.paused:
|
||||
self.protocol.resume_frames()
|
||||
self.paused = False
|
||||
if not self.get_in_progress: # no cov
|
||||
# This should be guarded against with the read_mutex,
|
||||
# exception is here as a failsafe
|
||||
raise ServerError(
|
||||
"State of Websocket frame assembler was modified while an "
|
||||
"asynchronous get was in progress."
|
||||
)
|
||||
self.get_in_progress = False
|
||||
if not self.message_complete.is_set(): # no cov
|
||||
# This should be guarded against with the read_mutex,
|
||||
# exception is here as a failsafe
|
||||
raise ServerError(
|
||||
"Websocket frame assembler chunks queue ended before "
|
||||
"message was complete."
|
||||
)
|
||||
self.message_complete.clear()
|
||||
if self.message_fetched.is_set(): # no cov
|
||||
# This should be guarded against with the read_mutex,
|
||||
# and get_in_progress check, this exception is
|
||||
# here as a failsafe
|
||||
raise ServerError(
|
||||
"Websocket get_iter() found a message when state was "
|
||||
"already fetched."
|
||||
)
|
||||
|
||||
self.message_fetched.set()
|
||||
# this should already be empty, but set it here for safety
|
||||
self.chunks = []
|
||||
self.chunks_queue = None
|
||||
|
||||
async def put(self, frame: Frame) -> None:
|
||||
"""
|
||||
Add ``frame`` to the next message.
|
||||
When ``frame`` is the final frame in a message, :meth:`put` waits
|
||||
until the message is fetched, either by calling :meth:`get` or by
|
||||
iterating the return value of :meth:`get_iter`.
|
||||
:meth:`put` assumes that the stream of frames respects the protocol.
|
||||
If it doesn't, the behavior is undefined.
|
||||
"""
|
||||
|
||||
async with self.write_mutex:
|
||||
if frame.opcode is Opcode.TEXT:
|
||||
self.decoder = UTF8Decoder(errors="strict")
|
||||
elif frame.opcode is Opcode.BINARY:
|
||||
self.decoder = None
|
||||
elif frame.opcode is Opcode.CONT:
|
||||
pass
|
||||
else:
|
||||
# Ignore control frames.
|
||||
return
|
||||
data: Data
|
||||
if self.decoder is not None:
|
||||
data = self.decoder.decode(frame.data, frame.fin)
|
||||
else:
|
||||
data = frame.data
|
||||
if self.chunks_queue is None:
|
||||
self.chunks.append(data)
|
||||
else:
|
||||
await self.chunks_queue.put(data)
|
||||
|
||||
if not frame.fin:
|
||||
return
|
||||
if not self.get_in_progress:
|
||||
# nobody is waiting for this frame, so try to pause subsequent
|
||||
# frames at the protocol level
|
||||
self.paused = self.protocol.pause_frames()
|
||||
# Message is complete. Wait until it's fetched to return.
|
||||
|
||||
if self.chunks_queue is not None:
|
||||
await self.chunks_queue.put(None)
|
||||
if self.message_complete.is_set():
|
||||
# This should be guarded against with the write_mutex
|
||||
raise ServerError(
|
||||
"Websocket put() got a new message when a message was "
|
||||
"already in its chamber."
|
||||
)
|
||||
self.message_complete.set() # Signal to get() it can serve the
|
||||
if self.message_fetched.is_set():
|
||||
# This should be guarded against with the write_mutex
|
||||
raise ServerError(
|
||||
"Websocket put() got a new message when the previous "
|
||||
"message was not yet fetched."
|
||||
)
|
||||
|
||||
# Allow get() to run and eventually set the event.
|
||||
await self.message_fetched.wait()
|
||||
self.message_fetched.clear()
|
||||
self.decoder = None
|
||||
@@ -0,0 +1,878 @@
|
||||
import asyncio
|
||||
import secrets
|
||||
|
||||
from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
|
||||
|
||||
from websockets.exceptions import (
|
||||
ConnectionClosed,
|
||||
ConnectionClosedError,
|
||||
ConnectionClosedOK,
|
||||
)
|
||||
from websockets.frames import Frame, Opcode
|
||||
|
||||
|
||||
try: # websockets >= 11.0
|
||||
from websockets.protocol import Event, State # type: ignore
|
||||
from websockets.server import ServerProtocol # type: ignore
|
||||
except ImportError: # websockets < 11.0
|
||||
from websockets.connection import Event, State # type: ignore
|
||||
from websockets.server import ServerConnection as ServerProtocol
|
||||
|
||||
from websockets.typing import Data
|
||||
|
||||
from sanic.log import websockets_logger
|
||||
from sanic.server.protocols.base_protocol import SanicProtocol
|
||||
|
||||
from ...exceptions import ServerError, WebsocketClosed
|
||||
from .frame import WebsocketFrameAssembler
|
||||
|
||||
|
||||
OPEN = State.OPEN
|
||||
CLOSING = State.CLOSING
|
||||
CLOSED = State.CLOSED
|
||||
|
||||
|
||||
class WebsocketImplProtocol:
|
||||
ws_proto: ServerProtocol
|
||||
io_proto: SanicProtocol | None
|
||||
loop: asyncio.AbstractEventLoop | None
|
||||
max_queue: int
|
||||
close_timeout: float
|
||||
ping_interval: float | None
|
||||
ping_timeout: float | None
|
||||
assembler: WebsocketFrameAssembler
|
||||
# dict[bytes, asyncio.Future[None]]
|
||||
pings: dict[bytes, asyncio.Future]
|
||||
conn_mutex: asyncio.Lock
|
||||
recv_lock: asyncio.Lock
|
||||
recv_cancel: asyncio.Future | None
|
||||
process_event_mutex: asyncio.Lock
|
||||
can_pause: bool
|
||||
# asyncio.Future[None] | None
|
||||
data_finished_fut: asyncio.Future | None
|
||||
# asyncio.Future[None] | None
|
||||
pause_frame_fut: asyncio.Future | None
|
||||
# asyncio.Future[None] | None
|
||||
connection_lost_waiter: asyncio.Future | None
|
||||
keepalive_ping_task: asyncio.Task | None
|
||||
auto_closer_task: asyncio.Task | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ws_proto,
|
||||
max_queue=None,
|
||||
ping_interval: float | None = 20,
|
||||
ping_timeout: float | None = 20,
|
||||
close_timeout: float = 10,
|
||||
loop=None,
|
||||
):
|
||||
self.ws_proto = ws_proto
|
||||
self.io_proto = None
|
||||
self.loop = None
|
||||
self.max_queue = max_queue
|
||||
self.close_timeout = close_timeout
|
||||
self.ping_interval = ping_interval
|
||||
self.ping_timeout = ping_timeout
|
||||
self.assembler = WebsocketFrameAssembler(self)
|
||||
self.pings = {}
|
||||
self.conn_mutex = asyncio.Lock()
|
||||
self.recv_lock = asyncio.Lock()
|
||||
self.recv_cancel = None
|
||||
self.process_event_mutex = asyncio.Lock()
|
||||
self.data_finished_fut = None
|
||||
self.can_pause = True
|
||||
self.pause_frame_fut = None
|
||||
self.keepalive_ping_task = None
|
||||
self.auto_closer_task = None
|
||||
self.connection_lost_waiter = None
|
||||
|
||||
@property
|
||||
def subprotocol(self):
|
||||
return self.ws_proto.subprotocol
|
||||
|
||||
def pause_frames(self):
|
||||
if not self.can_pause:
|
||||
return False
|
||||
if self.pause_frame_fut:
|
||||
websockets_logger.debug("Websocket connection already paused.")
|
||||
return False
|
||||
if (not self.loop) or (not self.io_proto):
|
||||
return False
|
||||
if self.io_proto.transport:
|
||||
self.io_proto.transport.pause_reading()
|
||||
self.pause_frame_fut = self.loop.create_future()
|
||||
websockets_logger.debug("Websocket connection paused.")
|
||||
return True
|
||||
|
||||
def resume_frames(self):
|
||||
if not self.pause_frame_fut:
|
||||
websockets_logger.debug("Websocket connection not paused.")
|
||||
return False
|
||||
if (not self.loop) or (not self.io_proto):
|
||||
websockets_logger.debug(
|
||||
"Websocket attempting to resume reading frames, "
|
||||
"but connection is gone."
|
||||
)
|
||||
return False
|
||||
if self.io_proto.transport:
|
||||
self.io_proto.transport.resume_reading()
|
||||
self.pause_frame_fut.set_result(None)
|
||||
self.pause_frame_fut = None
|
||||
websockets_logger.debug("Websocket connection unpaused.")
|
||||
return True
|
||||
|
||||
async def connection_made(
|
||||
self,
|
||||
io_proto: SanicProtocol,
|
||||
loop: asyncio.AbstractEventLoop | None = None,
|
||||
):
|
||||
if not loop:
|
||||
try:
|
||||
loop = getattr(io_proto, "loop")
|
||||
except AttributeError:
|
||||
loop = asyncio.get_event_loop()
|
||||
if not loop:
|
||||
# This catch is for mypy type checker
|
||||
# to assert loop is not None here.
|
||||
raise ServerError("Connection received with no asyncio loop.")
|
||||
if self.auto_closer_task:
|
||||
raise ServerError(
|
||||
"Cannot call connection_made more than once "
|
||||
"on a websocket connection."
|
||||
)
|
||||
self.loop = loop
|
||||
self.io_proto = io_proto
|
||||
self.connection_lost_waiter = self.loop.create_future()
|
||||
self.data_finished_fut = asyncio.shield(self.loop.create_future())
|
||||
|
||||
if self.ping_interval:
|
||||
self.keepalive_ping_task = asyncio.create_task(
|
||||
self.keepalive_ping()
|
||||
)
|
||||
self.auto_closer_task = asyncio.create_task(
|
||||
self.auto_close_connection()
|
||||
)
|
||||
|
||||
async def wait_for_connection_lost(self, timeout=None) -> bool:
|
||||
"""
|
||||
Wait until the TCP connection is closed or ``timeout`` elapses.
|
||||
If timeout is None, wait forever.
|
||||
Recommend you should pass in self.close_timeout as timeout
|
||||
|
||||
Return ``True`` if the connection is closed and ``False`` otherwise.
|
||||
|
||||
"""
|
||||
if not self.connection_lost_waiter:
|
||||
return False
|
||||
if self.connection_lost_waiter.done():
|
||||
return True
|
||||
else:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(self.connection_lost_waiter), timeout
|
||||
)
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
# Re-check self.connection_lost_waiter.done() synchronously
|
||||
# because connection_lost() could run between the moment the
|
||||
# timeout occurs and the moment this coroutine resumes running
|
||||
return self.connection_lost_waiter.done()
|
||||
|
||||
async def process_events(self, events: Sequence[Event]) -> None:
|
||||
"""
|
||||
Process a list of incoming events.
|
||||
"""
|
||||
# Wrapped in a mutex lock, to prevent other incoming events
|
||||
# from processing at the same time
|
||||
async with self.process_event_mutex:
|
||||
for event in events:
|
||||
if not isinstance(event, Frame):
|
||||
# Event is not a frame. Ignore it.
|
||||
continue
|
||||
if event.opcode == Opcode.PONG:
|
||||
await self.process_pong(event)
|
||||
elif event.opcode == Opcode.CLOSE:
|
||||
if self.recv_cancel:
|
||||
self.recv_cancel.cancel()
|
||||
else:
|
||||
await self.assembler.put(event)
|
||||
|
||||
async def process_pong(self, frame: Frame) -> None:
|
||||
if frame.data in self.pings:
|
||||
# Acknowledge all pings up to the one matching this pong.
|
||||
ping_ids = []
|
||||
for ping_id, ping in self.pings.items():
|
||||
ping_ids.append(ping_id)
|
||||
if not ping.done():
|
||||
ping.set_result(None)
|
||||
if ping_id == frame.data:
|
||||
break
|
||||
else: # noqa
|
||||
raise ServerError("ping_id is not in self.pings")
|
||||
# Remove acknowledged pings from self.pings.
|
||||
for ping_id in ping_ids:
|
||||
del self.pings[ping_id]
|
||||
|
||||
async def keepalive_ping(self) -> None:
|
||||
"""
|
||||
Send a Ping frame and wait for a Pong frame at regular intervals.
|
||||
This coroutine exits when the connection terminates and one of the
|
||||
following happens:
|
||||
- :meth:`ping` raises :exc:`ConnectionClosed`, or
|
||||
- :meth:`auto_close_connection` cancels :attr:`keepalive_ping_task`.
|
||||
"""
|
||||
if self.ping_interval is None:
|
||||
return
|
||||
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(self.ping_interval)
|
||||
|
||||
# ping() raises CancelledError if the connection is closed,
|
||||
# when auto_close_connection() cancels keepalive_ping_task.
|
||||
|
||||
# ping() raises ConnectionClosed if the connection is lost,
|
||||
# when connection_lost() calls abort_pings().
|
||||
|
||||
ping_waiter = await self.ping()
|
||||
|
||||
if self.ping_timeout is not None:
|
||||
try:
|
||||
await asyncio.wait_for(ping_waiter, self.ping_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
websockets_logger.warning(
|
||||
"Websocket timed out waiting for pong"
|
||||
)
|
||||
self.fail_connection(1011)
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
# It is expected for this task to be cancelled during during
|
||||
# normal operation, when the connection is closed.
|
||||
websockets_logger.debug(
|
||||
"Websocket keepalive ping task was cancelled."
|
||||
)
|
||||
except (ConnectionClosed, WebsocketClosed):
|
||||
websockets_logger.debug(
|
||||
"Websocket closed. Keepalive ping task exiting."
|
||||
)
|
||||
except Exception as e:
|
||||
websockets_logger.warning(
|
||||
"Unexpected exception in websocket keepalive ping task."
|
||||
)
|
||||
websockets_logger.debug(str(e))
|
||||
|
||||
def _force_disconnect(self) -> bool:
|
||||
"""
|
||||
Internal method used by end_connection and fail_connection
|
||||
only when the graceful auto-closer cannot be used
|
||||
"""
|
||||
if self.auto_closer_task and not self.auto_closer_task.done():
|
||||
self.auto_closer_task.cancel()
|
||||
if self.data_finished_fut and not self.data_finished_fut.done():
|
||||
self.data_finished_fut.cancel()
|
||||
self.data_finished_fut = None
|
||||
if self.keepalive_ping_task and not self.keepalive_ping_task.done():
|
||||
self.keepalive_ping_task.cancel()
|
||||
self.keepalive_ping_task = None
|
||||
if self.loop and self.io_proto and self.io_proto.transport:
|
||||
self.io_proto.transport.close()
|
||||
self.loop.call_later(
|
||||
self.close_timeout, self.io_proto.transport.abort
|
||||
)
|
||||
# We were never open, or already closed
|
||||
return True
|
||||
|
||||
def fail_connection(self, code: int = 1006, reason: str = "") -> bool:
|
||||
"""
|
||||
Fail the WebSocket Connection
|
||||
This requires:
|
||||
1. Stopping all processing of incoming data, which means cancelling
|
||||
pausing the underlying io protocol. The close code will be 1006
|
||||
unless a close frame was received earlier.
|
||||
2. Sending a close frame with an appropriate code if the opening
|
||||
handshake succeeded and the other side is likely to process it.
|
||||
3. Closing the connection. :meth:`auto_close_connection` takes care
|
||||
of this.
|
||||
(The specification describes these steps in the opposite order.)
|
||||
"""
|
||||
if self.io_proto and self.io_proto.transport:
|
||||
# Stop new data coming in
|
||||
# In Python Version 3.7: pause_reading is idempotent
|
||||
# ut can be called when the transport is already paused or closed
|
||||
self.io_proto.transport.pause_reading()
|
||||
|
||||
# Keeping fail_connection() synchronous guarantees it can't
|
||||
# get stuck and simplifies the implementation of the callers.
|
||||
# Not draining the write buffer is acceptable in this context.
|
||||
|
||||
# clear the send buffer
|
||||
_ = self.ws_proto.data_to_send()
|
||||
# If we're not already CLOSED or CLOSING, then send the close.
|
||||
if self.ws_proto.state is OPEN:
|
||||
if code in (1000, 1001):
|
||||
self.ws_proto.send_close(code, reason)
|
||||
else:
|
||||
self.ws_proto.fail(code, reason)
|
||||
try:
|
||||
data_to_send = self.ws_proto.data_to_send()
|
||||
while (
|
||||
len(data_to_send)
|
||||
and self.io_proto
|
||||
and self.io_proto.transport
|
||||
):
|
||||
frame_data = data_to_send.pop(0)
|
||||
self.io_proto.transport.write(frame_data)
|
||||
except Exception:
|
||||
# sending close frames may fail if the
|
||||
# transport closes during this period
|
||||
...
|
||||
if code == 1006:
|
||||
# Special case: 1006 consider the transport already closed
|
||||
self.ws_proto.state = CLOSED
|
||||
if self.data_finished_fut and not self.data_finished_fut.done():
|
||||
# We have a graceful auto-closer. Use it to close the connection.
|
||||
self.data_finished_fut.cancel()
|
||||
self.data_finished_fut = None
|
||||
if (not self.auto_closer_task) or self.auto_closer_task.done():
|
||||
return self._force_disconnect()
|
||||
return False
|
||||
|
||||
def end_connection(self, code=1000, reason=""):
|
||||
# This is like slightly more graceful form of fail_connection
|
||||
# Use this instead of close() when you need an immediate
|
||||
# close and cannot await websocket.close() handshake.
|
||||
|
||||
if code == 1006 or not self.io_proto or not self.io_proto.transport:
|
||||
return self.fail_connection(code, reason)
|
||||
|
||||
# Stop new data coming in
|
||||
# In Python Version 3.7: pause_reading is idempotent
|
||||
# i.e. it can be called when the transport is already paused or closed.
|
||||
self.io_proto.transport.pause_reading()
|
||||
if self.ws_proto.state == OPEN:
|
||||
data_to_send = self.ws_proto.data_to_send()
|
||||
self.ws_proto.send_close(code, reason)
|
||||
data_to_send.extend(self.ws_proto.data_to_send())
|
||||
try:
|
||||
while (
|
||||
len(data_to_send)
|
||||
and self.io_proto
|
||||
and self.io_proto.transport
|
||||
):
|
||||
frame_data = data_to_send.pop(0)
|
||||
self.io_proto.transport.write(frame_data)
|
||||
except Exception:
|
||||
# sending close frames may fail if the
|
||||
# transport closes during this period
|
||||
# But that doesn't matter at this point
|
||||
...
|
||||
if self.data_finished_fut and not self.data_finished_fut.done():
|
||||
# We have the ability to signal the auto-closer
|
||||
# try to trigger it to auto-close the connection
|
||||
self.data_finished_fut.cancel()
|
||||
self.data_finished_fut = None
|
||||
if (not self.auto_closer_task) or self.auto_closer_task.done():
|
||||
# Auto-closer is not running, do force disconnect
|
||||
return self._force_disconnect()
|
||||
return False
|
||||
|
||||
async def auto_close_connection(self) -> None:
|
||||
"""
|
||||
Close the WebSocket Connection
|
||||
When the opening handshake succeeds, :meth:`connection_open` starts
|
||||
this coroutine in a task. It waits for the data transfer phase to
|
||||
complete then it closes the TCP connection cleanly.
|
||||
When the opening handshake fails, :meth:`fail_connection` does the
|
||||
same. There's no data transfer phase in that case.
|
||||
"""
|
||||
try:
|
||||
# Wait for the data transfer phase to complete.
|
||||
if self.data_finished_fut:
|
||||
try:
|
||||
await self.data_finished_fut
|
||||
websockets_logger.debug(
|
||||
"Websocket task finished. Closing the connection."
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
# Cancelled error is called when data phase is cancelled
|
||||
# if an error occurred or the client closed the connection
|
||||
websockets_logger.debug(
|
||||
"Websocket handler cancelled. Closing the connection."
|
||||
)
|
||||
|
||||
# Cancel the keepalive ping task.
|
||||
if self.keepalive_ping_task:
|
||||
self.keepalive_ping_task.cancel()
|
||||
self.keepalive_ping_task = None
|
||||
|
||||
# Half-close the TCP connection if possible (when there's no TLS).
|
||||
if (
|
||||
self.io_proto
|
||||
and self.io_proto.transport
|
||||
and self.io_proto.transport.can_write_eof()
|
||||
):
|
||||
websockets_logger.debug(
|
||||
"Websocket half-closing TCP connection"
|
||||
)
|
||||
try:
|
||||
self.io_proto.transport.write_eof()
|
||||
except RuntimeError:
|
||||
...
|
||||
if self.connection_lost_waiter:
|
||||
if await self.wait_for_connection_lost(timeout=0):
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
...
|
||||
except BaseException:
|
||||
websockets_logger.exception("Error closing websocket connection")
|
||||
finally:
|
||||
# Does this still exist?
|
||||
if self.keepalive_ping_task:
|
||||
self.keepalive_ping_task.cancel()
|
||||
self.keepalive_ping_task = None
|
||||
# The try/finally ensures that the transport never remains open,
|
||||
# even if this coroutine is cancelled (for example).
|
||||
if (not self.io_proto) or (not self.io_proto.transport):
|
||||
# we were never open, or done. Can't do any finalization.
|
||||
return
|
||||
elif (
|
||||
self.connection_lost_waiter
|
||||
and self.connection_lost_waiter.done()
|
||||
):
|
||||
# connection confirmed closed already, proceed to abort waiter
|
||||
...
|
||||
elif self.io_proto.transport.is_closing():
|
||||
# Connection is already closing (due to half-close above)
|
||||
# proceed to abort waiter
|
||||
...
|
||||
else:
|
||||
self.io_proto.transport.close()
|
||||
if not self.connection_lost_waiter:
|
||||
# Our connection monitor task isn't running.
|
||||
try:
|
||||
await asyncio.sleep(self.close_timeout)
|
||||
except asyncio.CancelledError:
|
||||
...
|
||||
if self.io_proto and self.io_proto.transport:
|
||||
self.io_proto.transport.abort()
|
||||
else:
|
||||
if await self.wait_for_connection_lost(
|
||||
timeout=self.close_timeout
|
||||
):
|
||||
# Connection aborted before the timeout expired.
|
||||
return
|
||||
websockets_logger.warning(
|
||||
"Timeout waiting for TCP connection to close. Aborting"
|
||||
)
|
||||
if self.io_proto and self.io_proto.transport:
|
||||
self.io_proto.transport.abort()
|
||||
|
||||
def abort_pings(self) -> None:
|
||||
"""
|
||||
Raise ConnectionClosed in pending keepalive pings.
|
||||
They'll never receive a pong once the connection is closed.
|
||||
"""
|
||||
if self.ws_proto.state is not CLOSED:
|
||||
raise ServerError(
|
||||
"Webscoket about_pings should only be called "
|
||||
"after connection state is changed to CLOSED"
|
||||
)
|
||||
|
||||
for ping in self.pings.values():
|
||||
ping.set_exception(ConnectionClosedError(None, None))
|
||||
# If the exception is never retrieved, it will be logged when ping
|
||||
# is garbage-collected. This is confusing for users.
|
||||
# Given that ping is done (with an exception), canceling it does
|
||||
# nothing, but it prevents logging the exception.
|
||||
ping.cancel()
|
||||
|
||||
async def close(self, code: int = 1000, reason: str = "") -> None:
|
||||
"""
|
||||
Perform the closing handshake.
|
||||
This is a websocket-protocol level close.
|
||||
:meth:`close` waits for the other end to complete the handshake and
|
||||
for the TCP connection to terminate.
|
||||
:meth:`close` is idempotent: it doesn't do anything once the
|
||||
connection is closed.
|
||||
:param code: WebSocket close code
|
||||
:param reason: WebSocket close reason
|
||||
"""
|
||||
if code == 1006:
|
||||
self.fail_connection(code, reason)
|
||||
return
|
||||
async with self.conn_mutex:
|
||||
if self.ws_proto.state is OPEN:
|
||||
self.ws_proto.send_close(code, reason)
|
||||
data_to_send = self.ws_proto.data_to_send()
|
||||
await self.send_data(data_to_send)
|
||||
|
||||
async def recv(self, timeout: float | None = None) -> Data | None:
|
||||
"""
|
||||
Receive the next message.
|
||||
Return a :class:`str` for a text frame and :class:`bytes` for a binary
|
||||
frame.
|
||||
When the end of the message stream is reached, :meth:`recv` raises
|
||||
:exc:`~websockets.exceptions.ConnectionClosed`. Specifically, it
|
||||
raises :exc:`~websockets.exceptions.ConnectionClosedOK` after a normal
|
||||
connection closure and
|
||||
:exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
|
||||
error or a network failure.
|
||||
If ``timeout`` is ``None``, block until a message is received. Else,
|
||||
if no message is received within ``timeout`` seconds, return ``None``.
|
||||
Set ``timeout`` to ``0`` to check if a message was already received.
|
||||
:raises ~websockets.exceptions.ConnectionClosed: when the
|
||||
connection is closed
|
||||
:raises asyncio.CancelledError: if the websocket closes while waiting
|
||||
:raises ServerError: if two tasks call :meth:`recv` or
|
||||
:meth:`recv_streaming` concurrently
|
||||
"""
|
||||
|
||||
if self.recv_lock.locked():
|
||||
raise ServerError(
|
||||
"cannot call recv while another task is "
|
||||
"already waiting for the next message"
|
||||
)
|
||||
await self.recv_lock.acquire()
|
||||
if self.ws_proto.state is CLOSED:
|
||||
self.recv_lock.release()
|
||||
raise WebsocketClosed(
|
||||
"Cannot receive from websocket interface after it is closed."
|
||||
)
|
||||
assembler_get: asyncio.Task | None = None
|
||||
try:
|
||||
self.recv_cancel = asyncio.Future()
|
||||
assembler_get = asyncio.create_task(self.assembler.get(timeout))
|
||||
tasks = (self.recv_cancel, assembler_get)
|
||||
done, pending = await asyncio.wait(
|
||||
tasks,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
done_task = next(iter(done))
|
||||
if done_task is self.recv_cancel:
|
||||
# recv was cancelled
|
||||
for p in pending:
|
||||
p.cancel()
|
||||
raise asyncio.CancelledError()
|
||||
else:
|
||||
self.recv_cancel.cancel()
|
||||
return done_task.result()
|
||||
except asyncio.CancelledError:
|
||||
# recv was cancelled
|
||||
if assembler_get:
|
||||
assembler_get.cancel()
|
||||
raise
|
||||
finally:
|
||||
self.recv_cancel = None
|
||||
self.recv_lock.release()
|
||||
|
||||
async def recv_burst(self, max_recv=256) -> Sequence[Data]:
|
||||
"""
|
||||
Receive the messages which have arrived since last checking.
|
||||
Return a :class:`list` containing :class:`str` for a text frame
|
||||
and :class:`bytes` for a binary frame.
|
||||
When the end of the message stream is reached, :meth:`recv_burst`
|
||||
raises :exc:`~websockets.exceptions.ConnectionClosed`. Specifically,
|
||||
it raises :exc:`~websockets.exceptions.ConnectionClosedOK` after a
|
||||
normal connection closure and
|
||||
:exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
|
||||
error or a network failure.
|
||||
:raises ~websockets.exceptions.ConnectionClosed: when the
|
||||
connection is closed
|
||||
:raises ServerError: if two tasks call :meth:`recv_burst` or
|
||||
:meth:`recv_streaming` concurrently
|
||||
"""
|
||||
|
||||
if self.recv_lock.locked():
|
||||
raise ServerError(
|
||||
"cannot call recv_burst while another task is already waiting "
|
||||
"for the next message"
|
||||
)
|
||||
await self.recv_lock.acquire()
|
||||
if self.ws_proto.state is CLOSED:
|
||||
self.recv_lock.release()
|
||||
raise WebsocketClosed(
|
||||
"Cannot receive from websocket interface after it is closed."
|
||||
)
|
||||
messages = []
|
||||
assembler_get: asyncio.Task | None = None
|
||||
try:
|
||||
# Prevent pausing the transport when we're
|
||||
# receiving a burst of messages
|
||||
self.can_pause = False
|
||||
self.recv_cancel = asyncio.Future()
|
||||
while True:
|
||||
assembler_get = asyncio.create_task(self.assembler.get(0))
|
||||
tasks = (self.recv_cancel, assembler_get)
|
||||
done, pending = await asyncio.wait(
|
||||
tasks,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
done_task = next(iter(done))
|
||||
if done_task is self.recv_cancel:
|
||||
# recv_burst was cancelled
|
||||
for p in pending:
|
||||
p.cancel()
|
||||
raise asyncio.CancelledError()
|
||||
m = done_task.result()
|
||||
if m is None:
|
||||
# None left in the burst. This is good!
|
||||
break
|
||||
messages.append(m)
|
||||
if len(messages) >= max_recv:
|
||||
# Too much data in the pipe. Hit our burst limit.
|
||||
break
|
||||
# Allow an eventloop iteration for the
|
||||
# next message to pass into the Assembler
|
||||
await asyncio.sleep(0)
|
||||
self.recv_cancel.cancel()
|
||||
except asyncio.CancelledError:
|
||||
# recv_burst was cancelled
|
||||
if assembler_get:
|
||||
assembler_get.cancel()
|
||||
raise
|
||||
finally:
|
||||
self.recv_cancel = None
|
||||
self.can_pause = True
|
||||
self.recv_lock.release()
|
||||
return messages
|
||||
|
||||
async def recv_streaming(self) -> AsyncIterator[Data]:
|
||||
"""
|
||||
Receive the next message frame by frame.
|
||||
Return an iterator of :class:`str` for a text frame and :class:`bytes`
|
||||
for a binary frame. The iterator should be exhausted, or else the
|
||||
connection will become unusable.
|
||||
With the exception of the return value, :meth:`recv_streaming` behaves
|
||||
like :meth:`recv`.
|
||||
"""
|
||||
if self.recv_lock.locked():
|
||||
raise ServerError(
|
||||
"Cannot call recv_streaming while another task "
|
||||
"is already waiting for the next message"
|
||||
)
|
||||
await self.recv_lock.acquire()
|
||||
if self.ws_proto.state is CLOSED:
|
||||
self.recv_lock.release()
|
||||
raise WebsocketClosed(
|
||||
"Cannot receive from websocket interface after it is closed."
|
||||
)
|
||||
try:
|
||||
cancelled = False
|
||||
self.recv_cancel = asyncio.Future()
|
||||
self.can_pause = False
|
||||
async for m in self.assembler.get_iter():
|
||||
if self.recv_cancel.done():
|
||||
cancelled = True
|
||||
break
|
||||
yield m
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError()
|
||||
finally:
|
||||
self.can_pause = True
|
||||
self.recv_cancel = None
|
||||
self.recv_lock.release()
|
||||
|
||||
async def send(self, message: Data | Iterable[Data]) -> None:
|
||||
"""
|
||||
Send a message.
|
||||
A string (:class:`str`) is sent as a `Text frame`_. A bytestring or
|
||||
bytes-like object (:class:`bytes`, :class:`bytearray`, or
|
||||
:class:`memoryview`) is sent as a `Binary frame`_.
|
||||
.. _Text frame: https://tools.ietf.org/html/rfc6455#section-5.6
|
||||
.. _Binary frame: https://tools.ietf.org/html/rfc6455#section-5.6
|
||||
:meth:`send` also accepts an iterable of strings, bytestrings, or
|
||||
bytes-like objects. In that case the message is fragmented. Each item
|
||||
is treated as a message fragment and sent in its own frame. All items
|
||||
must be of the same type, or else :meth:`send` will raise a
|
||||
:exc:`TypeError` and the connection will be closed.
|
||||
:meth:`send` rejects dict-like objects because this is often an error.
|
||||
If you wish to send the keys of a dict-like object as fragments, call
|
||||
its :meth:`~dict.keys` method and pass the result to :meth:`send`.
|
||||
:raises TypeError: for unsupported inputs
|
||||
"""
|
||||
async with self.conn_mutex:
|
||||
if self.ws_proto.state in (CLOSED, CLOSING):
|
||||
raise WebsocketClosed(
|
||||
"Cannot write to websocket interface after it is closed."
|
||||
)
|
||||
if (not self.data_finished_fut) or self.data_finished_fut.done():
|
||||
raise ServerError(
|
||||
"Cannot write to websocket interface after it is finished."
|
||||
)
|
||||
|
||||
# Unfragmented message -- this case must be handled first because
|
||||
# strings and bytes-like objects are iterable.
|
||||
|
||||
if isinstance(message, str):
|
||||
self.ws_proto.send_text(message.encode("utf-8"))
|
||||
await self.send_data(self.ws_proto.data_to_send())
|
||||
|
||||
elif isinstance(message, (bytes, bytearray, memoryview)):
|
||||
self.ws_proto.send_binary(message)
|
||||
await self.send_data(self.ws_proto.data_to_send())
|
||||
|
||||
elif isinstance(message, Mapping):
|
||||
# Catch a common mistake -- passing a dict to send().
|
||||
raise TypeError("data is a dict-like object")
|
||||
|
||||
elif isinstance(message, Iterable):
|
||||
# Fragmented message -- regular iterator.
|
||||
raise NotImplementedError(
|
||||
"Fragmented websocket messages are not supported."
|
||||
)
|
||||
else:
|
||||
raise TypeError("Websocket data must be bytes, str.")
|
||||
|
||||
async def ping(self, data: Data | None = None) -> asyncio.Future:
|
||||
"""
|
||||
Send a ping.
|
||||
Return an :class:`~asyncio.Future` that will be resolved when the
|
||||
corresponding pong is received. You can ignore it if you don't intend
|
||||
to wait.
|
||||
A ping may serve as a keepalive or as a check that the remote endpoint
|
||||
received all messages up to this point::
|
||||
await pong_event = ws.ping()
|
||||
await pong_event # only if you want to wait for the pong
|
||||
By default, the ping contains four random bytes. This payload may be
|
||||
overridden with the optional ``data`` argument which must be a string
|
||||
(which will be encoded to UTF-8) or a bytes-like object.
|
||||
"""
|
||||
async with self.conn_mutex:
|
||||
if self.ws_proto.state in (CLOSED, CLOSING):
|
||||
raise WebsocketClosed(
|
||||
"Cannot send a ping when the websocket interface "
|
||||
"is closed."
|
||||
)
|
||||
if (not self.io_proto) or (not self.io_proto.loop):
|
||||
raise ServerError(
|
||||
"Cannot send a ping when the websocket has no I/O "
|
||||
"protocol attached."
|
||||
)
|
||||
if data is not None:
|
||||
if isinstance(data, str):
|
||||
data = data.encode("utf-8")
|
||||
elif isinstance(data, (bytearray, memoryview)):
|
||||
data = bytes(data)
|
||||
|
||||
# Protect against duplicates if a payload is explicitly set.
|
||||
if data in self.pings:
|
||||
raise ValueError(
|
||||
"already waiting for a pong with the same data"
|
||||
)
|
||||
|
||||
# Generate a unique random payload otherwise.
|
||||
while data is None or data in self.pings:
|
||||
data = secrets.token_bytes(4)
|
||||
|
||||
self.pings[data] = self.io_proto.loop.create_future()
|
||||
|
||||
self.ws_proto.send_ping(data)
|
||||
await self.send_data(self.ws_proto.data_to_send())
|
||||
|
||||
return asyncio.shield(self.pings[data])
|
||||
|
||||
async def pong(self, data: Data = b"") -> None:
|
||||
"""
|
||||
Send a pong.
|
||||
An unsolicited pong may serve as a unidirectional heartbeat.
|
||||
The payload may be set with the optional ``data`` argument which must
|
||||
be a string (which will be encoded to UTF-8) or a bytes-like object.
|
||||
"""
|
||||
async with self.conn_mutex:
|
||||
if self.ws_proto.state in (CLOSED, CLOSING):
|
||||
# Cannot send pong after transport is shutting down
|
||||
return
|
||||
if isinstance(data, str):
|
||||
data = data.encode("utf-8")
|
||||
elif isinstance(data, (bytearray, memoryview)):
|
||||
data = bytes(data)
|
||||
self.ws_proto.send_pong(data)
|
||||
await self.send_data(self.ws_proto.data_to_send())
|
||||
|
||||
async def send_data(self, data_to_send):
|
||||
for data in data_to_send:
|
||||
if data:
|
||||
await self.io_proto.send(data)
|
||||
else:
|
||||
# Send an EOF - We don't actually send it,
|
||||
# just trigger to autoclose the connection
|
||||
if (
|
||||
self.auto_closer_task
|
||||
and not self.auto_closer_task.done()
|
||||
and self.data_finished_fut
|
||||
and not self.data_finished_fut.done()
|
||||
):
|
||||
# Auto-close the connection
|
||||
self.data_finished_fut.set_result(None)
|
||||
else:
|
||||
# This will fail the connection appropriately
|
||||
SanicProtocol.close(self.io_proto, timeout=1.0)
|
||||
|
||||
async def async_data_received(self, data_to_send, events_to_process):
|
||||
if self.ws_proto.state in (OPEN, CLOSING) and len(data_to_send) > 0:
|
||||
# receiving data can generate data to send (eg, pong for a ping)
|
||||
# send connection.data_to_send()
|
||||
await self.send_data(data_to_send)
|
||||
if len(events_to_process) > 0:
|
||||
await self.process_events(events_to_process)
|
||||
|
||||
def data_received(self, data):
|
||||
self.ws_proto.receive_data(data)
|
||||
data_to_send = self.ws_proto.data_to_send()
|
||||
events_to_process = self.ws_proto.events_received()
|
||||
if len(data_to_send) > 0 or len(events_to_process) > 0:
|
||||
asyncio.create_task(
|
||||
self.async_data_received(data_to_send, events_to_process)
|
||||
)
|
||||
|
||||
async def async_eof_received(self, data_to_send, events_to_process):
|
||||
# receiving EOF can generate data to send
|
||||
# send connection.data_to_send()
|
||||
if self.ws_proto.state in (OPEN, CLOSING) and len(data_to_send) > 0:
|
||||
await self.send_data(data_to_send)
|
||||
if len(events_to_process) > 0:
|
||||
await self.process_events(events_to_process)
|
||||
if self.recv_cancel:
|
||||
self.recv_cancel.cancel()
|
||||
if (
|
||||
self.auto_closer_task
|
||||
and not self.auto_closer_task.done()
|
||||
and self.data_finished_fut
|
||||
and not self.data_finished_fut.done()
|
||||
):
|
||||
# Auto-close the connection
|
||||
self.data_finished_fut.set_result(None)
|
||||
# Cancel the running handler if its waiting
|
||||
else:
|
||||
# This will fail the connection appropriately
|
||||
SanicProtocol.close(self.io_proto, timeout=1.0)
|
||||
|
||||
def eof_received(self) -> bool | None:
|
||||
self.ws_proto.receive_eof()
|
||||
data_to_send = self.ws_proto.data_to_send()
|
||||
events_to_process = self.ws_proto.events_received()
|
||||
asyncio.create_task(
|
||||
self.async_eof_received(data_to_send, events_to_process)
|
||||
)
|
||||
return False
|
||||
|
||||
def connection_lost(self, exc):
|
||||
"""
|
||||
The WebSocket Connection is Closed.
|
||||
"""
|
||||
if not self.ws_proto.state == CLOSED:
|
||||
# signal to the websocket connection handler
|
||||
# we've lost the connection
|
||||
self.ws_proto.fail(code=1006)
|
||||
self.ws_proto.state = CLOSED
|
||||
|
||||
self.abort_pings()
|
||||
if self.connection_lost_waiter:
|
||||
self.connection_lost_waiter.set_result(None)
|
||||
|
||||
async def __aiter__(self):
|
||||
try:
|
||||
while True:
|
||||
yield await self.recv()
|
||||
except ConnectionClosedOK:
|
||||
return
|
||||
Reference in New Issue
Block a user