125 lines
4.5 KiB
Python
125 lines
4.5 KiB
Python
"""SecSocket: a quantumpipe Session pumped over a WebSocket (either the
|
|
`websockets` library or FastAPI/Starlette). send() picks the frame kind
|
|
from the object (str -> TEXT, bytes -> BINARY, Message passes through),
|
|
receive() maps it back. See docs/design.md."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections.abc import Awaitable, Callable, MutableMapping
|
|
from typing import Any
|
|
|
|
from .exceptions import ProtocolError
|
|
from .sansio import Established, Message, Session, Verification
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
Handler = Callable[[Any], Awaitable[Any]]
|
|
|
|
# this wrapper's user kinds; kind 0 is reserved for in-band protocol use
|
|
BINARY = 1 # opaque bytes, delivered as bytes
|
|
TEXT = 2 # utf-8 text, delivered decoded as str
|
|
|
|
|
|
class SecSocket:
|
|
"""A quantumpipe Session pumped over a WebSocket; the Session is
|
|
exposed as .session (state, verified, verification, confirm,
|
|
credential)."""
|
|
|
|
def __init__(
|
|
self,
|
|
ws: Any,
|
|
store: MutableMapping,
|
|
key_id: Any = None,
|
|
*,
|
|
ad: bytes = b"",
|
|
handlers: dict[type, Handler] | None = None,
|
|
):
|
|
self._ws = ws
|
|
self.session = Session(store, key_id, ad=ad)
|
|
self._established = asyncio.Event()
|
|
self._handlers = handlers or {}
|
|
|
|
async def accept(self) -> None:
|
|
"""Server side: accept the socket if it needs accepting."""
|
|
if hasattr(self._ws, "accept"):
|
|
await self._ws.accept()
|
|
await self._run(initiate=False)
|
|
|
|
async def connect(self) -> None:
|
|
"""Client side: initiate the handshake."""
|
|
await self._run(initiate=True)
|
|
|
|
async def send(self, data: bytes | str | Message) -> None:
|
|
"""str -> TEXT frame, bytes -> BINARY frame, Message passes
|
|
through with its own kind."""
|
|
if isinstance(data, Message):
|
|
msg = data
|
|
elif isinstance(data, str):
|
|
msg = Message(data.encode(), TEXT)
|
|
else:
|
|
msg = Message(bytes(data), BINARY)
|
|
await self._dispatch(self.session(msg))
|
|
|
|
async def receive(self) -> bytes | str | Message:
|
|
"""BINARY -> bytes, TEXT -> str, other kinds -> raw Message."""
|
|
(msg,) = self.session(await self._recv())
|
|
if not isinstance(msg, Message):
|
|
raise ProtocolError("expected a traffic frame")
|
|
if msg.kind == BINARY:
|
|
return msg.data
|
|
if msg.kind == TEXT:
|
|
try:
|
|
return bytes(msg.data).decode()
|
|
except UnicodeDecodeError as e:
|
|
raise ProtocolError(f"invalid utf-8 in text frame: {e}") from None
|
|
return msg
|
|
|
|
async def close(self) -> None:
|
|
await self._ws.close()
|
|
|
|
async def _run(self, initiate: bool) -> None:
|
|
try:
|
|
if initiate:
|
|
await self._dispatch(self.session())
|
|
while not self._established.is_set():
|
|
await self._dispatch(self.session(await self._recv()))
|
|
except BaseException:
|
|
await self._ws.close()
|
|
raise
|
|
|
|
async def _dispatch(self, results: list) -> None:
|
|
for r in results:
|
|
if isinstance(r, bytes):
|
|
_log.debug("ws send frame[%d]", len(r))
|
|
await self._send(r)
|
|
continue
|
|
_log.debug("handle %s", type(r).__name__)
|
|
if isinstance(r, Established):
|
|
self._established.set()
|
|
handler = self._handlers.get(type(r))
|
|
if handler is not None:
|
|
out = await handler(r)
|
|
if out is not None and isinstance(r, Verification):
|
|
r.answer = out
|
|
if isinstance(r, Verification): # the answered result resumes the session
|
|
await self._dispatch(self.session(r))
|
|
|
|
# the two supported WebSocket flavors: the `websockets` library
|
|
# (send/recv) and FastAPI/Starlette (accept/send_bytes/receive_bytes)
|
|
|
|
async def _send(self, data: bytes) -> None:
|
|
if hasattr(self._ws, "send_bytes"): # FastAPI/Starlette
|
|
await self._ws.send_bytes(data)
|
|
else: # websockets library: send(str | bytes)
|
|
await self._ws.send(data)
|
|
|
|
async def _recv(self) -> bytes:
|
|
if hasattr(self._ws, "receive_bytes"): # FastAPI/Starlette
|
|
return await self._ws.receive_bytes()
|
|
data = await self._ws.recv() # websockets library
|
|
if not isinstance(data, bytes):
|
|
raise ProtocolError(f"expected a binary frame, got {type(data).__name__}")
|
|
return data
|