robolib.netcode

 1"""
 2netcode.py - Defines the recv/encode/decode functions to interact with the msgpack packets in robosim
 3
 4Import level: 2
 5"""
 6from typing import Any
 7from dataclasses import dataclass
 8from enum import StrEnum
 9import socket
10
11from overrides import overrides
12from micronetcode import Codec, Client, Message, SourceType
13from microecs import EntityId
14import msgpack
15
16from robolib.utils import logger
17
18_SERIALIZABLE_TYPES = (str, int, float, bool, bytes, type(None))
19PACKET_MAXLEN = 100_000
20
21class RobosimClientState(StrEnum):
22    """The states in which a client can be"""
23    STAGED    = "staged"
24    CONNECTED = "connected"
25
26@dataclass
27class RobosimClient(Client):
28    """The robosim client: a regular micronetcode client + robosim client state"""
29    state: RobosimClientState = "staged"
30    robot_eid: EntityId | None = None
31
32class MsgpackCodec(Codec):
33    """Msgpack implementation for recv/encode and decode."""
34    def __init__(self, packet_maxlen: int | None = None):
35        self.packet_maxlen = packet_maxlen
36
37    @overrides
38    def recv(self, sock: socket.socket) -> bytes | None:
39        # First, read the packet length based on the first 4 bytes. Then conver this to an integer and read n bytes.
40        packet_len_in_bytes = sock.recv(4)
41        if len(packet_len_in_bytes) == 0:
42            return None
43        if len(packet_len_in_bytes) < 4:
44            logger.error(f"The unexpected happened: {sock}. Closing connection. TO FIX")
45            return None
46
47        try:
48            packet_len = int.from_bytes(packet_len_in_bytes, "big")
49        except Exception as e:
50            logger.error(e)
51            return None
52
53        if packet_len > PACKET_MAXLEN:
54            logger.debug(f"Client: {sock} sent a packet of {packet_len} > {PACKET_MAXLEN}. Closing connection")
55            return None
56
57        buf = bytearray(packet_len)
58        view = memoryview(buf)
59        received = 0
60        while received < packet_len:
61            chunk = sock.recv(packet_len - received)
62            if len(chunk) == 0:
63                return None
64            view[received: received + len(chunk)] = chunk
65            received += len(chunk)
66        return bytes(buf)
67
68    @overrides
69    def decode(self, data: bytes) -> list[Message]:
70        packet: dict = msgpack.unpackb(data, raw=False)
71        if not isinstance(packet, dict):
72            raise TypeError(f"packet is not dictionary: {type(packet)}")
73        return [Message(source_type=SourceType.SOCKET, data=packet)]
74
75    def encode(self, resp: Any) -> bytes:
76        if not isinstance(resp, dict):
77            raise TypeError(f"resp is not dictionary: {type(resp)}")
78
79        try:
80            packed_data = msgpack.packb(resp)
81        except TypeError as e:
82            _find_bad(resp)
83            raise e
84
85        # <SIZE-4b><packed_data>. We should look into avoiding this concatenation at each encode().
86        return len(packed_data).to_bytes(4, "big") + packed_data
87
88def _find_bad(d: Any, path: str="root"):
89    """ran in case of a TypeError in send_packet, so we know which part of the dict is not serializable"""
90    if isinstance(d, dict):
91        for k, v in d.items():
92            _find_bad(v, f"{path}[{k!r}]") # !r calls repr()
93    elif isinstance(d, (list, tuple)):
94        for i, v in enumerate(d):
95            _find_bad(v, f"{path}[{i}]")
96    elif not isinstance(d, _SERIALIZABLE_TYPES):
97        logger.error(f"Unserializable at {path}: {type(d).__name__}")
98    # the 'else' is implied because only 'good' types should be left here