"""One connection's protocol state: keys, sequence counters and the send lock."""
from dataclasses import dataclass
from logging import getLogger
from socket import SHUT_RDWR, socket
from threading import Lock
from edri.protocol.errors import ErrorReason, ProtocolError
from edri.protocol.frame import (
NONCE_LENGTH,
Frame,
FrameType,
decode_header,
encode_frame,
read_frame,
verify_frame,
)
__all__ = ["KeyedFrame", "SecureChannel"]
[docs]
@dataclass(frozen=True)
class KeyedFrame:
"""A received frame plus which key of a set authenticated it."""
frame: Frame
key: bytes
key_index: int
[docs]
class SecureChannel:
"""Frames in and out of one socket, authenticated per direction.
A channel starts **pre-key**, where it will accept exactly one unverified
``HELLO_SERVER`` — no key is agreed yet, and that frame is authenticated
retroactively because ``server_nonce`` sits inside the handshake transcript.
:meth:`activate` moves it to **established**, after which every frame is
verified and there is no path back short of a new connection.
"""
def __init__(self, sock: socket, max_message: int) -> None:
self.socket = sock
self.max_message = max_message
self.logger = getLogger(__name__)
self.send_key: bytes | None = None
self.receive_key: bytes | None = None
self.key_index = 0
self.send_seq = 0
self.recv_seq = 0
self.established = False
self._send_lock = Lock()
self._unkeyed_received = False
[docs]
def activate(self, send_key: bytes, receive_key: bytes, key_index: int = 0) -> None:
if self.established:
raise ProtocolError("channel is already established", ErrorReason.BAD_MAC)
self.send_key = send_key
self.receive_key = receive_key
self.key_index = key_index
self.established = True
[docs]
def send_unkeyed(self, frame_type: FrameType, target: bytes, payload: bytes) -> None:
"""Write the one frame that predates key agreement (``HELLO_SERVER``)."""
if self.established:
raise ProtocolError("unkeyed send after establishment", ErrorReason.BAD_MAC)
self._write(encode_frame(Frame(type=frame_type, seq=self.send_seq, target=target,
payload=payload), None))
# HELLO_SERVER is seq 0, so the HELLO_ACCEPT that follows it is seq 1.
self.send_seq += 1
[docs]
def receive_unkeyed(self) -> Frame:
"""Read the single permitted pre-key frame, with every field pinned.
The slot is consumed as soon as this call is allowed to proceed, not
only once the frame turns out to be a valid ``HELLO_SERVER`` - a
malformed or forged first attempt must burn the one chance the same
as a good one, so "at most one attempt" holds regardless of what a
caller does after a failure.
"""
if self.established or self._unkeyed_received:
raise ProtocolError("no further unkeyed frames are accepted", ErrorReason.BAD_MAC)
self._unkeyed_received = True
header, payload = read_frame(self.socket, self.max_message)
_, frame_type, flags, seq, _, target, _ = decode_header(header)
if frame_type != FrameType.HELLO_SERVER:
raise ProtocolError(f"expected HELLO_SERVER, got {frame_type.name}", ErrorReason.BAD_MAC)
if seq != 0:
raise ProtocolError(f"HELLO_SERVER must carry seq 0, got {seq}", ErrorReason.BAD_SEQ)
if len(payload) != NONCE_LENGTH:
raise ProtocolError(
f"HELLO_SERVER must carry a {NONCE_LENGTH} byte nonce, got {len(payload)}",
ErrorReason.BAD_MAC,
)
self.recv_seq = 1
return Frame(type=frame_type, seq=seq, target=target, payload=payload, flags=flags)
[docs]
def send(self, frame_type: FrameType, target: bytes, payload: bytes) -> None:
if not self.established:
raise ProtocolError("send before the channel was established", ErrorReason.BAD_MAC)
# The lock spans seq allocation and the write, so the numbers a peer sees
# are gap-free even with several threads sending (the Connector's shape).
with self._send_lock:
raw = encode_frame(
Frame(type=frame_type, seq=self.send_seq, target=target, payload=payload),
self.send_key,
)
self._write_locked(raw)
self.send_seq += 1
[docs]
def receive(self) -> Frame:
if not self.established:
raise ProtocolError("receive before the channel was established", ErrorReason.BAD_MAC)
header, payload = read_frame(self.socket, self.max_message)
assert self.receive_key is not None # established implies both keys
verify_frame(header, payload, self.receive_key)
_, frame_type, flags, seq, _, target, _ = decode_header(header)
if seq != self.recv_seq:
raise ProtocolError(f"expected seq {self.recv_seq}, got {seq}", ErrorReason.BAD_SEQ)
self.recv_seq += 1
return Frame(type=frame_type, seq=seq, target=target, payload=payload, flags=flags)
[docs]
def rekey(self, send_key: bytes, receive_key: bytes, key_index: int = 0) -> None:
"""Swap in the derived session keys once the handshake agrees on them.
Sequence counters deliberately survive: the handshake frames already
consumed 0 (and 1 downstream), and restarting the count would let a
captured handshake frame be replayed as traffic.
"""
if not self.established:
raise ProtocolError("rekey before the channel was established", ErrorReason.BAD_MAC)
self.send_key = send_key
self.receive_key = receive_key
self.key_index = key_index
[docs]
def receive_any_key(self, keys: list[bytes]) -> "KeyedFrame":
"""Read one frame, trying each key in turn.
Only the server's HELLO_CLIENT needs this: during a rotation the Switch
holds several keys and cannot know which one the client used until the
MAC verifies.
"""
if not self.established:
raise ProtocolError("receive before the channel was established", ErrorReason.BAD_MAC)
header, payload = read_frame(self.socket, self.max_message)
for key_index, key in enumerate(keys):
try:
verify_frame(header, payload, key)
except ProtocolError:
continue
break
else:
raise ProtocolError("frame matched none of the configured keys", ErrorReason.BAD_MAC)
_, frame_type, flags, seq, _, target, _ = decode_header(header)
if seq != self.recv_seq:
raise ProtocolError(f"expected seq {self.recv_seq}, got {seq}", ErrorReason.BAD_SEQ)
self.recv_seq += 1
frame = Frame(type=frame_type, seq=seq, target=target, payload=payload, flags=flags)
return KeyedFrame(frame=frame, key=key, key_index=key_index)
def _write(self, raw: bytes) -> None:
with self._send_lock:
self._write_locked(raw)
def _write_locked(self, raw: bytes) -> None:
# sendall, never send: a partial write silently truncates a frame and
# desynchronizes the stream for good.
self.socket.sendall(raw)
[docs]
def close(self) -> None:
try:
self.socket.shutdown(SHUT_RDWR)
except OSError:
pass
try:
self.socket.close()
except OSError:
pass