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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
|
import pytest
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.security.websocket import OriginValidator
from channels.testing import WebsocketCommunicator
@pytest.mark.django_db(transaction=True)
@pytest.mark.asyncio
async def test_origin_validator():
"""
Tests that OriginValidator correctly allows/denies connections.
"""
# Make our test application
application = OriginValidator(AsyncWebsocketConsumer(), ["allowed-domain.com"])
# Test a normal connection
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"http://allowed-domain.com")]
)
connected, _ = await communicator.connect()
assert connected
await communicator.disconnect()
# Test a bad connection
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"http://bad-domain.com")]
)
connected, _ = await communicator.connect()
assert not connected
await communicator.disconnect()
# Make our test application, bad pattern
application = OriginValidator(AsyncWebsocketConsumer(), ["*.allowed-domain.com"])
# Test a bad connection
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"http://allowed-domain.com")]
)
connected, _ = await communicator.connect()
assert not connected
await communicator.disconnect()
# Make our test application, good pattern
application = OriginValidator(AsyncWebsocketConsumer(), [".allowed-domain.com"])
# Test a normal connection
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"http://www.allowed-domain.com")]
)
connected, _ = await communicator.connect()
assert connected
await communicator.disconnect()
# Make our test application, with scheme://domain[:port] for http
application = OriginValidator(
AsyncWebsocketConsumer(), ["http://allowed-domain.com"]
)
# Test a normal connection
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"http://allowed-domain.com")]
)
connected, _ = await communicator.connect()
assert connected
await communicator.disconnect()
# Test a bad connection
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"https://bad-domain.com:443")]
)
connected, _ = await communicator.connect()
assert not connected
await communicator.disconnect()
# Make our test application, with all hosts allowed
application = OriginValidator(AsyncWebsocketConsumer(), ["*"])
# Test a connection without any headers
communicator = WebsocketCommunicator(application, "/", headers=[])
connected, _ = await communicator.connect()
assert connected
await communicator.disconnect()
# Make our test application, with no hosts allowed
application = OriginValidator(AsyncWebsocketConsumer(), [])
# Test a connection without any headers
communicator = WebsocketCommunicator(application, "/", headers=[])
connected, _ = await communicator.connect()
assert not connected
await communicator.disconnect()
# Test bug with subdomain and empty origin header
application = OriginValidator(AsyncWebsocketConsumer(), [".allowed-domain.com"])
communicator = WebsocketCommunicator(application, "/", headers=[(b"origin", b"")])
connected, _ = await communicator.connect()
assert not connected
await communicator.disconnect()
# Test bug with subdomain and invalid origin header
application = OriginValidator(AsyncWebsocketConsumer(), [".allowed-domain.com"])
communicator = WebsocketCommunicator(
application, "/", headers=[(b"origin", b"something-invalid")]
)
connected, _ = await communicator.connect()
assert not connected
await communicator.disconnect()
|