293 lines
10 KiB
Python
293 lines
10 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
from websockets import connect as ws_connect
|
|
from websockets import serve
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
from quantumpipe.exceptions import ProtocolError
|
|
from quantumpipe.sansio import (
|
|
_CREDENTIAL,
|
|
Message,
|
|
Verification,
|
|
generate_registration_secret,
|
|
)
|
|
from quantumpipe.ws import TEXT, SecSocket
|
|
|
|
AD = b"test app"
|
|
|
|
|
|
class FakeWS:
|
|
"""websockets-library style endpoint (send/recv) over memory queues."""
|
|
|
|
def __init__(self, inbox, outbox):
|
|
self._inbox = inbox
|
|
self._outbox = outbox
|
|
|
|
async def send(self, data):
|
|
self._outbox.put_nowait(data)
|
|
|
|
async def recv(self):
|
|
data = await self._inbox.get()
|
|
if data is None: # the peer closed
|
|
raise ConnectionClosed(None, None)
|
|
return data
|
|
|
|
async def close(self):
|
|
self._outbox.put_nowait(None) # wake the peer like a real close
|
|
|
|
|
|
class FakeStarletteWS(FakeWS):
|
|
"""FastAPI/Starlette style endpoint (accept/send_bytes/receive_bytes)."""
|
|
|
|
accepted = False
|
|
|
|
async def accept(self):
|
|
self.accepted = True
|
|
|
|
async def send_bytes(self, data):
|
|
self._outbox.put_nowait(data)
|
|
|
|
async def receive_bytes(self):
|
|
return await self.recv()
|
|
|
|
|
|
def ws_pair(starlette=False):
|
|
"""A connected pair of in-memory WebSocket endpoints (srv, cli)."""
|
|
to_srv, to_cli = asyncio.Queue(), asyncio.Queue()
|
|
cls = FakeStarletteWS if starlette else FakeWS
|
|
return cls(to_srv, to_cli), cls(to_cli, to_srv)
|
|
|
|
|
|
async def socket_pair(
|
|
cred_srv, cred_cli, verify_srv=None, verify_cli=None, ad=AD, starlette=False
|
|
):
|
|
"""Two SecSockets handshake over a fake ws pair; returns (srv, cli)."""
|
|
ws_srv, ws_cli = ws_pair(starlette)
|
|
srv = SecSocket(
|
|
ws_srv,
|
|
{"x": cred_srv},
|
|
ad=ad,
|
|
handlers={Verification: verify_srv} if verify_srv else None,
|
|
)
|
|
cli = SecSocket(
|
|
ws_cli,
|
|
{"x": cred_cli},
|
|
"x",
|
|
ad=ad,
|
|
handlers={Verification: verify_cli} if verify_cli else None,
|
|
)
|
|
await asyncio.gather(srv.accept(), cli.connect())
|
|
return srv, cli
|
|
|
|
|
|
async def test_roundtrip():
|
|
psk = generate_registration_secret()
|
|
srv, cli = await socket_pair(psk, psk)
|
|
await cli.send(b"ping")
|
|
assert await srv.receive() == b"ping"
|
|
await srv.send(b"pong")
|
|
assert await cli.receive() == b"pong"
|
|
|
|
|
|
async def test_send_picks_frame_type():
|
|
# str goes as a TEXT frame and is delivered decoded; bytes are opaque
|
|
psk = generate_registration_secret()
|
|
srv, cli = await socket_pair(psk, psk)
|
|
await cli.send("hello")
|
|
assert await srv.receive() == "hello"
|
|
await srv.send(b"\x00\xff")
|
|
assert await cli.receive() == b"\x00\xff"
|
|
await cli.send("")
|
|
assert await srv.receive() == "" # empty text frame
|
|
await cli.send(b"")
|
|
assert await srv.receive() == b"" # empty binary frame
|
|
|
|
|
|
async def test_message_passthrough():
|
|
# applications can use their own kinds (protobuf etc.): a Message
|
|
# passes through the wrapper untouched
|
|
psk = generate_registration_secret()
|
|
srv, cli = await socket_pair(psk, psk)
|
|
await cli.send(Message(b"\x89fake-protobuf", 7))
|
|
assert await srv.receive() == Message(b"\x89fake-protobuf", 7)
|
|
|
|
|
|
async def test_starlette_flavor():
|
|
# the FastAPI/Starlette duck-type: accept + send_bytes/receive_bytes
|
|
psk = generate_registration_secret()
|
|
ws_srv, ws_cli = ws_pair(starlette=True)
|
|
srv = SecSocket(ws_srv, {"x": psk}, ad=AD)
|
|
cli = SecSocket(ws_cli, {"x": psk}, "x", ad=AD)
|
|
await asyncio.gather(srv.accept(), cli.connect())
|
|
assert ws_srv.accepted
|
|
await cli.send(b"ping")
|
|
assert await srv.receive() == b"ping"
|
|
|
|
|
|
async def test_protocol_kind_rejected():
|
|
# kind 0 is reserved for in-band protocol messages; none implemented
|
|
psk = generate_registration_secret()
|
|
srv, cli = await socket_pair(psk, psk)
|
|
with pytest.raises(ProtocolError, match="reserved"):
|
|
await cli.send(Message(b"x", 0))
|
|
frame = bytes(cli.session._tx.encrypt(bytes([0]) + b"x"))
|
|
await cli._ws.send(frame) # bypass the wrapper
|
|
with pytest.raises(ProtocolError, match="reserved"):
|
|
await srv.receive()
|
|
|
|
|
|
async def test_empty_frame_rejected():
|
|
# a box with no plaintext at all carries no kind byte
|
|
psk = generate_registration_secret()
|
|
srv, cli = await socket_pair(psk, psk)
|
|
await cli._ws.send(bytes(cli.session._tx.encrypt(b"")))
|
|
with pytest.raises(ProtocolError, match="empty"):
|
|
await srv.receive()
|
|
|
|
|
|
async def test_bad_utf8_text_frame_rejected():
|
|
psk = generate_registration_secret()
|
|
srv, cli = await socket_pair(psk, psk)
|
|
(frame,) = cli.session(Message(b"\xff\xff", TEXT))
|
|
await cli._ws.send(frame)
|
|
with pytest.raises(ProtocolError, match="utf-8"):
|
|
await srv.receive()
|
|
|
|
|
|
async def test_traffic_before_establishment_rejected():
|
|
psk = generate_registration_secret()
|
|
_ws_srv, ws_cli = ws_pair()
|
|
cli = SecSocket(ws_cli, {"x": psk}, "x", ad=AD)
|
|
with pytest.raises(ProtocolError, match="no traffic"):
|
|
await cli.send(b"early")
|
|
|
|
|
|
async def test_verification_callback():
|
|
psk = generate_registration_secret()
|
|
seen = []
|
|
|
|
async def skip(v: Verification): # -> None: skip the comparison
|
|
seen.append(v.sas)
|
|
|
|
srv, cli = await socket_pair(psk, psk, skip, skip)
|
|
cred_srv, cred_cli = srv.session.credential, cli.session.credential
|
|
assert len(seen) == 2 and seen[0] == seen[1]
|
|
# skipped comparison: established, but flagged unverified
|
|
assert not srv.session.verified and not cli.session.verified
|
|
assert cli.session.verification.sas == seen[1]
|
|
|
|
# the resumption flow never prompts: it is silent even while
|
|
# unverified; the stored digest is exposed and the parties confirm
|
|
# whenever they wish
|
|
seen.clear()
|
|
srv, cli = await socket_pair(cred_srv, cred_cli, skip, skip)
|
|
await cli.send(b"resumed")
|
|
assert await srv.receive() == b"resumed"
|
|
assert seen == [] # no comparison asked
|
|
assert not srv.session.verified and not cli.session.verified
|
|
assert cli.session.verification.sas == _CREDENTIAL.unpack(cred_cli)[3]
|
|
cli.session.confirm(True) # compared out of band: mark verified
|
|
srv.session.confirm(True)
|
|
assert cli.session.verified and srv.session.verified
|
|
cred_srv, cred_cli = srv.session.credential, cli.session.credential
|
|
|
|
# a verified credential resumes silently, flagged verified
|
|
srv, cli = await socket_pair(cred_srv, cred_cli, skip, skip)
|
|
await cli.send(b"again")
|
|
assert await srv.receive() == b"again"
|
|
assert seen == []
|
|
assert srv.session.verified and cli.session.verified
|
|
|
|
|
|
async def test_verification_rejected():
|
|
psk = generate_registration_secret()
|
|
|
|
async def reject(v: Verification):
|
|
return False
|
|
|
|
ws_srv, ws_cli = ws_pair()
|
|
srv_store = {"x": psk}
|
|
srv = SecSocket(ws_srv, srv_store, ad=AD)
|
|
cli_store = {"x": psk}
|
|
cli = SecSocket(ws_cli, cli_store, "x", ad=AD, handlers={Verification: reject})
|
|
srv_task = asyncio.create_task(srv.accept())
|
|
# the responder's default answer releases msg4, but the initiator's
|
|
# reject holds back msg5: her side fails, the responder never swaps
|
|
with pytest.raises(ProtocolError, match="rejected"):
|
|
await cli.connect()
|
|
assert cli.session.state == "failed"
|
|
assert cli_store == {} # rejected: the credential is deleted
|
|
with pytest.raises(ConnectionClosed):
|
|
await srv_task
|
|
assert srv.session.state == "reg:msg5" # parked: no msg5 ever came
|
|
assert srv_store == {"x": psk} # the registration key is intact
|
|
|
|
|
|
async def test_parallel_registration_race():
|
|
"""Two registrations racing on one registration key: the swap at the
|
|
comparison answer arbitrates -- the first answer (the app default
|
|
skip, right after msg3) wins the store slot, the loser's server
|
|
session fails with "registration key already spent". Both clients
|
|
still establish (the loser's composite is orphaned server-side)."""
|
|
psk = generate_registration_secret()
|
|
srv_store = {"x": psk}
|
|
outcomes = []
|
|
|
|
async def server_side(ws):
|
|
sock = SecSocket(ws, srv_store, ad=AD)
|
|
try:
|
|
await sock.accept()
|
|
except Exception as e: # noqa: BLE001 -- recorded and asserted below
|
|
outcomes.append(e)
|
|
else:
|
|
outcomes.append("established")
|
|
|
|
async def client_side():
|
|
ws_srv, ws_cli = ws_pair()
|
|
srv_task = asyncio.create_task(server_side(ws_srv))
|
|
store = {"x": psk}
|
|
sock = SecSocket(ws_cli, store, "x", ad=AD)
|
|
await sock.connect()
|
|
assert sock.session.state == "established"
|
|
return srv_task, store
|
|
|
|
(t1, s1), (t2, s2) = await asyncio.gather(client_side(), client_side())
|
|
await asyncio.gather(t1, t2)
|
|
|
|
# exactly one winner server-side; the loser is rejected at the swap
|
|
fails = [o for o in outcomes if isinstance(o, ProtocolError)]
|
|
assert outcomes.count("established") == 1
|
|
assert len(fails) == 1 and "registration key already spent" in str(fails[0])
|
|
# the winner's composite replaced the registration key; the loser's
|
|
# burn/delete machinery never touched the slot again
|
|
assert len(srv_store) == 1 and len(srv_store["x"]) == _CREDENTIAL.size
|
|
# both clients rotated their own stores and established
|
|
assert all(len(s["x"]) == _CREDENTIAL.size for s in (s1, s2))
|
|
|
|
|
|
async def test_real_websockets_roundtrip():
|
|
"""Integration over a real localhost websockets server/client."""
|
|
psk = generate_registration_secret()
|
|
srv_sock = None
|
|
|
|
async def handler(ws):
|
|
nonlocal srv_sock
|
|
srv_sock = SecSocket(ws, {"x": psk}, ad=AD)
|
|
await srv_sock.accept()
|
|
await ws.wait_closed()
|
|
|
|
server = await serve(handler, "127.0.0.1", 0)
|
|
port = server.sockets[0].getsockname()[1]
|
|
ws = await ws_connect(f"ws://127.0.0.1:{port}")
|
|
cli = SecSocket(ws, {"x": psk}, "x", ad=AD)
|
|
await cli.connect()
|
|
while srv_sock is None or srv_sock.session.state != "established":
|
|
await asyncio.sleep(0.01)
|
|
await cli.send("ping")
|
|
assert await srv_sock.receive() == "ping"
|
|
await srv_sock.send(b"pong")
|
|
assert await cli.receive() == b"pong"
|
|
await ws.close()
|
|
server.close()
|