hat.drivers.cotp

22class TcpConnectionInfo(typing.NamedTuple):
23    name: str | None
24    local_addr: net.TcpAddress
25    local_tsel: int | None
26    remote_addr: net.TcpAddress
27    remote_tsel: int | None

30class UnixConnectionInfo(typing.NamedTuple):
31    name: str | None
32    addr: net.UnixAddress
33    local_tsel: int | None
34    remote_tsel: int | None

43async def connect(addr: net.StreamAddress,
44                  *,
45                  local_tsel: int | None = None,
46                  remote_tsel: int | None = None,
47                  cotp_receive_queue_size: int = 1024,
48                  cotp_send_queue_size: int = 1024,
49                  **kwargs
50                  ) -> 'Connection':
51    """Create new COTP connection
52
53    Additional arguments are passed directly to `hat.drivers.tpkt.connect`.
54
55    """
56    conn = await tpkt.connect(addr, **kwargs)
57
58    try:
59        cr_tpdu = common.CR(src=next(_next_srcs),
60                            cls=0,
61                            calling_tsel=local_tsel,
62                            called_tsel=remote_tsel,
63                            max_tpdu=2048,
64                            pref_max_tpdu=None)
65        cr_tpdu_bytes = encoder.encode(cr_tpdu)
66        await conn.send(cr_tpdu_bytes)
67
68        cc_tpdu_bytes = await conn.receive()
69        cc_tpdu = encoder.decode(memoryview(cc_tpdu_bytes))
70        _validate_connect_response(cr_tpdu, cc_tpdu)
71
72        max_tpdu = _calculate_max_tpdu(cr_tpdu, cc_tpdu)
73        calling_tsel, called_tsel = _get_tsels(cr_tpdu, cc_tpdu)
74
75        return Connection(conn, max_tpdu, calling_tsel, called_tsel,
76                          cotp_receive_queue_size, cotp_send_queue_size)
77
78    except BaseException:
79        await aio.uncancellable(conn.async_close())
80        raise

 83async def listen(connection_cb: ConnectionCb,
 84                 addr: net.StreamAddress = net.TcpAddress('0.0.0.0', 102),
 85                 *,
 86                 cotp_receive_queue_size: int = 1024,
 87                 cotp_send_queue_size: int = 1024,
 88                 **kwargs
 89                 ) -> 'Server':
 90    """Create new COTP listening server
 91
 92    Additional arguments are passed directly to `hat.drivers.tpkt.listen`.
 93
 94    """
 95    server = Server()
 96    server._connection_cb = connection_cb
 97    server._receive_queue_size = cotp_receive_queue_size
 98    server._send_queue_size = cotp_send_queue_size
 99    server._log = mlog
100
101    server._srv = await tpkt.listen(server._on_connection, addr, **kwargs)
102
103    server._log = _create_server_logger(server._srv.info)
104
105    return server

108class Server(aio.Resource):
109    """COTP listening server
110
111    For creation of new instance see `listen` coroutine.
112
113    """
114
115    @property
116    def async_group(self) -> aio.Group:
117        """Async group"""
118        return self._srv.async_group
119
120    @property
121    def info(self) -> net.ServerInfo:
122        """Server info"""
123        return self._srv.info
124
125    async def _on_connection(self, tpkt_conn):
126        try:
127            try:
128                cr_tpdu_bytes = await tpkt_conn.receive()
129                cr_tpdu = encoder.decode(memoryview(cr_tpdu_bytes))
130                _validate_connect_request(cr_tpdu)
131
132                cc_tpdu = common.CC(dst=cr_tpdu.src,
133                                    src=next(_next_srcs),
134                                    cls=0,
135                                    calling_tsel=cr_tpdu.calling_tsel,
136                                    called_tsel=cr_tpdu.called_tsel,
137                                    max_tpdu=_calculate_cc_max_tpdu(cr_tpdu),
138                                    pref_max_tpdu=None)
139                cc_tpdu_bytes = encoder.encode(cc_tpdu)
140                await tpkt_conn.send(cc_tpdu_bytes)
141
142                max_tpdu = _calculate_max_tpdu(cr_tpdu, cc_tpdu)
143                calling_tsel, called_tsel = _get_tsels(cr_tpdu, cc_tpdu)
144                conn = Connection(tpkt_conn, max_tpdu,
145                                  called_tsel, calling_tsel,
146                                  self._receive_queue_size,
147                                  self._send_queue_size)
148
149            except BaseException:
150                await aio.uncancellable(tpkt_conn.async_close())
151                raise
152
153            try:
154                await aio.call(self._connection_cb, conn)
155
156            except BaseException:
157                await aio.uncancellable(conn.async_close())
158                raise
159
160        except Exception as e:
161            self._log.error("error creating new incomming connection: %s",
162                            e, exc_info=e)

165class Connection(aio.Resource):
166    """COTP connection
167
168    For creation of new instance see `connect` or `listen`.
169
170    """
171
172    def __init__(self,
173                 conn: tpkt.Connection,
174                 max_tpdu: int,
175                 local_tsel: int | None,
176                 remote_tsel: int | None,
177                 receive_queue_size: int,
178                 send_queue_size: int):
179        self._conn = conn
180        self._max_tpdu = max_tpdu
181        self._loop = asyncio.get_running_loop()
182        self._info = _connection_info_from_net(info=conn.info,
183                                               local_tsel=local_tsel,
184                                               remote_tsel=remote_tsel)
185        self._receive_queue = aio.Queue(receive_queue_size)
186        self._send_queue = aio.Queue(send_queue_size)
187        self._log = _create_connection_logger(self._info)
188
189        self.async_group.spawn(self._receive_loop)
190        self.async_group.spawn(self._send_loop)
191
192    @property
193    def async_group(self) -> aio.Group:
194        """Async group"""
195        return self._conn.async_group
196
197    @property
198    def info(self) -> ConnectionInfo:
199        """Connection info"""
200        return self._info
201
202    async def receive(self) -> util.Bytes:
203        """Receive data"""
204        try:
205            return await self._receive_queue.get()
206
207        except aio.QueueClosedError:
208            raise ConnectionError()
209
210    async def send(self, data: util.Bytes):
211        """Send data"""
212        try:
213            await self._send_queue.put((data, None))
214
215        except aio.QueueClosedError:
216            raise ConnectionError()
217
218    async def drain(self):
219        """Drain output buffer"""
220        try:
221            future = self._loop.create_future()
222            await self._send_queue.put((None, future))
223            await future
224
225        except aio.QueueClosedError:
226            raise ConnectionError()
227
228    async def _receive_loop(self):
229        try:
230            data_queue = collections.deque()
231            while True:
232                tpdu_bytes = await self._conn.receive()
233                tpdu = encoder.decode(memoryview(tpdu_bytes))
234
235                if isinstance(tpdu, (common.DR, common.ER)):
236                    self._log.info("received disconnect request / error")
237                    break
238
239                if not isinstance(tpdu, common.DT):
240                    continue
241
242                data_queue.append(tpdu.data)
243
244                if not tpdu.eot:
245                    continue
246
247                data = bytes(itertools.chain.from_iterable(data_queue))
248                data_queue.clear()
249
250                await self._receive_queue.put(data)
251
252        except ConnectionError:
253            pass
254
255        except Exception as e:
256            self._log.error("receive loop error: %s", e, exc_info=e)
257
258        finally:
259            self.close()
260            self._receive_queue.close()
261
262    async def _send_loop(self):
263        future = None
264        try:
265            while True:
266                data, future = await self._send_queue.get()
267
268                if data is None:
269                    await self._conn.drain()
270
271                else:
272                    data = memoryview(data)
273                    max_size = self._max_tpdu - 3
274
275                    while data:
276                        single_data, data = data[:max_size], data[max_size:]
277
278                        tpdu = common.DT(eot=not data, data=single_data)
279                        tpdu_bytes = encoder.encode(tpdu)
280
281                        await self._conn.send(tpdu_bytes)
282
283                if future and not future.done():
284                    future.set_result(None)
285
286        except ConnectionError:
287            pass
288
289        except Exception as e:
290            self._log.error("send loop error: %s", e, exc_info=e)
291
292        finally:
293            self.close()
294            self._send_queue.close()
295
296            while True:
297                if future and not future.done():
298                    future.set_result(None)
299                if self._send_queue.empty():
300                    break
301                _, future = self._send_queue.get_nowait()