276 lines
9.4 KiB
Python
276 lines
9.4 KiB
Python
import time
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from asgiref.sync import sync_to_async
|
|
from channels.layers import get_channel_layer
|
|
from channels.testing import WebsocketCommunicator
|
|
|
|
from api.auth.models import UserAuthSession
|
|
from api.auth.websocket_tickets import (
|
|
WebSocketTicketError,
|
|
consume_websocket_ticket,
|
|
issue_websocket_ticket,
|
|
)
|
|
from communications.models import Chat, ChatMembership, Message, MessageReaction
|
|
from eusrr_backend.asgi import application
|
|
from realtime.protocol import ConnectionRateLimiter
|
|
|
|
pytestmark = pytest.mark.django_db(transaction=True)
|
|
|
|
|
|
def test_connection_rate_limiter_has_action_specific_limits():
|
|
limiter = ConnectionRateLimiter()
|
|
assert all(limiter.allow("typing") for _ in range(30))
|
|
assert limiter.allow("typing") is False
|
|
|
|
limiter = ConnectionRateLimiter()
|
|
assert all(limiter.allow("open_chat") for _ in range(120))
|
|
assert limiter.allow("open_chat") is False
|
|
|
|
limiter = ConnectionRateLimiter()
|
|
assert all(limiter.allow(f"unknown-{index}") for index in range(180))
|
|
assert limiter.allow("one-more-unique-action") is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_inactive_membership_is_not_subscribed(ws_communicator, user):
|
|
chat = await sync_to_async(Chat.objects.create)(name="Revoked", type="group")
|
|
await sync_to_async(ChatMembership.objects.create)(
|
|
chat=chat,
|
|
user=user,
|
|
role="member",
|
|
is_active=False,
|
|
)
|
|
communicator = await ws_communicator(user=user)
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
|
|
await get_channel_layer().group_send(
|
|
f"chat_{chat.id}",
|
|
{
|
|
"type": "chat_message",
|
|
"chat_id": chat.id,
|
|
"payload": {"id": 1, "content": "private"},
|
|
},
|
|
)
|
|
assert await communicator.receive_nothing(timeout=0.1)
|
|
await communicator.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reaction_requires_chat_access(ws_communicator, user, user2):
|
|
chat = await sync_to_async(Chat.objects.create)(name="Private", type="group")
|
|
await sync_to_async(ChatMembership.objects.create)(chat=chat, user=user2)
|
|
message = await sync_to_async(Message.objects.create)(
|
|
chat=chat,
|
|
author=user2,
|
|
content="secret",
|
|
)
|
|
communicator = await ws_communicator(user=user)
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
|
|
await communicator.send_json_to(
|
|
{"action": "add_reaction", "message_id": message.id, "emoji": "👍"}
|
|
)
|
|
response = await communicator.receive_json_from()
|
|
assert response["code"] == "access_denied"
|
|
reaction_exists = await sync_to_async(
|
|
MessageReaction.objects.filter(message=message).exists
|
|
)()
|
|
assert not reaction_exists
|
|
await communicator.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_membership_revocation_removes_live_group(ws_communicator, user):
|
|
chat = await sync_to_async(Chat.objects.create)(name="Live", type="group")
|
|
membership = await sync_to_async(ChatMembership.objects.create)(
|
|
chat=chat,
|
|
user=user,
|
|
)
|
|
communicator = await ws_communicator(user=user)
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
|
|
membership.is_active = False
|
|
await sync_to_async(membership.save)(update_fields=["is_active"])
|
|
access_event = await communicator.receive_json_from()
|
|
assert access_event == {
|
|
"type": "chat_access_changed",
|
|
"chat_id": chat.id,
|
|
"can_access": False,
|
|
}
|
|
|
|
await get_channel_layer().group_send(
|
|
f"chat_{chat.id}",
|
|
{
|
|
"type": "chat_message",
|
|
"chat_id": chat.id,
|
|
"payload": {"id": 2, "content": "after revoke"},
|
|
},
|
|
)
|
|
assert await communicator.receive_nothing(timeout=0.1)
|
|
await communicator.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_membership_addition_joins_live_group(ws_communicator, user):
|
|
chat = await sync_to_async(Chat.objects.create)(name="Added live", type="group")
|
|
communicator = await ws_communicator(user=user)
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
|
|
await sync_to_async(ChatMembership.objects.create)(chat=chat, user=user)
|
|
access_event = await communicator.receive_json_from()
|
|
assert access_event == {
|
|
"type": "chat_access_changed",
|
|
"chat_id": chat.id,
|
|
"can_access": True,
|
|
}
|
|
|
|
await get_channel_layer().group_send(
|
|
f"chat_{chat.id}",
|
|
{
|
|
"type": "chat_message",
|
|
"chat_id": chat.id,
|
|
"payload": {"id": 3, "content": "after grant"},
|
|
},
|
|
)
|
|
update = await communicator.receive_json_from()
|
|
assert update["type"] == "list_update"
|
|
await communicator.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_message_honors_membership_permission(ws_communicator, user):
|
|
chat = await sync_to_async(Chat.objects.create)(name="Read only", type="group")
|
|
await sync_to_async(ChatMembership.objects.create)(
|
|
chat=chat,
|
|
user=user,
|
|
role="guest",
|
|
)
|
|
communicator = await ws_communicator(user=user)
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
await communicator.send_json_to({"action": "open_chat", "chat_id": chat.id})
|
|
assert (await communicator.receive_json_from())["type"] == "chat_opened"
|
|
|
|
await communicator.send_json_to({"action": "send_message", "content": "denied"})
|
|
response = await communicator.receive_json_from()
|
|
assert response["code"] == "access_denied"
|
|
message_exists = await sync_to_async(
|
|
Message.objects.filter(chat=chat, content="denied").exists
|
|
)()
|
|
assert not message_exists
|
|
await communicator.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_command_does_not_terminate_connection(ws_communicator, user):
|
|
communicator = await ws_communicator(user=user)
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
|
|
await communicator.send_json_to({"action": "open_chat", "chat_id": "invalid"})
|
|
first_error = await communicator.receive_json_from()
|
|
assert first_error["code"] == "invalid_field"
|
|
|
|
await communicator.send_json_to({"action": "not_supported"})
|
|
second_error = await communicator.receive_json_from()
|
|
assert second_error["code"] == "unknown_action"
|
|
|
|
await communicator.send_to(text_data="{not-json")
|
|
malformed_error = await communicator.receive_json_from()
|
|
assert malformed_error["code"] == "invalid_json"
|
|
|
|
await communicator.send_to(bytes_data=b"binary")
|
|
binary_error = await communicator.receive_json_from()
|
|
assert binary_error["code"] == "binary_not_supported"
|
|
await communicator.disconnect()
|
|
|
|
|
|
def test_websocket_ticket_is_single_use(user):
|
|
ticket = issue_websocket_ticket(user=user, session_id=None)
|
|
identity = consume_websocket_ticket(ticket)
|
|
assert identity.user.pk == user.pk
|
|
|
|
with pytest.raises(WebSocketTicketError, match="ticket_already_used"):
|
|
consume_websocket_ticket(ticket)
|
|
|
|
|
|
def test_websocket_ticket_expires(user):
|
|
ticket = issue_websocket_ticket(user=user, session_id=None)
|
|
future = time.time() + 31
|
|
with (
|
|
patch("django.core.signing.time.time", return_value=future),
|
|
pytest.raises(WebSocketTicketError, match="invalid_or_expired_ticket"),
|
|
):
|
|
consume_websocket_ticket(ticket)
|
|
|
|
|
|
def test_websocket_ticket_rejects_revoked_session(user):
|
|
session = UserAuthSession.objects.create(user=user, refresh_token_hash="stub")
|
|
ticket = issue_websocket_ticket(user=user, session_id=str(session.session_id))
|
|
session.revoke(reason="test")
|
|
with pytest.raises(WebSocketTicketError, match="session_revoked"):
|
|
consume_websocket_ticket(ticket)
|
|
|
|
|
|
def test_websocket_ticket_endpoint_requires_authentication(
|
|
api_client,
|
|
auth_client_factory,
|
|
user,
|
|
):
|
|
anonymous_response = api_client.post("/api/auth/ws-ticket/")
|
|
assert anonymous_response.status_code == 401
|
|
|
|
response = auth_client_factory(user).post("/api/auth/ws-ticket/")
|
|
assert response.status_code == 200
|
|
assert response.data["expires_in"] == 30
|
|
assert response.data["ticket"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_untrusted_websocket_origin_is_rejected(user):
|
|
ticket = await sync_to_async(issue_websocket_ticket)(user=user, session_id=None)
|
|
communicator = WebsocketCommunicator(
|
|
application,
|
|
f"/ws/?ticket={ticket}",
|
|
headers=[(b"origin", b"https://attacker.example")],
|
|
)
|
|
connected, _ = await communicator.connect()
|
|
assert connected is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ticket_connects_from_trusted_origin(user):
|
|
ticket = await sync_to_async(issue_websocket_ticket)(user=user, session_id=None)
|
|
communicator = WebsocketCommunicator(
|
|
application,
|
|
f"/ws/?ticket={ticket}",
|
|
headers=[(b"origin", b"http://localhost:3000")],
|
|
)
|
|
connected, _ = await communicator.connect()
|
|
assert connected is True
|
|
await communicator.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_matching_session_revocation_closes_live_socket(ws_communicator, user):
|
|
communicator = await ws_communicator(user=user)
|
|
communicator.scope["auth_session_id"] = "session-to-revoke"
|
|
connected, _ = await communicator.connect()
|
|
assert connected
|
|
|
|
await get_channel_layer().group_send(
|
|
f"user_{user.id}",
|
|
{
|
|
"type": "auth_session_revoked",
|
|
"session_id": "session-to-revoke",
|
|
},
|
|
)
|
|
close_event = await communicator.receive_output()
|
|
assert close_event == {"type": "websocket.close", "code": 4401}
|