327 lines
10 KiB
Python
327 lines
10 KiB
Python
import asyncio
|
|
import pathlib
|
|
import socket
|
|
import sys
|
|
import warnings
|
|
|
|
__all__ = ("TransportClosed", "bind", "connect", "from_socket")
|
|
|
|
_windows = sys.platform == "win32"
|
|
|
|
|
|
class TransportClosed(Exception):
|
|
"""
|
|
Raised when the asyncio.DatagramTransport underlying a DatagramStream is
|
|
closed.
|
|
"""
|
|
|
|
|
|
class DatagramStream:
|
|
"""
|
|
Representation of a Datagram socket attached via either bind() or
|
|
connect() returned to consumers of this module. Provides simple
|
|
wrappers around sending and receiving bytes.
|
|
|
|
Due to the stateless nature of datagram protocols, errors are not
|
|
immediately available to this class at the point an action was performed
|
|
that will generate it. Rather, successive calls will raise exceptions if
|
|
there are any. Checking for exceptions can be done explicitly by using the
|
|
exception property.
|
|
|
|
For instance, failure to connect to a remote endpoint will not be noticed
|
|
until some point in time later, at which point ConnectionRefused will be
|
|
raised.
|
|
"""
|
|
|
|
def __init__(self, transport, recvq, excq, drained):
|
|
"""
|
|
@param transport - asyncio transport
|
|
@param recvq - asyncio queue that gets populated by the
|
|
DatagramProtocol with received datagrams.
|
|
@param excq - asyncio queue that gets populated with any errors
|
|
detected by the DatagramProtocol.
|
|
@param drained - asyncio event that is unset when writing is
|
|
paused and set otherwise.
|
|
"""
|
|
self._transport = transport
|
|
self._recvq = recvq
|
|
self._excq = excq
|
|
self._drained = drained
|
|
|
|
def __del__(self):
|
|
self._transport.close()
|
|
|
|
@property
|
|
def exception(self):
|
|
"""
|
|
If the underlying protocol detected an error, raise the first
|
|
unconsumed exception it noticed, otherwise returns None.
|
|
"""
|
|
try:
|
|
exc = self._excq.get_nowait()
|
|
raise exc
|
|
except asyncio.queues.QueueEmpty:
|
|
pass
|
|
|
|
@property
|
|
def sockname(self):
|
|
"""
|
|
The associated socket's own address
|
|
"""
|
|
r = self._transport.get_extra_info("sockname")
|
|
return None if r == "" else r
|
|
|
|
@property
|
|
def peername(self):
|
|
"""
|
|
The address the associated socket is connected to
|
|
"""
|
|
r = self._transport.get_extra_info("peername")
|
|
return None if r == "" else r
|
|
|
|
@property
|
|
def socket(self):
|
|
"""
|
|
The socket instance used by the stream. In python <3.8 this is a
|
|
socket.socket instance, after it is an asyncio.TransportSocket
|
|
instance.
|
|
"""
|
|
return self._transport.get_extra_info("socket")
|
|
|
|
def close(self):
|
|
"""
|
|
Close the underlying transport.
|
|
"""
|
|
self._transport.close()
|
|
|
|
async def _send(self, data, addr=None):
|
|
"""
|
|
@param data - bytes to send
|
|
@param addr - remote address to send data to, if unspecified then the
|
|
underlying socket has to have been been connected to a
|
|
remote address previously.
|
|
|
|
@raises TransportClosed - DatagramTransport closed.
|
|
"""
|
|
if self._transport.is_closing():
|
|
raise TransportClosed()
|
|
|
|
_ = self.exception
|
|
self._transport.sendto(data, addr)
|
|
await self._drained.wait()
|
|
|
|
async def recv(self):
|
|
"""
|
|
Receive data on the local socket.
|
|
|
|
@return - tuple of the bytes received and the address (ip, port) that
|
|
the data was received from.
|
|
|
|
@raises TransportClosed - DatagramTransport closed.
|
|
"""
|
|
if self._transport.is_closing():
|
|
raise TransportClosed()
|
|
|
|
_ = self.exception
|
|
data, addr = await self._recvq.get()
|
|
if data is None:
|
|
raise TransportClosed()
|
|
|
|
return data, addr
|
|
|
|
|
|
class DatagramServer(DatagramStream):
|
|
"""
|
|
Datagram socket bound to an address on the local machine.
|
|
"""
|
|
|
|
async def send(self, data, addr):
|
|
"""
|
|
@param data - bytes to send
|
|
@param addr - remote address to send data to.
|
|
"""
|
|
await super()._send(data, addr)
|
|
|
|
|
|
class DatagramClient(DatagramStream):
|
|
"""
|
|
Datagram socket connected to a remote address.
|
|
"""
|
|
|
|
async def send(self, data):
|
|
"""
|
|
@param data - bytes to send
|
|
"""
|
|
await super()._send(data)
|
|
|
|
|
|
class Protocol(asyncio.DatagramProtocol):
|
|
"""
|
|
asyncio.DatagramProtocol for feeding received packets into the
|
|
Datagram{Client,Server} which handles converting the lower level callback
|
|
based asyncio into higher level coroutines.
|
|
"""
|
|
|
|
def __init__(self, recvq, excq, drained):
|
|
"""
|
|
@param recvq - asyncio.Queue for new datagrams
|
|
@param excq - asyncio.Queue for exceptions
|
|
@param drained - asyncio.Event set when the write buffer is below the
|
|
high watermark.
|
|
"""
|
|
self._recvq = recvq
|
|
self._excq = excq
|
|
self._drained = drained
|
|
|
|
self._drained.set()
|
|
|
|
# Transports are connected at the time a connection is made.
|
|
self._transport = None
|
|
|
|
def connection_made(self, transport):
|
|
if self._transport is not None:
|
|
old_peer = self._transport.get_extra_info("peername")
|
|
new_peer = transport.get_extra_info("peername")
|
|
warnings.warn(
|
|
"Reinitializing transport connection from %s to %s", old_peer, new_peer
|
|
)
|
|
|
|
self._transport = transport
|
|
|
|
def connection_lost(self, exc):
|
|
if exc is not None:
|
|
self._excq.put_nowait(exc)
|
|
|
|
self._recvq.put_nowait((None, None))
|
|
|
|
if self._transport is not None:
|
|
self._transport.close()
|
|
self._transport = None
|
|
|
|
def datagram_received(self, data, addr):
|
|
self._recvq.put_nowait((data, addr))
|
|
|
|
def error_received(self, exc):
|
|
self._excq.put_nowait(exc)
|
|
|
|
def pause_writing(self):
|
|
self._drained.clear()
|
|
super().pause_writing()
|
|
|
|
def resume_writing(self):
|
|
self._drained.set()
|
|
super().resume_writing()
|
|
|
|
|
|
async def bind(addr, reuse_port=None):
|
|
"""
|
|
Bind a socket to a local address for datagrams. The socket will be either
|
|
AF_INET, AF_INET6 or AF_UNIX depending upon the type of address specified.
|
|
@param addr - For AF_INET or AF_INET6, a tuple with the the host and port to
|
|
to bind; port may be set to 0 to get any free port.
|
|
For AF_UNIX the path at which to bind (with a leading \0 for
|
|
abstract sockets).
|
|
@param reuse_port - Tells the kernel to allow this endpoint to be bound to
|
|
the same port as other existing endpoints are bound to, so long as
|
|
they all set this flag when being created. This option is not
|
|
supported on Windows and some UNIX's. If the
|
|
:py:data:`~socket.SO_REUSEPORT` constant is not defined then this
|
|
capability is unsupported.
|
|
@return - A DatagramServer instance
|
|
"""
|
|
loop = asyncio.get_event_loop()
|
|
recvq = asyncio.Queue()
|
|
excq = asyncio.Queue()
|
|
drained = asyncio.Event()
|
|
|
|
if not _windows and not isinstance(addr, tuple):
|
|
family = socket.AF_UNIX
|
|
if isinstance(addr, pathlib.Path):
|
|
addr = str(addr)
|
|
else:
|
|
family = 0
|
|
|
|
transport, protocol = await loop.create_datagram_endpoint(
|
|
lambda: Protocol(recvq, excq, drained),
|
|
local_addr=addr,
|
|
family=family,
|
|
reuse_port=reuse_port,
|
|
)
|
|
|
|
return DatagramServer(transport, recvq, excq, drained)
|
|
|
|
|
|
async def connect(addr):
|
|
"""
|
|
Connect a socket to a remote address for datagrams. The socket will be
|
|
either AF_INET, AF_INET6 or AF_UNIX depending upon the type of host
|
|
specified.
|
|
|
|
@param addr - For AF_INET or AF_INET6, a tuple with the the host and port to
|
|
to connect to.
|
|
For AF_UNIX the path at which to connect (with a leading \0
|
|
for abstract sockets).
|
|
@return - A DatagramClient instance
|
|
"""
|
|
loop = asyncio.get_event_loop()
|
|
recvq = asyncio.Queue()
|
|
excq = asyncio.Queue()
|
|
drained = asyncio.Event()
|
|
|
|
if not _windows and not isinstance(addr, tuple):
|
|
family = socket.AF_UNIX
|
|
if isinstance(addr, pathlib.Path):
|
|
addr = str(addr)
|
|
else:
|
|
family = 0
|
|
|
|
transport, protocol = await loop.create_datagram_endpoint(
|
|
lambda: Protocol(recvq, excq, drained),
|
|
remote_addr=addr,
|
|
family=family,
|
|
)
|
|
|
|
return DatagramClient(transport, recvq, excq, drained)
|
|
|
|
|
|
async def from_socket(sock):
|
|
"""
|
|
Create a DatagramStream from a socket. This is meant to be used in cases
|
|
where the defaults set by `bind()` and `connect()` are not desired and/or
|
|
sufficient. If `socket.connect()` was previously called on the socket,
|
|
then an instance of DatagramClient will be returned, otherwise an instance
|
|
of DatagramServer.
|
|
|
|
@param sock - socket to use in the DatagramStream.
|
|
@return - A DatagramClient for connected sockets, otherwise a
|
|
DatagramServer.
|
|
"""
|
|
loop = asyncio.get_event_loop()
|
|
recvq = asyncio.Queue()
|
|
excq = asyncio.Queue()
|
|
drained = asyncio.Event()
|
|
|
|
if not _windows:
|
|
supported_families = tuple((socket.AF_INET, socket.AF_INET6, socket.AF_UNIX))
|
|
else:
|
|
supported_families = tuple((socket.AF_INET, socket.AF_INET6))
|
|
|
|
if sock.family not in supported_families:
|
|
raise TypeError(
|
|
"socket family not one of %s"
|
|
% (", ".join(str(f) for f in supported_families))
|
|
)
|
|
|
|
if sock.type != socket.SOCK_DGRAM:
|
|
raise TypeError("socket type must be %s" % (socket.SOCK_DGRAM,))
|
|
|
|
transport, protocol = await loop.create_datagram_endpoint(
|
|
lambda: Protocol(recvq, excq, drained), sock=sock
|
|
)
|
|
|
|
if transport.get_extra_info("peername") is not None:
|
|
# Workaround transport ignoring the peer address of the socket.
|
|
transport._address = transport.get_extra_info("peername")
|
|
return DatagramClient(transport, recvq, excq, drained)
|
|
else:
|
|
return DatagramServer(transport, recvq, excq, drained)
|