Commit 68b67867 authored by pgjones's avatar pgjones
Browse files

Support authenticating WebSocket connections

This allows for the login_required decorator and current_user to be
used for WebSocket connections but due to limitations with setting
cookies on WebSocket connections the login_user and logout_user
functionality cannot be used.
parent c1894e0c
Pipeline #280769759 failed with stages
in 3 minutes and 48 seconds
......@@ -7,3 +7,4 @@ Discussions
design_choices.rst
resolving.rst
websocket.rst
.. _websocket:
WebSocket Authentication
========================
Quart-Auth can be used to authenticate WebSocket connections, however
as cookies cannot be reliably set with a WebSocket connection it is
not possible to login or logout users via a WebSocket connection.
......@@ -12,7 +12,7 @@ in users.
.. code-block:: python
from quart import Quart, render_template_string
from quart import Quart, render_template_string, websocket
from quart_auth import (
AuthUser, AuthManager, current_user, login_required, login_user, logout_user
)
......@@ -50,3 +50,8 @@ in users.
Hello logged out user
{% endif %}
""")
@app.websocket("/ws")
async def ws():
await websocket.send(f"Hello current_user.auth_id")
...
import warnings
from enum import auto, Enum
from functools import wraps
from hashlib import sha512
from typing import Any, Callable, Dict, Optional
from itsdangerous import BadSignature, URLSafeSerializer
from quart import current_app, has_request_context, Quart, request, Response
from quart import (
current_app,
has_request_context,
has_websocket_context,
Quart,
request,
Response,
websocket,
)
from quart.exceptions import Unauthorized as QuartUnauthorized
from quart.globals import _request_ctx_stack
from quart.globals import _request_ctx_stack, _websocket_ctx_stack
from quart.local import LocalProxy
QUART_AUTH_USER_ATTRIBUTE = "_quart_auth_user"
......@@ -71,6 +80,7 @@ class AuthManager:
def init_app(self, app: Quart) -> None:
app.auth_manager = self # type: ignore
app.after_request(self.after_request)
app.after_websocket(self.after_websocket)
app.context_processor(_template_context)
def resolve_user(self) -> AuthUser:
......@@ -80,7 +90,11 @@ class AuthManager:
def load_cookie(self) -> Optional[str]:
try:
token = request.cookies[_get_config_or_default("QUART_AUTH_COOKIE_NAME")]
token = ""
if has_request_context():
token = request.cookies[_get_config_or_default("QUART_AUTH_COOKIE_NAME")]
elif has_websocket_context():
token = websocket.cookies[_get_config_or_default("QUART_AUTH_COOKIE_NAME")]
except KeyError:
return None
else:
......@@ -125,6 +139,18 @@ class AuthManager:
)
return response
def after_websocket(self, response: Optional[Response]) -> Optional[Response]:
if current_user.action != Action.PASS:
if response is not None:
warnings.warn(
"The auth cookie may not be set by the client. "
"Cookies are unreliably set on websocket responses."
)
else:
warnings.warn("The auth cookie cannot be set by the client.")
return response
def login_required(func: Callable) -> Callable:
"""A decorator to restrict route access to authenticated users.
......@@ -172,14 +198,20 @@ def login_user(user: AuthUser, remember: bool = False) -> None:
user.action = Action.WRITE_PERMANENT
else:
user.action = Action.WRITE
setattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
if has_request_context():
setattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
else:
raise RuntimeError("Cannot login unless within a request context")
def logout_user() -> None:
"""Use this to end the session of the current_user."""
user = current_app.auth_manager.user_class(None)
user.action = Action.DELETE
setattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
if has_request_context():
setattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
else:
raise RuntimeError("Cannot logout unless within a request context")
def renew_login() -> None:
......@@ -188,15 +220,28 @@ def renew_login() -> None:
def _load_user() -> AuthUser:
if has_request_context() and not hasattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE):
user = current_app.auth_manager.resolve_user()
setattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
return getattr(
_request_ctx_stack.top,
QUART_AUTH_USER_ATTRIBUTE,
current_app.auth_manager.user_class(None),
)
if has_request_context():
if not hasattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE):
user = current_app.auth_manager.resolve_user()
setattr(_request_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
return getattr(
_request_ctx_stack.top,
QUART_AUTH_USER_ATTRIBUTE,
current_app.auth_manager.user_class(None),
)
elif has_websocket_context():
if not hasattr(_websocket_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE):
user = current_app.auth_manager.resolve_user()
setattr(_websocket_ctx_stack.top, QUART_AUTH_USER_ATTRIBUTE, user)
return getattr(
_websocket_ctx_stack.top,
QUART_AUTH_USER_ATTRIBUTE,
current_app.auth_manager.user_class(None),
)
else:
return current_app.auth_manager.user_class(None)
def _get_config_or_default(config_key: str) -> Any:
......
import pytest
from quart import Quart, redirect, render_template_string, ResponseReturnValue, url_for
from quart import Quart, redirect, render_template_string, ResponseReturnValue, url_for, websocket
from werkzeug.datastructures import Headers
from quart_auth import (
......@@ -40,6 +40,12 @@ def _app() -> Quart:
login_user(AuthUser("2"))
return "login"
@app.websocket("/ws")
@login_required
async def ws() -> None:
data = await websocket.receive()
await websocket.send(f"{data} {current_user.auth_id}")
@app.route("/renew")
async def renew() -> ResponseReturnValue:
renew_login()
......@@ -132,3 +138,12 @@ async def test_renew_login(app: Quart) -> None:
assert next(cookie for cookie in test_client.cookie_jar).expires is None
await test_client.get("/renew")
assert next(cookie for cookie in test_client.cookie_jar).expires is not None
@pytest.mark.asyncio
async def test_websocket_login(app: Quart) -> None:
test_client = app.test_client()
await test_client.get("/login")
async with test_client.websocket("/ws") as ws:
await ws.send("Hello")
assert (await ws.receive()) == "Hello 2"
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment