jmri_openlcb.datagrams

  1"""OpenLCB Datagram Protocol (per S-9.7.4.1).
  2
  3Datagrams are variable-length messages up to 72 bytes carried across
  41..N CAN frames. Every datagram has one *source* node and one
  5*destination* node; each direction must be acknowledged by a
  6"Datagram Received OK" reply or rejected with "Datagram Rejected".
  7
  8CAN header layout (29-bit ID, high bits first) for datagram frames::
  9
 10    bit 28       = 1  (extended CAN frame)
 11    bits 24..27  = frame kind:
 12                     0x2  Complete datagram in one frame
 13                     0x3  First frame of a multi-frame datagram
 14                     0x4  Middle frame(s)
 15                     0x5  Final frame
 16    bits 12..23  = destination node alias
 17    bits 0..11   = source node alias
 18
 19Reply MTIs use the *global-messages* framing (bits 24..27 = 0x1,
 20bits 12..23 = 12-bit MTI). The two reply MTIs we care about are:
 21
 22    0x4C48  Datagram Received OK
 23    0x4E48  Datagram Rejected
 24
 25The engine below reads inbound frames from an
 26:class:`~jmri_openlcb.traffic.OlcbTrafficController` via its
 27``frames`` :class:`~jmri_core.events.Signal`, reassembles them, and
 28either exposes the resulting :class:`Datagram` to callers via
 29:meth:`DatagramEngine.wait_for` or dispatches to a registered
 30handler. Send side is symmetric: split the payload into 8-byte
 31chunks, emit the right frame kinds, and (optionally) wait for the
 32peer's Datagram Received OK reply.
 33"""
 34from __future__ import annotations
 35
 36import asyncio
 37import logging
 38from collections.abc import Awaitable, Callable
 39from dataclasses import dataclass, field
 40from typing import TYPE_CHECKING
 41
 42from jmri_core.events import Signal
 43
 44from jmri_openlcb.codec import OlcbFrame, encode_frame
 45
 46if TYPE_CHECKING:
 47    from jmri_openlcb.traffic import OlcbTrafficController
 48
 49_log = logging.getLogger(__name__)
 50
 51# ---- header helpers ----------------------------------------------------
 52
 53DATAGRAM_COMPLETE = 0x2
 54DATAGRAM_FIRST = 0x3
 55DATAGRAM_MIDDLE = 0x4
 56DATAGRAM_FINAL = 0x5
 57
 58# Global-message MTIs used for datagram replies (per S-9.7.4.1).
 59MTI_DATAGRAM_RECEIVED_OK = 0xA28
 60MTI_DATAGRAM_REJECTED = 0xA48
 61
 62# Max datagram body size — the standard caps at 72 bytes (9 frames of 8).
 63MAX_DATAGRAM_BYTES = 72
 64
 65
 66def _extract_frame_kind(header: int) -> int:
 67    """Bits 24..27 identify the frame kind."""
 68    return (header >> 24) & 0xF
 69
 70
 71def _extract_dest_alias(header: int) -> int:
 72    return (header >> 12) & 0xFFF
 73
 74
 75def _extract_source_alias(header: int) -> int:
 76    return header & 0xFFF
 77
 78
 79def _extract_mti(header: int) -> int:
 80    """The 12-bit MTI (bits 12..23), used for message-type frames.
 81
 82    Only meaningful when the frame kind (bits 24..27) is 0x1
 83    (OpenLCB global/addressed message). For datagram-kind frames the
 84    same bit range holds the destination alias, not an MTI.
 85    """
 86    return (header >> 12) & 0xFFF
 87
 88
 89def make_datagram_header(
 90    kind: int, dest_alias: int, source_alias: int,
 91) -> int:
 92    if kind not in (
 93        DATAGRAM_COMPLETE, DATAGRAM_FIRST,
 94        DATAGRAM_MIDDLE, DATAGRAM_FINAL,
 95    ):
 96        raise ValueError(f"invalid datagram frame kind: {kind:#x}")
 97    if not 0 <= dest_alias <= 0xFFF:
 98        raise ValueError(f"dest alias must be 12-bit: {dest_alias:#x}")
 99    if not 0 <= source_alias <= 0xFFF:
100        raise ValueError(f"source alias must be 12-bit: {source_alias:#x}")
101    return (
102        (1 << 28)
103        | (kind << 24)
104        | (dest_alias << 12)
105        | source_alias
106    )
107
108
109def make_reply_header(
110    mti: int, source_alias: int,
111) -> int:
112    """CAN header for a global/addressed message frame (frame kind 1).
113
114    For addressed messages, the destination alias travels in the
115    *data* payload (bytes 0..1) — the header itself only carries
116    the MTI + source alias.
117    """
118    return (
119        (1 << 28)
120        | (0x1 << 24)
121        | (mti << 12)
122        | source_alias
123    )
124
125
126def _encode_dest_prefix(dest_alias: int) -> bytes:
127    """The 2-byte destination-alias prefix that leads every addressed
128    message's data field."""
129    return bytes([(dest_alias >> 8) & 0xFF, dest_alias & 0xFF])
130
131
132def _decode_dest_alias_from_body(data: bytes) -> int | None:
133    """Peel the destination alias off the first two bytes of an
134    addressed-message body. Returns None when the frame doesn't
135    carry one (short frame, corrupted header)."""
136    if len(data) < 2:
137        return None
138    return ((data[0] & 0x0F) << 8) | data[1]
139
140
141# ---- dataclasses --------------------------------------------------------
142
143
144@dataclass(frozen=True, slots=True)
145class Datagram:
146    """A fully-reassembled datagram — the body plus the two node aliases."""
147
148    source_alias: int
149    dest_alias: int
150    data: bytes
151
152
153@dataclass(frozen=True, slots=True)
154class DatagramRejected(Exception):
155    """Raised when the peer replies with ``Datagram Rejected`` to a
156    :meth:`DatagramEngine.send` call. The 2-byte error code is
157    exposed so callers can distinguish "permanent" (bit 15 set)
158    from "temporary" and log accordingly."""
159
160    error_code: int
161
162    def __str__(self) -> str:  # pragma: no cover — trivial
163        return f"datagram rejected (code {self.error_code:#06x})"
164
165
166# ---- engine -------------------------------------------------------------
167
168
169@dataclass
170class _Pending:
171    """One in-flight outgoing datagram, waiting for the peer's reply."""
172
173    fut: asyncio.Future[None]
174    # The engine assumes one outstanding datagram per (src → dst)
175    # pair; the standard allows more but no known real node needs it
176    # and the extra bookkeeping is fragile.
177    dest_alias: int
178
179
180class DatagramEngine:
181    """Reassembles inbound datagrams, sends outbound ones with
182    fragmentation, and pairs sends with the peer's ``Received OK`` /
183    ``Rejected`` reply.
184
185    The reader hooks into an :class:`OlcbTrafficController`'s frames
186    signal via :meth:`attach`; the caller is responsible for
187    :meth:`detach` (or calling ``aclose``) at shutdown so the signal
188    listener doesn't outlive the engine.
189    """
190
191    def __init__(
192        self, tc: OlcbTrafficController,
193    ) -> None:
194        self._tc = tc
195        self._unsubscribe: Callable[[], None] | None = None
196        # Per-peer reassembly buffers keyed by source alias.
197        self._partial: dict[int, bytearray] = {}
198        # Pending sends keyed by destination alias.
199        self._pending: dict[int, _Pending] = {}
200        # Fires with every fully-received *inbound* datagram after the
201        # engine has already sent Received OK to the peer.
202        self.inbound: Signal[Datagram] = Signal()
203
204    async def attach(self) -> None:
205        if self._unsubscribe is None:
206            self._unsubscribe = self._tc.frames.connect(self._on_frame)
207
208    async def detach(self) -> None:
209        if self._unsubscribe is not None:
210            self._unsubscribe()
211            self._unsubscribe = None
212        # Fail any outstanding sends so callers don't hang forever.
213        for pending in list(self._pending.values()):
214            if not pending.fut.done():
215                pending.fut.set_exception(
216                    RuntimeError("datagram engine closed"),
217                )
218        self._pending.clear()
219        self._partial.clear()
220
221    # ---- send -----------------------------------------------------------
222
223    async def send(
224        self,
225        dest_alias: int,
226        data: bytes,
227        *,
228        timeout: float = 3.0,
229        wait_for_reply: bool = True,
230    ) -> None:
231        """Send *data* to *dest_alias*. Awaits the peer's Datagram
232        Received OK when ``wait_for_reply`` is True (default). Raises
233        :class:`DatagramRejected` on Rejected, :class:`TimeoutError`
234        on no reply, :class:`ValueError` on bad inputs.
235
236        Data over 8 bytes is fragmented into First / Middle... / Final
237        frames per S-9.7.4.1. Zero-byte payloads are permitted and
238        travel as a single Complete frame with an empty data field.
239        """
240        if not 0 <= dest_alias <= 0xFFF:
241            raise ValueError(f"dest alias must be 12-bit: {dest_alias:#x}")
242        if len(data) > MAX_DATAGRAM_BYTES:
243            raise ValueError(
244                f"datagram body {len(data)} exceeds "
245                f"{MAX_DATAGRAM_BYTES}-byte cap",
246            )
247        if dest_alias in self._pending:
248            raise RuntimeError(
249                f"a datagram to alias {dest_alias:#x} is already in flight",
250            )
251
252        source_alias = self._tc.source_alias
253
254        # Arm the reply-waiter BEFORE sending so a very fast peer
255        # can't beat us to it. Skipped when the caller opts out.
256        pending: _Pending | None = None
257        if wait_for_reply:
258            loop = asyncio.get_event_loop()
259            pending = _Pending(fut=loop.create_future(), dest_alias=dest_alias)
260            self._pending[dest_alias] = pending
261
262        try:
263            await self._transmit(source_alias, dest_alias, data)
264            if pending is not None:
265                await asyncio.wait_for(pending.fut, timeout=timeout)
266        finally:
267            if pending is not None:
268                self._pending.pop(dest_alias, None)
269
270    async def _transmit(
271        self, source_alias: int, dest_alias: int, data: bytes,
272    ) -> None:
273        """Write out the datagram frames. Single Complete when body ≤ 8,
274        First + N × Middle + Final otherwise."""
275        if len(data) <= 8:
276            # Zero-length datagrams still go on the wire as a Complete
277            # frame — the protocol permits them even if no real
278            # sub-protocol uses them.
279            frame = OlcbFrame(
280                header=make_datagram_header(
281                    DATAGRAM_COMPLETE, dest_alias, source_alias,
282                ),
283                data=data,
284            )
285            await self._tc.send_frame(frame)
286            return
287
288        # Multi-frame: First, then Middle...Middle, then Final.
289        chunks: list[bytes] = [
290            data[i:i + 8] for i in range(0, len(data), 8)
291        ]
292        for i, chunk in enumerate(chunks):
293            if i == 0:
294                kind = DATAGRAM_FIRST
295            elif i == len(chunks) - 1:
296                kind = DATAGRAM_FINAL
297            else:
298                kind = DATAGRAM_MIDDLE
299            frame = OlcbFrame(
300                header=make_datagram_header(kind, dest_alias, source_alias),
301                data=chunk,
302            )
303            await self._tc.send_frame(frame)
304
305    # ---- receive --------------------------------------------------------
306
307    async def _on_frame(self, frame: OlcbFrame) -> None:
308        kind = _extract_frame_kind(frame.header)
309        if kind in (
310            DATAGRAM_COMPLETE, DATAGRAM_FIRST,
311            DATAGRAM_MIDDLE, DATAGRAM_FINAL,
312        ):
313            await self._handle_datagram_frame(kind, frame)
314            return
315        # Not a datagram frame — is it a reply to something WE sent?
316        if kind == 0x1:
317            mti = _extract_mti(frame.header)
318            if mti in (MTI_DATAGRAM_RECEIVED_OK, MTI_DATAGRAM_REJECTED):
319                await self._handle_reply(mti, frame)
320
321    async def _handle_datagram_frame(
322        self, kind: int, frame: OlcbFrame,
323    ) -> None:
324        dest_alias = _extract_dest_alias(frame.header)
325        source_alias = _extract_source_alias(frame.header)
326
327        # We accept datagrams only when they're addressed to our own
328        # alias — the same rule real hardware follows. Everything else
329        # is on the bus but not for us.
330        if dest_alias != self._tc.source_alias:
331            return
332
333        if kind == DATAGRAM_COMPLETE:
334            body = bytes(frame.data)
335            await self._deliver(source_alias, dest_alias, body)
336            return
337
338        buf = self._partial.setdefault(source_alias, bytearray())
339        if kind == DATAGRAM_FIRST:
340            # Reset any dangling partial (peer aborted a previous send).
341            buf.clear()
342            buf.extend(frame.data)
343            return
344        if kind == DATAGRAM_MIDDLE:
345            buf.extend(frame.data)
346            return
347        if kind == DATAGRAM_FINAL:
348            buf.extend(frame.data)
349            body = bytes(buf)
350            self._partial.pop(source_alias, None)
351            await self._deliver(source_alias, dest_alias, body)
352
353    async def _deliver(
354        self, source_alias: int, dest_alias: int, body: bytes,
355    ) -> None:
356        # Send Received OK before firing our own signal so the reply
357        # is on the wire even if a subscriber raises.
358        await self._send_received_ok(source_alias)
359        dg = Datagram(
360            source_alias=source_alias,
361            dest_alias=dest_alias,
362            data=body,
363        )
364        await self.inbound.emit(dg)
365
366    async def _send_received_ok(self, dest_alias: int) -> None:
367        frame = OlcbFrame(
368            header=make_reply_header(
369                MTI_DATAGRAM_RECEIVED_OK,
370                self._tc.source_alias,
371            ),
372            # Addressed messages carry the destination alias in the
373            # first two body bytes, followed by protocol-specific
374            # fields (here: a single 0x00 = "no reply pending").
375            data=_encode_dest_prefix(dest_alias) + b"\x00",
376        )
377        await self._tc.send_frame(frame)
378
379    async def _handle_reply(self, mti: int, frame: OlcbFrame) -> None:
380        # The reply's *source* is the peer we sent to. Match on that.
381        peer_alias = _extract_source_alias(frame.header)
382        # For addressed messages, the destination alias travels in
383        # the first 2 body bytes, not the header.
384        addressed_dest = _decode_dest_alias_from_body(frame.data)
385        if addressed_dest != self._tc.source_alias:
386            return
387        pending = self._pending.get(peer_alias)
388        if pending is None or pending.fut.done():
389            return
390        # The status/error bytes follow the 2-byte destination
391        # prefix. Received OK carries one byte (0x00 or reply-pending
392        # flag); Rejected carries the 2-byte error code.
393        body = frame.data[2:]
394        if mti == MTI_DATAGRAM_RECEIVED_OK:
395            pending.fut.set_result(None)
396        else:  # MTI_DATAGRAM_REJECTED
397            code = 0
398            if len(body) >= 2:
399                code = (body[0] << 8) | body[1]
400            pending.fut.set_exception(DatagramRejected(error_code=code))