django-socket 0.2.0__py3-none-any.whl

django_socket/dispatch.py

@@ -0,0 +1,208 @@
+"""El nucleo: atiende los scopes 'websocket' y 'lifespan'.
+
+Vive aparte de `asgi.py` porque hay dos caminos que llegan aqui: el parche
+sobre `ASGIHandler` (modo cero-configuracion) y `ASGIApplication` explicito.
+Los dos comparten este codigo.
+"""
+
+from __future__ import annotations
+
+import logging
+from urllib.parse import urlparse
+
+from asgiref.sync import ThreadSensitiveContext
+
+from . import authentication
+from . import middleware as mw
+from . import groups, routing
+from .websocket import InvalidJSON, RateLimited, WebSocket, WebSocketDisconnect
+
+logger = logging.getLogger("django_socket")
+
+# Codigos de cierre propios (rango privado 4000-4999).
+CLOSE_NO_ROUTE = 4404
+CLOSE_BAD_DATA = 4400      # el cliente mando algo que no se puede parsear
+CLOSE_SERVER_ERROR = 1011
+
+_layer_started = False
+
+
+def _settings():
+    from django.conf import settings
+
+    return getattr(settings, "DJANGO_SOCKET", {}) or {}
+
+
+async def _start_layer() -> None:
+    global _layer_started
+    if not _layer_started:
+        await groups.get_layer().startup()
+        _layer_started = True
+
+
+async def _stop_layer() -> None:
+    global _layer_started
+    if _layer_started:
+        await groups.get_layer().shutdown()
+        _layer_started = False
+
+
+# ---------------------------------------------------------------- lifespan
+
+
+async def handle_lifespan(scope, receive, send) -> None:
+    while True:
+        message = await receive()
+        if message["type"] == "lifespan.startup":
+            try:
+                await _start_layer()
+            except Exception as exc:
+                logger.exception("django_socket: fallo en el arranque")
+                await send({"type": "lifespan.startup.failed", "message": str(exc)})
+                return
+            await send({"type": "lifespan.startup.complete"})
+        elif message["type"] == "lifespan.shutdown":
+            try:
+                await _stop_layer()
+            finally:
+                await send({"type": "lifespan.shutdown.complete"})
+            return
+
+
+# --------------------------------------------------------------- websocket
+
+
+async def handle_websocket(scope, receive, send) -> None:
+    # El primer evento del protocolo siempre es websocket.connect.
+    event = await receive()
+    if event["type"] != "websocket.connect":
+        return
+
+    await _start_layer()  # por si el servidor no soporta lifespan
+    sock = WebSocket(scope, receive, send, layer=groups.get_layer())
+
+    if not origin_allowed(sock):
+        logger.warning(
+            "django_socket: origen rechazado %r para %s",
+            sock.headers.get("origin"),
+            sock.path,
+        )
+        # En seco: un origen ajeno no debe tener un socket abierto ni un instante.
+        await sock.deny()
+        return
+
+    match = routing.resolve(sock.path)
+    if match is None:
+        logger.warning(
+            "django_socket: ninguna ruta casa con %s. Registradas: %s",
+            sock.path,
+            ", ".join(f"/{r.route}" for r in routing.get_routes()) or "(ninguna)",
+        )
+        await sock.close(CLOSE_NO_ROUTE, "No route")
+        return
+
+    route, kwargs = match
+    sock.path_params = kwargs
+
+    async with ThreadSensitiveContext():
+        if route.auth is not False:
+            await authentication.resolve(sock, route.auth)
+
+        from . import ratelimit
+
+        sock._rate = ratelimit.crear(route.rate_limit, route.burst)
+
+        async def ejecutar():
+            if route.group:
+                # group="room:{room}" se rellena con los parametros de la ruta.
+                await sock.join(route.group.format(**kwargs))
+            await route.handler(sock, **kwargs)
+
+        try:
+            # El middleware va por fuera del handler pero por dentro de este
+            # try, para que un fallo suyo se trate igual que uno del handler.
+            await mw.aplicar(sock, ejecutar)
+        except WebSocketDisconnect:
+            pass  # el cliente se fue; salida normal
+        except RateLimited as exc:
+            logger.warning(
+                "django_socket: %s va demasiado rapido en %s (%s)",
+                sock.client, sock.path, exc,
+            )
+            await sock.close(exc.code, f"Rate limit; retry in {exc.espera:.0f}s")
+            return
+        except InvalidJSON as exc:
+            # Culpa del cliente, no del servidor: un aviso y un codigo que lo
+            # diga. Nada de traceback ni de 1011, que harian pensar que el bug
+            # es tuyo cada vez que alguien mande basura por el socket.
+            logger.warning(
+                "django_socket: %s en %s (cliente %s)", exc, sock.path, sock.client
+            )
+            await sock.close(CLOSE_BAD_DATA, "Invalid JSON")
+            return
+        except Exception:
+            logger.exception("django_socket: excepcion en el handler de %s", sock.path)
+            await sock.close(CLOSE_SERVER_ERROR, "Internal error")
+            return
+        await sock.close()
+
+
+# ------------------------------------------------------------------ origen
+
+
+def origin_allowed(sock) -> bool:
+    """
+    Los WebSockets no estan sujetos a la politica de mismo origen: sin esta
+    comprobacion cualquier web podria abrir un socket autenticado contra la
+    tuya (cross-site WebSocket hijacking).
+
+    Un Origin ausente se acepta: los navegadores siempre lo mandan, asi que
+    solo lo omiten clientes nativos. Ponlo estricto con REQUIRE_ORIGIN.
+    """
+    from django.conf import settings
+
+    conf = _settings()
+    origin = sock.headers.get("origin")
+    if origin is None:
+        return not conf.get("REQUIRE_ORIGIN", False)
+
+    allowed = conf.get("ALLOWED_ORIGINS")
+    if allowed is not None:
+        if "*" in allowed:
+            return True
+        return origin in allowed or _host_of(origin) in allowed
+
+    # Por defecto: ALLOWED_HOSTS + CSRF_TRUSTED_ORIGINS, como hace Django.
+    host = _host_of(origin)
+    if not host:
+        return False
+
+    trusted = {
+        _host_of(o) for o in getattr(settings, "CSRF_TRUSTED_ORIGINS", []) or []
+    }
+    if host in trusted:
+        return True
+
+    hosts = list(getattr(settings, "ALLOWED_HOSTS", []) or [])
+    if settings.DEBUG and not hosts:
+        hosts = ["localhost", "127.0.0.1", "[::1]"]
+    return any(_host_matches(host, pattern) for pattern in hosts)
+
+
+def _host_of(origin: str) -> str:
+    """'https://ejemplo.com:8000' -> 'ejemplo.com'."""
+    try:
+        return (urlparse(origin).hostname or "").lower()
+    except ValueError:
+        return ""
+
+
+def _host_matches(host: str, pattern: str) -> bool:
+    pattern = pattern.lower()
+    if pattern == "*":
+        return True
+    if pattern.startswith("."):  # ".ejemplo.com" cubre subdominios y el apex
+        return host == pattern[1:] or host.endswith(pattern)
+    if pattern.startswith("*."):
+        return host == pattern[2:] or host.endswith(pattern[1:])
+    return host == pattern

django_socket/events.py

@@ -0,0 +1,143 @@
+"""Enrutado de mensajes JSON por tipo.
+
+Casi toda app que habla JSON por el socket manda `{"type": "algo", ...}` y
+acaba con un if/elif largo dentro del bucle. Esto lo convierte en funciones con
+nombre, sin que dejes de ver el flujo:
+
+    from django_socket import Events, ws
+
+    chat = Events()
+
+    @chat.on("mensaje")
+    async def mensaje(sock, datos):
+        await sock.broadcast({"type": "mensaje", "texto": datos["texto"]})
+
+    @chat.on("escribiendo")
+    async def escribiendo(sock):            # si no usas los datos, no los pidas
+        await sock.broadcast({"type": "escribiendo"}, exclude_self=True)
+
+    @ws("chat/<str:room>/", group="room:{room}")
+    async def handler(sock, room):
+        await chat.run(sock)
+
+Es opcional del todo: `async for msg in sock` sigue estando ahi y no necesitas
+saber que esto existe.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+
+from .websocket import InvalidJSON
+
+logger = logging.getLogger("django_socket")
+
+CUALQUIERA = "*"
+
+
+class Events:
+    """
+    Despacha mensajes JSON segun un campo (por defecto `"type"`).
+
+    `key`     -- el campo que decide el tipo. Usa `Events(key="action")` si tu
+                 protocolo se llama de otra forma.
+    `strict`  -- que hacer con un tipo que nadie maneja. Por defecto se ignora
+                 (y se anota en el log); con `strict=True` se cierra con 4400.
+    """
+
+    def __init__(self, key: str = "type", strict: bool = False):
+        self.key = key
+        self.strict = strict
+        self._handlers: dict[str, tuple] = {}
+
+    def on(self, *tipos: str):
+        """
+        Registra un handler para uno o varios tipos.
+
+            @chat.on("mensaje")
+            @chat.on("entrar", "salir")     # varios de golpe
+            @chat.on("*")                   # lo que no case con nada mas
+
+        El handler puede pedir los datos o no:
+
+            async def handler(sock, datos)  -> datos = el mensaje sin el campo key
+            async def handler(sock)         -> te basta con saber que llego
+        """
+        if not tipos:
+            raise TypeError("@on() necesita al menos un tipo: @on('mensaje')")
+
+        def decorator(fn):
+            if not inspect.iscoroutinefunction(fn):
+                raise TypeError(
+                    f"@on espera 'async def', y {fn.__name__} es una funcion "
+                    f"normal."
+                )
+            quiere_datos = _quiere_datos(fn)
+            for tipo in tipos:
+                if tipo in self._handlers:
+                    anterior = self._handlers[tipo][0].__name__
+                    raise ValueError(
+                        f"El tipo {tipo!r} ya lo maneja {anterior}."
+                    )
+                self._handlers[tipo] = (fn, quiere_datos)
+            return fn
+
+        return decorator
+
+    @property
+    def tipos(self) -> list[str]:
+        return sorted(self._handlers)
+
+    async def run(self, sock) -> None:
+        """Consume mensajes hasta que el cliente cierra."""
+        async for msg in sock:
+            await self.handle(sock, msg.json())
+
+    async def handle(self, sock, datos) -> None:
+        """Despacha un mensaje ya parseado. Util para testear un handler suelto."""
+        if not isinstance(datos, dict):
+            raise InvalidJSON(datos, f"se esperaba un objeto con {self.key!r}")
+
+        tipo = datos.get(self.key)
+        entrada = self._handlers.get(tipo) or self._handlers.get(CUALQUIERA)
+
+        if entrada is None:
+            if self.strict:
+                raise InvalidJSON(
+                    datos, f"tipo {tipo!r} desconocido; hay: {self.tipos}"
+                )
+            # A nivel WARNING a proposito: casi siempre es una errata en el
+            # nombre del tipo, y en DEBUG no la ve nadie. Si tu protocolo manda
+            # tipos que de verdad quieres ignorar, registra un @on("*") vacio.
+            logger.warning(
+                "django_socket: nadie maneja %s=%r en %s (registrados: %s)",
+                self.key, tipo, sock.path, self.tipos or "ninguno",
+            )
+            return
+
+        handler, quiere_datos = entrada
+        if quiere_datos:
+            await handler(sock, {k: v for k, v in datos.items() if k != self.key})
+        else:
+            await handler(sock)
+
+
+def _quiere_datos(fn) -> bool:
+    """
+    Mira si el handler pide el payload ademas del socket.
+
+    Se resuelve una vez al registrar, no en cada mensaje.
+    """
+    params = [
+        p for p in inspect.signature(fn).parameters.values()
+        if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
+    ]
+    if len(params) == 1:
+        return False
+    if len(params) == 2:
+        return True
+    raise TypeError(
+        f"{fn.__name__} debe aceptar (sock) o (sock, datos), no "
+        f"{len(params)} parametros posicionales."
+    )

django_socket/groups.py

@@ -0,0 +1,287 @@
+"""Capa de difusion: grupos y broadcast.
+
+`MemoryLayer` sirve para un solo proceso (dev, o un unico worker).
+`RedisLayer` reparte el fan-out entre procesos via pub/sub.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+logger = logging.getLogger("django_socket")
+
+
+class BaseLayer:
+    async def startup(self) -> None: ...
+
+    async def shutdown(self) -> None: ...
+
+    async def add(self, group: str, sock) -> None:
+        raise NotImplementedError
+
+    async def discard(self, group: str, sock) -> None:
+        raise NotImplementedError
+
+    async def send(self, group: str, data: Any, *, exclude=None) -> None:
+        raise NotImplementedError
+
+    async def size(self, group: str) -> int:
+        raise NotImplementedError
+
+
+class MemoryLayer(BaseLayer):
+    """Grupos en memoria del proceso. Sin dependencias externas."""
+
+    def __init__(self):
+        self._groups: dict[str, set] = {}
+
+    async def add(self, group: str, sock) -> None:
+        self._groups.setdefault(group, set()).add(sock)
+
+    async def discard(self, group: str, sock) -> None:
+        members = self._groups.get(group)
+        if members:
+            members.discard(sock)
+            if not members:
+                del self._groups[group]
+
+    async def size(self, group: str) -> int:
+        return len(self._groups.get(group, ()))
+
+    async def send(self, group: str, data: Any, *, exclude=None) -> None:
+        await self._deliver_local(group, data, exclude=exclude)
+
+    async def _deliver_local(self, group: str, data: Any, *, exclude=None) -> None:
+        """
+        Reparte encolando, sin esperar a que nadie lea.
+
+        Esperar seria el bug: un solo miembro que no consume dejaria colgado
+        para siempre al que difunde. Cada socket tiene un buzon acotado; si se
+        llena, ese cliente va demasiado atrasado y se le echa en vez de dejar
+        que arrastre a los demas.
+        """
+        miembros = [s for s in self._groups.get(group, ()) if s is not exclude]
+        if not miembros:
+            return
+
+        lentos = []
+        for sock in miembros:
+            if not await sock.enqueue(data):
+                lentos.append(sock)
+
+        # Cede el turno una vez por difusion, no una por miembro: es donde
+        # corren los escritores. Sin esto, un handler que difunde en bucle
+        # (`for fila in lote: await sock.broadcast(fila)`) no soltaria nunca el
+        # loop, los buzones se llenarian y acabaria echando a clientes sanos.
+        await asyncio.sleep(0)
+
+        for sock in lentos:
+            logger.warning(
+                "django_socket: %r no consume (buzon lleno); se le echa del "
+                "grupo %r. Sube DJANGO_SOCKET['SEND_QUEUE_MAX'] si tu caso "
+                "manda rafagas legitimas.",
+                sock, group,
+            )
+            sock.evict()
+            await self.discard(group, sock)
+
+
+class RedisLayer(MemoryLayer):
+    """
+    Mantiene los miembros locales igual que MemoryLayer, pero publica cada
+    broadcast en Redis para que los demas procesos entreguen a los suyos.
+
+    Sobrevive a que Redis se caiga: la entrega local sigue funcionando, los
+    sockets de los usuarios no se cierran, y al volver Redis el proceso se
+    resuscribe solo. Lo que se publique mientras esta caido se pierde -- esto
+    es pub/sub, no una cola.
+    """
+
+    ESPERA_MAX = 10.0          # tope del backoff al reconectar
+    LATIDO = 15.0              # health check de redis-py, en segundos
+
+    def __init__(self, url: str = "redis://localhost:6379/0", prefix: str = "djws"):
+        super().__init__()
+        self.url = url
+        self.channel = f"{prefix}:broadcast"
+        self._redis = None
+        self._pubsub = None
+        self._listener: asyncio.Task | None = None
+        self._conectado = False
+        # Identifica a este proceso para no entregar dos veces lo que ya
+        # entregamos localmente.
+        self._origin = f"{id(self)}"
+
+    async def startup(self) -> None:
+        try:
+            import redis.asyncio as aioredis
+        except ImportError as exc:  # pragma: no cover
+            raise RuntimeError(
+                "RedisLayer necesita el paquete 'redis'. Instalalo con: pip install redis"
+            ) from exc
+        self._redis = aioredis.from_url(
+            self.url,
+            health_check_interval=self.LATIDO,   # sin esto no detecta la caida
+            retry_on_timeout=True,
+        )
+        self._pubsub = self._redis.pubsub()
+        await self._pubsub.subscribe(self.channel)
+        self._conectado = True
+        self._listener = asyncio.create_task(self._listen())
+        logger.info("django_socket: RedisLayer conectada a %s", self.url)
+
+    async def shutdown(self) -> None:
+        self._conectado = False
+        if self._listener:
+            self._listener.cancel()
+            try:
+                await self._listener
+            except asyncio.CancelledError:
+                pass
+        if self._pubsub:
+            try:
+                await self._pubsub.unsubscribe(self.channel)
+            except Exception:
+                pass
+            await self._pubsub.aclose()
+        if self._redis:
+            await self._redis.aclose()
+
+    async def send(self, group: str, data: Any, *, exclude=None) -> None:
+        # Entrega local inmediata (asi `exclude` funciona por identidad)...
+        await self._deliver_local(group, data, exclude=exclude)
+        # ...y avisa al resto de procesos.
+        carga = json.dumps(
+            {"group": group, "data": data, "origin": self._origin}, default=str
+        )
+        try:
+            await self._redis.publish(self.channel, carga)
+        except Exception as exc:
+            # Que Redis falle no puede tumbar la conexion del usuario: la
+            # entrega local ya se hizo y el handler debe seguir vivo. Se pierde
+            # el fan-out a los demas procesos, y por eso se loguea como error.
+            logger.error(
+                "django_socket: no se pudo publicar en Redis (%s: %s). "
+                "El grupo %r solo recibio la entrega local.",
+                type(exc).__name__, exc, group,
+            )
+
+    async def _listen(self) -> None:
+        """
+        Bucle de escucha resistente.
+
+        Se usa `get_message(timeout=...)` en vez de `listen()` para tener un
+        despertar periodico: ahi es donde redis-py corre su health check y
+        donde podemos detectar que la conexion se fue. Con `listen()` a secas
+        el proceso se queda sordo para siempre tras una caida.
+        """
+        espera = 0.5
+        while True:
+            try:
+                mensaje = await self._pubsub.get_message(
+                    ignore_subscribe_messages=True, timeout=1.0
+                )
+                espera = 0.5                      # todo bien, resetea el backoff
+                if mensaje is not None and mensaje.get("type") == "message":
+                    await self._entregar(mensaje)
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                if not self._conectado:
+                    return
+                logger.warning(
+                    "django_socket: se perdio la conexion con Redis (%s). "
+                    "Reintentando en %.1fs.", type(exc).__name__, espera,
+                )
+                await asyncio.sleep(espera)
+                espera = min(espera * 2, self.ESPERA_MAX)
+                await self._resuscribir()
+
+    async def _entregar(self, mensaje) -> None:
+        try:
+            payload = json.loads(mensaje["data"])
+        except (ValueError, KeyError, TypeError):
+            logger.warning("django_socket: mensaje ilegible en %s", self.channel)
+            return
+        if payload.get("origin") == self._origin:
+            return  # ya lo entregamos localmente
+        await self._deliver_local(payload["group"], payload["data"])
+
+    async def _resuscribir(self) -> None:
+        """Vuelve a suscribirse tras un corte. Si Redis sigue caido, lo dira el bucle."""
+        try:
+            await self._pubsub.aclose()
+        except Exception:
+            pass
+        self._pubsub = self._redis.pubsub()
+        await self._pubsub.subscribe(self.channel)
+        logger.info("django_socket: resuscrito a %s", self.channel)
+
+
+# --------------------------------------------------------------- layer global
+
+_layer: BaseLayer | None = None
+
+
+def get_layer() -> BaseLayer:
+    global _layer
+    if _layer is None:
+        _layer = _build_layer()
+    return _layer
+
+
+def set_layer(layer: BaseLayer) -> None:
+    global _layer
+    _layer = layer
+
+
+def _build_layer() -> BaseLayer:
+    from django.conf import settings
+
+    conf = getattr(settings, "DJANGO_SOCKET", {}) or {}
+    backend = conf.get("LAYER", "memory")
+    if backend == "memory":
+        return MemoryLayer()
+    if backend == "redis":
+        return RedisLayer(
+            url=conf.get("REDIS_URL", "redis://localhost:6379/0"),
+            prefix=conf.get("PREFIX", "djws"),
+        )
+    if callable(backend):
+        return backend()
+    raise ValueError(
+        f"DJANGO_SOCKET['LAYER'] invalido: {backend!r}. Usa 'memory', 'redis' "
+        f"o un callable que devuelva una BaseLayer."
+    )
+
+
+# ------------------------------------------------------------- API de usuario
+
+
+async def broadcast(data: Any, *, to: str) -> None:
+    """
+    Envia a todos los miembros de un grupo, desde fuera de un handler.
+
+        await broadcast({"aviso": "mantenimiento"}, to="room:1")
+    """
+    await get_layer().send(to, data)
+
+
+async def group_size(group: str) -> int:
+    """Cuantos sockets locales hay en el grupo."""
+    return await get_layer().size(group)
+
+
+def broadcast_sync(data: Any, *, to: str) -> None:
+    """
+    Igual que `broadcast`, para vistas sincronas, señales o tareas Celery.
+
+    Con la capa 'memory' solo alcanza a los sockets del mismo proceso; para
+    llegar a todos los workers necesitas LAYER='redis'.
+    """
+    from asgiref.sync import async_to_sync
+
+    async_to_sync(broadcast)(data, to=to)