Spaces:
Running on Zero
Running on Zero
File size: 2,328 Bytes
36cdb93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Shared token gate for the demo server.
Starlette's `BaseHTTPMiddleware` never sees websocket scopes, and the frame stream is
a websocket, so this is a plain ASGI wrapper rather than a FastAPI middleware.
A token supplied in the query string is echoed back as a cookie, so visiting
`/?token=...` once is enough: the websocket handshake then carries it automatically
and the URL can be shared without the query string surviving in every request.
"""
from urllib.parse import parse_qs
COOKIE = "livewan_token"
class TokenAuth:
def __init__(self, app, token):
self.app, self.token = app, token
def _supplied(self, scope):
headers = {k.lower(): v for k, v in (scope.get("headers") or [])}
token = parse_qs(scope.get("query_string", b"").decode()).get("token", [None])[0]
if token:
return token
auth = headers.get(b"authorization", b"").decode()
if auth.startswith("Bearer "):
return auth[7:]
for part in headers.get(b"cookie", b"").decode().split(";"):
k, _, v = part.strip().partition("=")
if k == COOKIE:
return v
return None
async def __call__(self, scope, receive, send):
if not self.token or scope["type"] not in ("http", "websocket"):
return await self.app(scope, receive, send)
if self._supplied(scope) != self.token:
if scope["type"] == "websocket":
return await send({"type": "websocket.close", "code": 1008})
await send({"type": "http.response.start", "status": 401,
"headers": [(b"content-type", b"text/plain; charset=utf-8")]})
return await send({"type": "http.response.body",
"body": b"unauthorized -- append ?token=<token> to the URL"})
if scope["type"] == "http":
async def _send(msg):
if msg["type"] == "http.response.start":
msg["headers"] = list(msg.get("headers") or []) + [
(b"set-cookie",
f"{COOKIE}={self.token}; Path=/; SameSite=Lax; Max-Age=604800".encode())
]
await send(msg)
return await self.app(scope, receive, _send)
await self.app(scope, receive, send)
|