from os import getpid
from socket import socket, AF_INET, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR
from logging import getLogger, DEBUG, StreamHandler
from threading import active_count
from time import monotonic
from typing import Any, Dict
from uuid import UUID
from edri.config.setting import (
SWITCH_HOST,
SWITCH_MAX_MESSAGE,
SWITCH_PORT,
SWITCH_STATUS_ENABLED,
SWITCH_STATUS_HOST,
SWITCH_STATUS_PORT,
is_loopback_only,
)
from edri.protocol.channel import SecureChannel
from edri.protocol.errors import ProtocolError
from edri.protocol.handshake import (
HANDSHAKE_TIMEOUT,
server_handshake,
server_handshake_unauthenticated,
)
from edri.protocol.keys import KeySet
from edri.switch.connection import Connection
from edri.switch import Forwarder, Receiver
from edri.switch.forwarder import ForwardMessage
from edri.switch.sender import Sender
from edri.switch.mailbox import Mailbox
from edri.switch.status import serve_status
from edri.utility.function import (
cpu_seconds_total,
current_rss_bytes,
open_fd_count,
peak_rss_bytes,
)
from queue import Queue
[docs]
class Switch:
def __init__(self) -> None:
self.logger = getLogger(__name__)
self.keyset = KeySet.from_settings()
if self.keyset is None:
if not is_loopback_only(SWITCH_HOST):
raise RuntimeError(
f"refusing to listen on {SWITCH_HOST} without authentication: "
"set EDRI_SWITCH_SECRET or EDRI_SWITCH_SECRET_FILE, or bind to loopback"
)
self.logger.warning(
"Starting without authentication on %s - every peer that reaches this "
"port is trusted. Development only; set EDRI_SWITCH_SECRET otherwise.",
SWITCH_HOST,
)
if SWITCH_STATUS_ENABLED and not is_loopback_only(SWITCH_STATUS_HOST):
if self.keyset is None:
raise RuntimeError(
f"refusing to publish the status endpoint on {SWITCH_STATUS_HOST} "
"without authentication: bind EDRI_SWITCH_STATUS_HOST to loopback, "
"set EDRI_SWITCH_SECRET, or set EDRI_SWITCH_STATUS_ENABLED=false"
)
# Beyond the spec's rule, and deliberately: the status endpoint has
# no key check of its own, so a secret makes the Switch safe while
# leaving this port readable by anyone who can reach it.
self.logger.warning(
"Status endpoint on %s is reachable off-host and has no authentication; "
"it publishes router ids and peer addresses.", SWITCH_STATUS_HOST,
)
self.socket = socket(AF_INET, SOCK_STREAM)
self.connections: Dict[UUID, Connection] = {}
self.mailboxes: Dict[UUID, Mailbox] = {}
self.forwarder_queue = Queue[ForwardMessage]()
self.forwarder: Forwarder = Forwarder(self.forwarder_queue, self.connections)
self.forwarder.start()
self._started_at = monotonic()
[docs]
def status(self) -> dict[str, Any]:
"""Read-only operational snapshot of the Switch, returned as JSON-able dict.
Mirrors the kind of information the Router exposes at
``/edri/health-check-status``: one record per connected router plus
process-level vitals. Reads live structures without locking — fine for
best-effort telemetry, where a momentarily stale count is acceptable.
"""
return {
"connections": [
{
"router_id": str(router_id),
"peer": _peername(connection.channel.socket if connection.channel else None),
"queue_depth": connection.queue.qsize(),
"has_demands": connection.demands is not None,
"demands_bytes": len(connection.demands) if connection.demands else 0,
}
for router_id, connection in self.connections.items()
],
"connections_total": len(self.connections),
"forwarder_queue_depth": self.forwarder_queue.qsize(),
"vitals": {
"rss_bytes": current_rss_bytes(),
"rss_peak_bytes": peak_rss_bytes(),
"cpu_seconds": cpu_seconds_total(),
"threads": active_count(),
"open_fds": open_fd_count(),
"pid": getpid(),
"uptime_seconds": monotonic() - self._started_at,
},
}
[docs]
def handle_new_connection(self, connection_socket: socket) -> None:
connection_socket.settimeout(HANDSHAKE_TIMEOUT)
channel = SecureChannel(connection_socket, SWITCH_MAX_MESSAGE)
if self.keyset is None:
# Loopback-only development mode; Task 8 guarantees no other bind
# reaches this point.
router_id, key_index = server_handshake_unauthenticated(channel)
else:
router_id, key_index = server_handshake(channel, self.keyset)
connection_socket.settimeout(None)
previous = self.connections.get(router_id)
if previous is not None:
# The newcomer proved key possession; the incumbent may be a
# half-open socket from a node that restarted, which cannot be told
# from a healthy one. Evicting keeps a restarted node able to
# reclaim its own identity. Its mailbox is kept (see below).
self.logger.warning(
"Router %s reconnected from %s, evicting the previous connection from %s",
router_id, _peername(connection_socket),
_peername(previous.channel.socket) if previous.channel else None,
)
previous.disconnect()
# The mailbox belongs to the stable router identity, not a TCP session.
# Its unacknowledged head and queued tail survive a reconnect unchanged.
mailbox = self.mailboxes.setdefault(router_id, Mailbox())
connection = Connection(mailbox, channel, key_index=key_index)
self.connections[router_id] = connection
self.logger.info("Connection for %s from %s", router_id, _peername(connection_socket))
Receiver(self.forwarder_queue, router_id, connection, self.connections).start()
Sender(connection, router_id).start()
for other_router_id, other_connection in list(self.connections.items()):
if other_router_id != router_id and other_connection.demands is not None:
mailbox.put(other_connection.demands)
[docs]
def run(self) -> None:
self.logger.info("Starting switch on %s:%s", SWITCH_HOST, SWITCH_PORT)
self.socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
self.socket.bind((SWITCH_HOST, SWITCH_PORT))
self.socket.listen()
if SWITCH_STATUS_ENABLED:
try:
serve_status(self, SWITCH_STATUS_HOST, SWITCH_STATUS_PORT, self.logger)
except OSError as e:
# A taken status port must not stop the Switch from forwarding.
self.logger.warning("Status endpoint unavailable: %s", e)
self.logger.info("Ready for new connections")
while True:
connection, address = self.socket.accept()
self.logger.debug("New incoming connection: %s", address)
try:
self.handle_new_connection(connection)
except (ProtocolError, ConnectionError, OSError) as e:
# A peer that fails the handshake - a wrong secret, garbage
# bytes, an oversized declared length, a port probe, a scanner -
# must not take the whole Switch down. ProtocolError is not an
# OSError, so it needs naming explicitly. Drop just that
# connection and keep serving.
self.logger.warning("Dropping bad incoming connection %s: %s", address, e)
try:
connection.close()
except OSError:
pass
def _peername(connection_socket: socket | None) -> str | None:
"""Best-effort ``host:port`` of a connected peer; ``None`` once closed."""
try:
host, port, *_ = connection_socket.getpeername() # type: ignore[union-attr]
return f"{host}:{port}"
except (OSError, AttributeError, ValueError):
# ValueError: an AF_UNIX peer (a socketpair in tests) has no host/port
# pair to unpack.
return None
if __name__ == "__main__":
stream_handler = StreamHandler()
logger = getLogger()
logger.setLevel(DEBUG)
logger.addHandler(stream_handler)
switch = Switch()
switch.run()