hat.drivers.net

20class UdpAddress(typing.NamedTuple):
21    host: str
22    port: int

28class UdpEndpointInfo(typing.NamedTuple):
29    name: str | None
30    local_addr: UdpAddress
31    remote_addr: UdpAddress | None

34class UnixEndpointInfo(typing.NamedTuple):
35    name: str | None
36    local_addr: common.UnixAddress | None
37    remote_addr: common.UnixAddress | None

43class DatagramType(enum.Enum):
44    UDP = 0
45    UNIX = 1

48def endpoint_info_to_json(info: EndpointInfo) -> json.Data:
49    if isinstance(info, UdpEndpointInfo):
50        data = {'name': info.name,
51                'local_addr': {'host': info.local_addr.host,
52                               'port': info.local_addr.port}}
53
54        if info.remote_addr is not None:
55            data['remote_addr'] = {'host': info.remote_addr.host,
56                                   'port': info.remote_addr.port}
57
58        return data
59
60    if isinstance(info, UnixEndpointInfo):
61        data = {'name': info.name}
62
63        if info.local_addr is not None:
64            data['local_addr'] = str(info.local_addr)
65
66        if info.remote_addr is not None:
67            data['remote_addr'] = str(info.remote_addr)
68
69        return data
70
71    raise TypeError('unsupported info type')

 74async def create_endpoint(datagram_type: DatagramType = DatagramType.UDP,
 75                          local_addr: DatagramAddress | None = None,
 76                          remote_addr: DatagramAddress | None = None,
 77                          *,
 78                          name: str | None = None,
 79                          receive_queue_size: int = 0,
 80                          **kwargs
 81                          ) -> 'Endpoint':
 82    """Create new UDP or Unix Domain Socket endpoint
 83
 84    Args:
 85        datagram_type: datagram protocol type
 86        local_addr: local address
 87        remote_addr: remote address
 88        name: endpoint name
 89        receive_queue_size: receive queue max size
 90        kwargs: additional arguments passed to
 91            `asyncio.AbstractEventLoop.create_datagram_endpoint`
 92
 93    """
 94    loop = asyncio.get_running_loop()
 95
 96    create_protocol = functools.partial(_Protocol, datagram_type, name,
 97                                        receive_queue_size)
 98    family = _get_address_family(datagram_type)
 99
100    if isinstance(local_addr, common.UnixAddress):
101        local_addr = str(local_addr)
102
103    if isinstance(remote_addr, common.UnixAddress):
104        remote_addr = str(remote_addr)
105
106    _, protocol = await loop.create_datagram_endpoint(create_protocol,
107                                                      local_addr,
108                                                      remote_addr,
109                                                      family=family,
110                                                      **kwargs)
111
112    return _create_endpoint(protocol)

127class Endpoint(aio.Resource):
128    """UDP or Unix Domain Socket endpoint"""
129
130    @property
131    def async_group(self) -> aio.Group:
132        """Async group"""
133        return self._async_group
134
135    @property
136    def info(self) -> EndpointInfo:
137        """Endpoint info"""
138        return self._protocol.info
139
140    @property
141    def empty(self) -> bool:
142        """Is receive queue empty"""
143        return self._protocol.empty
144
145    def send(self,
146             data: util.Bytes,
147             remote_addr: DatagramAddress | None = None):
148        """Send datagram
149
150        If `remote_addr` is not set, `remote_addr` passed to `create` is used.
151
152        """
153        if not self.is_open:
154            raise ConnectionError()
155
156        self._protocol.send(data, remote_addr)
157
158    async def receive(self) -> tuple[util.Bytes, DatagramAddress | None]:
159        """Receive datagram"""
160        return await self._protocol.receive()

24class TcpAddress(typing.NamedTuple):
25    host: str
26    port: int

32class TcpConnectionInfo(typing.NamedTuple):
33    name: str | None
34    local_addr: TcpAddress
35    remote_addr: TcpAddress

38class UnixConnectionInfo(typing.NamedTuple):
39    name: str | None
40    addr: common.UnixAddress

43class TcpServerInfo(typing.NamedTuple):
44    name: str | None
45    addresses: list[TcpAddress]

48class UnixServerInfo(typing.NamedTuple):
49    name: str | None
50    addresses: list[common.UnixAddress]

58class StreamType(enum.Enum):
59    TCP = 0
60    UNIX = 1

67def connection_info_to_json(info: ConnectionInfo) -> json.Data:
68    if isinstance(info, TcpConnectionInfo):
69        return {'name': info.name,
70                'local_addr': {'host': info.local_addr.host,
71                               'port': info.local_addr.port},
72                'remote_addr': {'host': info.remote_addr.host,
73                                'port': info.remote_addr.port}}
74
75    if isinstance(info, UnixConnectionInfo):
76        return {'name': info.name,
77                'addr': str(info.addr)}
78
79    raise TypeError('unsupported info type')

82def server_info_to_json(info: ServerInfo) -> json.Data:
83    if isinstance(info, TcpServerInfo):
84        return {'name': info.name,
85                'addresses': [{'host': addr.host,
86                               'port': addr.port}
87                              for addr in info.addresses]}
88
89    if isinstance(info, UnixServerInfo):
90        return {'name': info.name,
91                'addresses': [str(addr) for addr in info.addresses]}
92
93    raise TypeError('unsupported info type')

 96async def connect(addr: StreamAddress,
 97                  *,
 98                  name: str | None = None,
 99                  input_buffer_limit: int = 64 * 1024,
100                  **kwargs
101                  ) -> 'Connection':
102    """Create TCP or Unix Domain Socket connection
103
104    Argument `addr` specifies remote server listening address.
105
106    Argument `name` defines connection name available in property `info`.
107
108    Argument `input_buffer_limit` defines number of bytes in input buffer
109    that whill temporary pause data receiving. Once number of bytes
110    drops bellow `input_buffer_limit`, data receiving is resumed. If this
111    argument is ``0``, data receive pausing is disabled.
112
113    Additional arguments are passed directly to `asyncio.create_connection` or
114    `asyncio.create_unix_connection`.
115
116    """
117    loop = asyncio.get_running_loop()
118    stream_type = _get_stream_type(addr)
119    create_protocol = functools.partial(_Protocol, stream_type, None, name,
120                                        input_buffer_limit)
121
122    if stream_type == StreamType.TCP:
123        _, protocol = await loop.create_connection(create_protocol,
124                                                   addr.host, addr.port,
125                                                   **kwargs)
126
127    elif stream_type == StreamType.UNIX:
128        _, protocol = await loop.create_unix_connection(create_protocol, addr,
129                                                        **kwargs)
130
131    else:
132        raise ValueError('unsupported stream type')
133
134    return _create_connection(protocol)

137async def listen(connection_cb: ConnectionCb,
138                 addr: StreamAddress,
139                 *,
140                 name: str | None = None,
141                 bind_connections: bool = False,
142                 input_buffer_limit: int = 64 * 1024,
143                 **kwargs
144                 ) -> 'Server':
145    """Create listening TCP or Unix Domain Socket server
146
147    Argument `name` defines server name available in property `info`. This
148    name is used for all incomming connections.
149
150    If `bind_connections` is ``True``, closing server will close all open
151    incoming connections.
152
153    Argument `input_buffer_limit` is associated with newly created connections
154    (see `connect`).
155
156    Additional arguments are passed directly to `asyncio.create_server` or
157    `asyncio.create_unix_server`.
158
159    """
160    server = Server()
161    server._connection_cb = connection_cb
162    server._bind_connections = bind_connections
163    server._stream_type = _get_stream_type(addr)
164    server._async_group = aio.Group()
165    server._log = mlog
166
167    loop = asyncio.get_running_loop()
168    on_connection = functools.partial(server.async_group.spawn,
169                                      server._on_connection)
170    create_protocol = functools.partial(_Protocol, server._stream_type,
171                                        on_connection, name,
172                                        input_buffer_limit)
173
174    if server._stream_type == StreamType.TCP:
175        server._srv = await loop.create_server(create_protocol,
176                                               addr.host, addr.port, **kwargs)
177
178    elif server._stream_type == StreamType.UNIX:
179        server._srv = await loop.create_unix_server(create_protocol,
180                                                    addr, **kwargs)
181
182    else:
183        raise ValueError('unsupported stream type')
184
185    server.async_group.spawn(aio.call_on_cancel, server._on_close)
186
187    try:
188        server._info = _get_server_info(server._stream_type, name, server._srv)
189        server._log = _create_server_logger(server._info)
190
191    except Exception:
192        await aio.uncancellable(server.async_close())
193        raise
194
195    server._log.debug('listening for incomming connections')
196
197    return server

200class Server(aio.Resource):
201    """TCP or Unix Domain Socket listening server
202
203    Closing server will cancel all running `connection_cb` coroutines.
204
205    """
206
207    @property
208    def async_group(self) -> aio.Group:
209        """Async group"""
210        return self._async_group
211
212    @property
213    def info(self) -> ServerInfo:
214        """Server info"""
215        return self._info
216
217    async def _on_close(self):
218        self._srv.close()
219
220        if self._bind_connections or sys.version_info[:2] < (3, 12):
221            await self._srv.wait_closed()
222
223    async def _on_connection(self, protocol):
224        self._log.debug('new incomming connection')
225
226        conn = _create_connection(protocol)
227
228        try:
229            await aio.call(self._connection_cb, conn)
230
231            if self._bind_connections:
232                await conn.wait_closing()
233
234            else:
235                conn = None
236
237        except Exception as e:
238            self._log.warning('connection callback error: %s', e, exc_info=e)
239
240        finally:
241            if conn:
242                await aio.uncancellable(conn.async_close())

257class Connection(aio.Resource):
258    """TCP or Unix Domain Socket connection"""
259
260    @property
261    def async_group(self) -> aio.Group:
262        """Async group"""
263        return self._async_group
264
265    @property
266    def info(self) -> ConnectionInfo:
267        """Connection info"""
268        return self._protocol.info
269
270    @property
271    def ssl_object(self) -> ssl.SSLObject | ssl.SSLSocket | None:
272        """SSL Object"""
273        return self._protocol.ssl_object
274
275    async def write(self, data: util.Bytes):
276        """Write data
277
278        This coroutine will wait until `data` can be added to output buffer.
279
280        """
281        if not self.is_open:
282            raise ConnectionError()
283
284        await self._protocol.write(data)
285
286    async def drain(self):
287        """Drain output buffer"""
288        await self._protocol.drain()
289
290    async def read(self, n: int = -1) -> util.Bytes:
291        """Read up to `n` bytes
292
293        If EOF is detected and no new bytes are available, `ConnectionError`
294        is raised.
295
296        """
297        return await self._protocol.read(n)
298
299    async def readexactly(self, n: int) -> util.Bytes:
300        """Read exactly `n` bytes
301
302        If exact number of bytes could not be read, `ConnectionError` is
303        raised.
304
305        """
306        return await self._protocol.readexactly(n)
307
308    def clear_input_buffer(self) -> int:
309        """Clear input buffer
310
311        Returns number of bytes cleared from buffer.
312
313        """
314        return self._protocol.clear_input_buffer()