Initial work on Channel

RNS/Channel.py

@@ -0,0 +1,350 @@
+from __future__ import annotations
+import collections
+import enum
+import threading
+import time
+from types import TracebackType
+from typing import Type, Callable, TypeVar
+import abc
+import contextlib
+import struct
+import RNS
+from abc import ABC, abstractmethod
+_TPacket = TypeVar("_TPacket")
+
+
+class ChannelOutletBase(ABC):
+    @abstractmethod
+    def send(self, raw: bytes) -> _TPacket:
+        raise NotImplemented()
+
+    @abstractmethod
+    def resend(self, packet: _TPacket) -> _TPacket:
+        raise NotImplemented()
+
+    @property
+    @abstractmethod
+    def mdu(self):
+        raise NotImplemented()
+
+    @property
+    @abstractmethod
+    def rtt(self):
+        raise NotImplemented()
+
+    @property
+    @abstractmethod
+    def is_usable(self):
+        raise NotImplemented()
+
+    @abstractmethod
+    def get_packet_state(self, packet: _TPacket) -> MessageState:
+        raise NotImplemented()
+
+    @abstractmethod
+    def timed_out(self):
+        raise NotImplemented()
+
+    @abstractmethod
+    def __str__(self):
+        raise NotImplemented()
+
+    @abstractmethod
+    def set_packet_timeout_callback(self, packet: _TPacket, callback: Callable[[_TPacket], None] | None,
+                                    timeout: float | None = None):
+        raise NotImplemented()
+
+    @abstractmethod
+    def set_packet_delivered_callback(self, packet: _TPacket, callback: Callable[[_TPacket], None] | None):
+        raise NotImplemented()
+
+    @abstractmethod
+    def get_packet_id(self, packet: _TPacket) -> any:
+        raise NotImplemented()
+
+
+class CEType(enum.IntEnum):
+    ME_NO_MSG_TYPE      = 0
+    ME_INVALID_MSG_TYPE = 1
+    ME_NOT_REGISTERED   = 2
+    ME_LINK_NOT_READY   = 3
+    ME_ALREADY_SENT     = 4
+    ME_TOO_BIG          = 5
+
+
+class ChannelException(Exception):
+    def __init__(self, ce_type: CEType, *args):
+        super().__init__(args)
+        self.type = ce_type
+
+
+class MessageState(enum.IntEnum):
+    MSGSTATE_NEW       = 0
+    MSGSTATE_SENT      = 1
+    MSGSTATE_DELIVERED = 2
+    MSGSTATE_FAILED    = 3
+
+
+class MessageBase(abc.ABC):
+    MSGTYPE = None
+
+    @abstractmethod
+    def pack(self) -> bytes:
+        raise NotImplemented()
+
+    @abstractmethod
+    def unpack(self, raw):
+        raise NotImplemented()
+
+
+class Envelope:
+    def unpack(self, message_factories: dict[int, Type]) -> MessageBase:
+        msgtype, self.sequence, length = struct.unpack(">HHH", self.raw[:6])
+        raw = self.raw[6:]
+        ctor = message_factories.get(msgtype, None)
+        if ctor is None:
+            raise ChannelException(CEType.ME_NOT_REGISTERED, f"Unable to find constructor for Channel MSGTYPE {hex(msgtype)}")
+        message = ctor()
+        message.unpack(raw)
+        return message
+
+    def pack(self) -> bytes:
+        if self.message.__class__.MSGTYPE is None:
+            raise ChannelException(CEType.ME_NO_MSG_TYPE, f"{self.message.__class__} lacks MSGTYPE")
+        data = self.message.pack()
+        self.raw = struct.pack(">HHH", self.message.MSGTYPE, self.sequence, len(data)) + data
+        return self.raw
+
+    def __init__(self, outlet: ChannelOutletBase, message: MessageBase = None, raw: bytes = None, sequence: int = None):
+        self.ts = time.time()
+        self.id = id(self)
+        self.message = message
+        self.raw = raw
+        self.packet: _TPacket = None
+        self.sequence = sequence
+        self.outlet = outlet
+        self.tries = 0
+        self.tracked = False
+
+
+class Channel(contextlib.AbstractContextManager):
+    def __init__(self, outlet: ChannelOutletBase):
+        self._outlet = outlet
+        self._lock = threading.RLock()
+        self._tx_ring: collections.deque[Envelope] = collections.deque()
+        self._rx_ring: collections.deque[Envelope] = collections.deque()
+        self._message_callback: Callable[[MessageBase], None] | None = None
+        self._next_sequence = 0
+        self._message_factories: dict[int, Type[MessageBase]] = self._get_msg_constructors()
+        self._max_tries = 5
+
+    def __enter__(self) -> Channel:
+        return self
+
+    def __exit__(self, __exc_type: Type[BaseException] | None, __exc_value: BaseException | None,
+                 __traceback: TracebackType | None) -> bool | None:
+        self.shutdown()
+        return False
+
+    @staticmethod
+    def _get_msg_constructors() -> (int, Type[MessageBase]):
+        subclass_tuples = []
+        for subclass in MessageBase.__subclasses__():
+            with contextlib.suppress(Exception):
+                subclass()  # verify constructor works with no arguments, needed for unpacking
+                subclass_tuples.append((subclass.MSGTYPE, subclass))
+        message_factories = dict(subclass_tuples)
+        return message_factories
+
+    def register_message_type(self, message_class: Type[MessageBase]):
+        if not issubclass(message_class, MessageBase):
+            raise ChannelException(CEType.ME_INVALID_MSG_TYPE, f"{message_class} is not a subclass of {MessageBase}.")
+        if message_class.MSGTYPE is None:
+            raise ChannelException(CEType.ME_INVALID_MSG_TYPE, f"{message_class} has invalid MSGTYPE class attribute.")
+        try:
+            message_class()
+        except Exception as ex:
+            raise ChannelException(CEType.ME_INVALID_MSG_TYPE,
+                                     f"{message_class} raised an exception when constructed with no arguments: {ex}")
+
+        self._message_factories[message_class.MSGTYPE] = message_class
+
+    def set_message_callback(self, callback: Callable[[MessageBase], None]):
+        self._message_callback = callback
+
+    def shutdown(self):
+        self.clear_rings()
+
+    def clear_rings(self):
+        with self._lock:
+            for envelope in self._tx_ring:
+                if envelope.packet is not None:
+                    self._outlet.set_packet_timeout_callback(envelope.packet, None)
+                    self._outlet.set_packet_delivered_callback(envelope.packet, None)
+            self._tx_ring.clear()
+            self._rx_ring.clear()
+
+    def emplace_envelope(self, envelope: Envelope, ring: collections.deque[Envelope]) -> bool:
+        with self._lock:
+            i = 0
+            for env in ring:
+                if env.sequence < envelope.sequence:
+                    ring.insert(i, envelope)
+                    return True
+                if env.sequence == envelope.sequence:
+                    RNS.log(f"Envelope: Emplacement of duplicate envelope sequence.", RNS.LOG_EXTREME)
+                    return False
+                i += 1
+            envelope.tracked = True
+            ring.append(envelope)
+            return True
+
+    def prune_rx_ring(self):
+        with self._lock:
+            # Implementation for fixed window = 1
+            stale = list(sorted(self._rx_ring, key=lambda env: env.sequence, reverse=True))[1:]
+            for env in stale:
+                env.tracked = False
+                self._rx_ring.remove(env)
+
+    def receive(self, raw: bytes):
+        try:
+            envelope = Envelope(outlet=self._outlet, raw=raw)
+            message = envelope.unpack(self._message_factories)
+            with self._lock:
+                is_new = self.emplace_envelope(envelope, self._rx_ring)
+                self.prune_rx_ring()
+            if not is_new:
+                RNS.log("Channel: Duplicate message received", RNS.LOG_DEBUG)
+                return
+            RNS.log(f"Message received: {message}", RNS.LOG_DEBUG)
+            if self._message_callback:
+                threading.Thread(target=self._message_callback, name="Message Callback", args=[message], daemon=True)\
+                    .start()
+        except Exception as ex:
+            RNS.log(f"Channel: Error receiving data: {ex}")
+
+    def is_ready_to_send(self) -> bool:
+        if not self._outlet.is_usable:
+            RNS.log("Channel: Link is not usable.", RNS.LOG_EXTREME)
+            return False
+
+        with self._lock:
+            for envelope in self._tx_ring:
+                if envelope.outlet == self._outlet and (not envelope.packet
+                                                        or self._outlet.get_packet_state(envelope.packet) == MessageState.MSGSTATE_SENT):
+                    RNS.log("Channel: Link has a pending message.", RNS.LOG_EXTREME)
+                    return False
+        return True
+
+    def _packet_tx_op(self, packet: _TPacket, op: Callable[[_TPacket], bool]):
+        with self._lock:
+            envelope = next(filter(lambda e: self._outlet.get_packet_id(e.packet) == self._outlet.get_packet_id(packet),
+                                   self._tx_ring), None)
+            if envelope and op(envelope):
+                envelope.tracked = False
+                if envelope in self._tx_ring:
+                    self._tx_ring.remove(envelope)
+                else:
+                    RNS.log("Channel: Envelope not found in TX ring", RNS.LOG_DEBUG)
+        if not envelope:
+            RNS.log("Channel: Spurious message received.", RNS.LOG_EXTREME)
+
+    def _packet_delivered(self, packet: _TPacket):
+        self._packet_tx_op(packet, lambda env: True)
+
+    def _packet_timeout(self, packet: _TPacket):
+        def retry_envelope(envelope: Envelope) -> bool:
+            if envelope.tries >= self._max_tries:
+                RNS.log("Channel: Retry count exceeded, tearing down Link.", RNS.LOG_ERROR)
+                self.shutdown()  # start on separate thread?
+                self._outlet.timed_out()
+                return True
+            envelope.tries += 1
+            self._outlet.resend(envelope.packet)
+            return False
+
+        self._packet_tx_op(packet, retry_envelope)
+
+    def send(self, message: MessageBase) -> Envelope:
+        envelope: Envelope | None = None
+        with self._lock:
+            if not self.is_ready_to_send():
+                raise ChannelException(CEType.ME_LINK_NOT_READY, f"Link is not ready")
+            envelope = Envelope(self._outlet, message=message, sequence=self._next_sequence)
+            self._next_sequence = (self._next_sequence + 1) % 0x10000
+            self.emplace_envelope(envelope, self._tx_ring)
+        if envelope is None:
+            raise BlockingIOError()
+
+        envelope.pack()
+        if len(envelope.raw) > self._outlet.mdu:
+            raise ChannelException(CEType.ME_TOO_BIG, f"Packed message too big for packet: {len(envelope.raw)} > {self._outlet.mdu}")
+        envelope.packet = self._outlet.send(envelope.raw)
+        envelope.tries += 1
+        self._outlet.set_packet_delivered_callback(envelope.packet, self._packet_delivered)
+        self._outlet.set_packet_timeout_callback(envelope.packet, self._packet_timeout)
+        return envelope
+
+
+class LinkChannelOutlet(ChannelOutletBase):
+    def __init__(self, link: RNS.Link):
+        self.link = link
+
+    def send(self, raw: bytes) -> RNS.Packet:
+        packet = RNS.Packet(self.link, raw, context=RNS.Packet.CHANNEL)
+        packet.send()
+        return packet
+
+    def resend(self, packet: RNS.Packet) -> RNS.Packet:
+        if not packet.resend():
+            RNS.log("Failed to resend packet", RNS.LOG_ERROR)
+        return packet
+
+    @property
+    def mdu(self):
+        return self.link.MDU
+
+    @property
+    def rtt(self):
+        return self.link.rtt
+
+    @property
+    def is_usable(self):
+        return True  # had issues looking at Link.status
+
+    def get_packet_state(self, packet: _TPacket) -> MessageState:
+        status = packet.receipt.get_status()
+        if status == RNS.PacketReceipt.SENT:
+            return MessageState.MSGSTATE_SENT
+        if status == RNS.PacketReceipt.DELIVERED:
+            return MessageState.MSGSTATE_DELIVERED
+        if status == RNS.PacketReceipt.FAILED:
+            return MessageState.MSGSTATE_FAILED
+        else:
+            raise Exception(f"Unexpected receipt state: {status}")
+
+    def timed_out(self):
+        self.link.teardown()
+
+    def __str__(self):
+        return f"{self.__class__.__name__}({self.link})"
+
+    def set_packet_timeout_callback(self, packet: RNS.Packet, callback: Callable[[RNS.Packet], None] | None,
+                                    timeout: float | None = None):
+        if timeout:
+            packet.receipt.set_timeout(timeout)
+
+        def inner(receipt: RNS.PacketReceipt):
+            callback(packet)
+
+        packet.receipt.set_timeout_callback(inner if callback else None)
+
+    def set_packet_delivered_callback(self, packet: RNS.Packet, callback: Callable[[RNS.Packet], None] | None):
+        def inner(receipt: RNS.PacketReceipt):
+            callback(packet)
+
+        packet.receipt.set_delivery_callback(inner if callback else None)
+
+    def get_packet_id(self, packet: RNS.Packet) -> any:
+        return packet.get_hash()