saf_train_link.py

Go to the documentation of this file.

"""SafTrainLink - Subset-037 SaF protection of a train sim's ERTMS traffic (the train = SaF initiator).
 
Python side of RBC_GP's train-link SaF termination (``ab_gp_saf_train.c``; root
``docs/architecture/SAFE_COMM_INTEGRATION.dox`` plan steps 2/3). Wraps either transport a train sim uses
(``DualHomedLink`` or ``SafeCommLink``) and keeps their ``ensure_connected()`` / ``send_all()`` / ``poll_recv()``
surface, so ``train_sim`` only swaps the object:
 
* An outgoing ``RBC_MSG_ERTMS_ENVELOPE`` (kind 22: kind, train_id, codec bytes, zero padding) becomes one Sa-DATA
  frame of that train's SaF connection, carried in an ``RBC_MSG_SAF_ENVELOPE`` (kind 28: kind, train_id, 2-byte
  big-endian frame length, frame, zero padding). The PDU is exactly the codec bytes (``l_message`` long).
* The first ERTMS send for a train starts its handshake (AU1); PDUs queue until AR arrives, then go out in order.
* An incoming kind-28 frame is verified; a Sa-DATA PDU is handed to the caller as an ordinary kind-22 frame, so the
  catalog decoder and everything above it are unchanged.
* Anything else (the flat P0/M136/TRAIN_CONNECT scheme) passes through untouched - the RBC side protects only the
  ERTMS envelope too.
* A handshake past ``t_connect_ms`` or a DATA connection silent past ``t_supervision_ms`` is dropped; the next ERTMS
  send starts a new handshake (this is also how a site failover is recovered: the newly ONLINE site has no session).
* A Sa-DISCONNECT from the RBC (it answers a Sa-DATA frame for a connection it no longer holds - e.g. its processes
  restarted) drops the session at once: the PDUs sent in the last ``resend_window_ms`` were swallowed by the dead
  session, so they are queued again and a new handshake starts immediately (a real train's application would redo
  whatever the lost messages carried, e.g. Start of Mission).
 
Keys: ``keys`` maps a train's ETCS identity (its nid_engine) to its 24-byte KMAC - the same test-only key file the
RBC reads through ``RTE_SAFECOMM_TRAIN_KEYS_FILE``.
"""
import collections
import time
 
from . import saf
 
ERTMS_KIND = 22
SAF_KIND = 28
_SAF_HEADER = 4
ETY_ENGINE = 2  # Subset-037-2 Table 8
 
 
def load_keys_file(path):
    """Parse ``<etcs_id> <48 hex>`` lines (``#`` comments, blank lines ignored) into {etcs_id: kmac}."""
    keys = {}
    with open(path, encoding="ascii") as fh:
        for lineno, raw in enumerate(fh, 1):
            line = raw.split("#", 1)[0].strip()
            if not line:
                continue
            parts = line.split()
            if len(parts) != 2 or len(parts[1]) != 2 * saf.KEY_LEN:
                raise ValueError(f"{path}:{lineno}: expected '<etcs_id> <48 hex chars>'")
            keys[int(parts[0], 0) & 0xFFFFFF] = bytes.fromhex(parts[1])
    return keys
 
 
def encode_envelope(train_id, frame, wire_size):
    if not frame or len(frame) > wire_size - _SAF_HEADER:
        raise ValueError(f"SaF frame of {len(frame)} bytes does not fit a {wire_size}-byte envelope")
    env = bytes([SAF_KIND, train_id & 0xFF]) + len(frame).to_bytes(2, "big") + bytes(frame)
    return env + bytes(wire_size - len(env))
 
 
def decode_envelope(env):
    """(train_id, frame) of a kind-28 envelope, or None if malformed."""
    if len(env) < _SAF_HEADER or env[0] != SAF_KIND:
        return None
    length = int.from_bytes(env[2:4], "big")
    if length == 0 or _SAF_HEADER + length > len(env):
        return None
    return env[1], bytes(env[_SAF_HEADER:_SAF_HEADER + length])
 
 
class SafTrainLink:
    def __init__(self, inner, log, rbc_etcs_id, keys, wire_size, *, t_connect_ms=10000, t_supervision_ms=20000,
                 clock=None, rng=None, resend_window_ms=3000):
        self._inner = inner
        self._log = log
        self._rbc_id = rbc_etcs_id & 0xFFFFFF
        self._keys = dict(keys)
        self._wire_size = wire_size
        self._t_connect_ms = t_connect_ms
        self._t_supervision_ms = t_supervision_ms
        self._clock = clock
        self._rng = rng
        self._sessions = {}   # train_id -> SafSession
        self._queues = {}     # train_id -> [pdu, ...] waiting for the handshake
        self._recent = {}     # train_id -> deque of (sent_ms, pdu) sent in DATA, for a resend after a Sa-DISCONNECT
        self._resend_window_ms = resend_window_ms
 
    def __getattr__(self, name):
        # sockets / disconnect / connected_sites and anything else belong to the wrapped transport
        return getattr(self._inner, name)
 
    # -- transport surface -------------------------------------------------------------------------------------
    def ensure_connected(self):
        self._inner.ensure_connected()
        self._supervise()
 
    def send_all(self, payload):
        if len(payload) < 5 or payload[0] != ERTMS_KIND:
            self._inner.send_all(payload)
            return
        tid = payload[1]
        codec = bytes(payload[2:])
        l_message = int.from_bytes(codec[1:3], "big")
        if l_message == 0 or l_message > len(codec):
            self._log.warning("[SAF] [Train %d] ERTMS frame with bad l_message=%d not sent", tid, l_message)
            return
        pdu = codec[:l_message]
        sess = self._sessions.get(tid)
        if sess is not None and sess.state == saf.DATA:
            self._send_pdu(tid, sess, pdu)
            return
        self._queues.setdefault(tid, []).append(pdu)
        if sess is None or sess.state == saf.IDLE:
            self._start_handshake(tid)
 
    def poll_recv(self, timeout=0.0):
        for site, data in self._inner.poll_recv(timeout):
            if not data or data[0] != SAF_KIND:
                yield site, data
                continue
            decoded = decode_envelope(data)
            if decoded is None:
                self._log.warning("[SAF] malformed SaF envelope from %s dropped", site)
                continue
            tid, frame = decoded
            pdu = self._on_frame(tid, frame)
            if pdu is not None:
                env = bytes([ERTMS_KIND, tid]) + pdu
                yield site, env + bytes(max(0, self._wire_size - len(env)))
 
    # -- internals ----------------------------------------------------------------------------------------------
    def _new_session(self, tid):
        kmac = self._keys.get(tid)
        if kmac is None:
            return None
        kwargs = {"ety": ETY_ENGINE, "t_connect_ms": self._t_connect_ms, "t_supervision_ms": self._t_supervision_ms}
        if self._clock is not None:
            kwargs["clock"] = self._clock
        if self._rng is not None:
            kwargs["rng"] = self._rng
        return saf.SafSession("initiator", tid, self._rbc_id, kmac, self._raw_send_unused, self._raw_recv_unused,
                              **kwargs)
 
    @staticmethod
    def _raw_send_unused(_frame):
        raise RuntimeError("frame-driven session: frames are sent by SafTrainLink")
 
    @staticmethod
    def _raw_recv_unused():
        raise TimeoutError
 
    def _send_frame(self, tid, frame):
        self._inner.send_all(encode_envelope(tid, frame, self._wire_size))
 
    def _now_ms(self):
        return self._clock() if self._clock is not None else int(time.monotonic() * 1000)
 
    def _send_pdu(self, tid, sess, pdu):
        self._send_frame(tid, sess.build_frame(pdu))
        sess.seq_tx = (sess.seq_tx + 1) & 0xFFFFFFFF
        recent = self._recent.setdefault(tid, collections.deque())
        now = self._now_ms()
        recent.append((now, pdu))
        while recent and now - recent[0][0] > self._resend_window_ms:
            recent.popleft()
 
    def _start_handshake(self, tid):
        sess = self._new_session(tid)
        if sess is None:
            self._log.error("[SAF] [Train %d] no KMAC for ETCS id %d - ERTMS traffic not sent", tid, tid)
            self._queues.pop(tid, None)
            return
        self._sessions[tid] = sess
        self._send_frame(tid, sess.initiator_start())
        self._log.info("[SAF] [Train %d] AU1 sent (rbc_etcs_id=%d)", tid, self._rbc_id)
 
    def _on_frame(self, tid, frame):
        sess = self._sessions.get(tid)
        if sess is None:
            self._log.warning("[SAF] [Train %d] SaF frame with no session dropped", tid)
            return None
        try:
            if sess.state in (saf.WAIT_AU2, saf.WAIT_AR):
                reply = sess.initiator_on_frame(frame)
                if reply is not None:
                    self._send_frame(tid, reply)
                if sess.state == saf.DATA:
                    self._log.info("[SAF] [Train %d] safe connection established", tid)
                    for pdu in self._queues.pop(tid, []):
                        self._send_pdu(tid, sess, pdu)
                return None
            return sess.accept_frame(frame)
        except saf.SafError as err:
            if err.status == "PEER_DISCONNECT":
                now = self._now_ms()
                lost = [pdu for sent, pdu in self._recent.pop(tid, ()) if now - sent <= self._resend_window_ms]
                self._queues[tid] = lost + self._queues.get(tid, [])
                self._log.info("[SAF] [Train %d] released by the RBC - reconnecting, %d recent message(s) queued "
                               "for resend", tid, len(lost))
            else:
                self._log.warning("[SAF] [Train %d] connection dropped: %s", tid, err)
            if self._queues.get(tid):
                self._start_handshake(tid)
            return None
 
    def _supervise(self):
        for tid, sess in list(self._sessions.items()):
            if sess.handshake_expired():
                self._log.warning("[SAF] [Train %d] handshake timed out - retrying", tid)
                sess.state = saf.IDLE
                if self._queues.get(tid):
                    self._start_handshake(tid)
            elif sess.state == saf.DATA:
                try:
                    sess.supervise()
                except saf.SafError:
                    self._log.warning("[SAF] [Train %d] supervision expired - link lost, will re-handshake", tid)