Create dedicated datagram socket for server

The previous datagram socket for client use only allowed receiving a single message before it has been read (implementation via future).
For server use cases we would like to receive messages continuously.
This commit implements a dummy socket embedding received data such that dns.asyncquery.receive_udp() can still be used to read the message
in order to keep the interface as close as possible to the TCP implementation.
Also implementing a proper server class as public interface.

Author: darkdragon-001

dns/_asyncbackend.py

@@ -65,6 +65,20 @@ async def recv(self, size, timeout):
         raise NotImplementedError
 
 
+class Server:  # pragma: no cover
+    async def serve_forever(self):
+        raise NotImplementedError
+
+    async def close(self):
+        raise NotImplementedError
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.close()
+
+
 class NullTransport:
     async def connect_tcp(self, host, port, timeout, local_address):
         raise NotImplementedError

dns/_asyncio_backend.py

@@ -21,14 +21,21 @@ def _get_running_loop():
         return asyncio.get_event_loop()
 
 
-class _DatagramProtocol:
+class _BaseDatagramProtocol(asyncio.DatagramProtocol):
     def __init__(self):
         self.transport = None
-        self.recvfrom = None
 
     def connection_made(self, transport):
         self.transport = transport
 
+    def close(self):
+        if self.transport is not None:
+            self.transport.close()
+
+class _SocketDatagramProtocol(_BaseDatagramProtocol):
+    def __init__(self):
+        self.recvfrom = None
+
     def datagram_received(self, data, addr):
         if self.recvfrom and not self.recvfrom.done():
             self.recvfrom.set_result((data, addr))
@@ -48,9 +55,18 @@ def connection_lost(self, exc):
             else:
                 self.recvfrom.set_exception(exc)
 
-    def close(self):
-        if self.transport is not None:
-            self.transport.close()
+class _ServerDatagramProtocol(_BaseDatagramProtocol):
+    def __init__(self, datagram_received_cb, serving_future):
+        self.datagram_received_cb = datagram_received_cb
+        self.serving = serving_future
+
+    def datagram_received(self, data, addr):
+        if self.datagram_received_cb:
+            asyncio.ensure_future(self.datagram_received_cb(addr, self.transport, self, data))
+
+    def connection_lost(self, exc):
+        if self.serving and not self.serving.done():
+            self.serving.set_result(True)
 
 
 async def _maybe_wait_for(awaitable, timeout):
@@ -63,7 +79,7 @@ async def _maybe_wait_for(awaitable, timeout):
         return await awaitable
 
 
-class _DatagramSocket(dns._asyncbackend.DatagramSocket):
+class _BaseDatagramSocket(dns._asyncbackend.DatagramSocket):
     def __init__(self, family, transport, protocol):
         super().__init__(family, socket.SOCK_DGRAM)
         self.transport = transport
@@ -74,6 +90,20 @@ async def sendto(self, what, destination, timeout):  # pragma: no cover
         self.transport.sendto(what, destination)
         return len(what)
 
+    async def close(self):
+        self.protocol.close()
+
+    async def getpeername(self):
+        return self.transport.get_extra_info("peername")
+
+    async def getsockname(self):
+        return self.transport.get_extra_info("sockname")
+
+    async def getpeercert(self, timeout):
+        raise NotImplementedError
+
+
+class _SocketDatagramSocket(_BaseDatagramSocket):
     async def recvfrom(self, size, timeout):
         # ignore size as there's no way I know to tell protocol about it
         done = _get_running_loop().create_future()
@@ -85,17 +115,20 @@ async def recvfrom(self, size, timeout):
         finally:
             self.protocol.recvfrom = None
 
-    async def close(self):
-        self.protocol.close()
 
-    async def getpeername(self):
-        return self.transport.get_extra_info("peername")
+class _ServerDatagramSocket(_BaseDatagramSocket):
+    def __init__(self, family, transport, protocol, addr, data):
+        super().__init__(family, transport, protocol)
+        self.addr = addr
+        self.data = data
 
-    async def getsockname(self):
-        return self.transport.get_extra_info("sockname")
-
-    async def getpeercert(self, timeout):
-        raise NotImplementedError
+    async def recvfrom(self, size, timeout):
+        if self.data is None:
+            raise EOFError("EOF")
+        # data always contains exactly one messaage
+        result = (self.data, self.addr)
+        self.data = None
+        return result
 
 
 class _StreamSocket(dns._asyncbackend.StreamSocket):
@@ -124,6 +157,30 @@ async def getpeercert(self, timeout):
         return self.writer.get_extra_info("peercert")
 
 
+class DatagramServer(dns._asyncbackend.Server):
+    def __init__(self, transport, protocol):
+        self.transport = transport
+        self.protocol = protocol
+
+    async def serve_forever(self):
+        await self.protocol.serving
+
+    async def close(self):
+        if self.transport:
+            self.transport.close()
+
+
+class StreamServer(dns._asyncbackend.Server):
+    def __init__(self, server):
+        self.server = server
+
+    async def serve_forever(self):
+        await self.server.serve_forever()
+
+    async def close(self):
+        await self.server.close()
+
+
 if dns._features.have("doh"):
     import anyio
     import httpcore
@@ -234,13 +291,13 @@ async def make_socket(
                 # proper fix for [#637].
                 source = (dns.inet.any_for_af(af), 0)
             transport, protocol = await loop.create_datagram_endpoint(
-                _DatagramProtocol,  # type: ignore
+                _SocketDatagramProtocol,  # type: ignore
                 local_addr=source,
                 family=af,
                 proto=proto,
                 remote_addr=destination,
             )
-            return _DatagramSocket(af, transport, protocol)
+            return _SocketDatagramSocket(af, transport, protocol)
         elif socktype == socket.SOCK_STREAM:
             if destination is None:
                 # This shouldn't happen, but we check to make code analysis software
@@ -270,10 +327,22 @@ async def make_server(
         socktype,
         addr,
     ):
+        loop = _get_running_loop()
         if socktype == socket.SOCK_DGRAM:
-            raise NotImplementedError(
-                "server not necessary for datagram, use make_socket() instead"
-            )  # pragma: no cover
+            if _is_win32 and addr is None:
+                # Win32 wants explicit binding before recvfrom().  This is the
+                # proper fix for [#637].
+                addr = (dns.inet.any_for_af(af), 0)
+            async def handle_udp(addr, transport, protocol, data):
+                sock_udp = _ServerDatagramSocket(af, transport, protocol, addr, data)
+                await client_connected_cb(sock_udp)
+            done = _get_running_loop().create_future()
+            transport, protocol = await loop.create_datagram_endpoint(
+                lambda: _ServerDatagramProtocol(handle_udp, done),
+                local_addr=addr,
+                family=af,
+            )
+            return DatagramServer(transport, protocol)
         elif socktype == socket.SOCK_STREAM:
             async def handle_tcp(r, w):
                 sock_tcp = _StreamSocket(af, r, w)

dns/asyncserver.py

@@ -20,12 +20,7 @@ async def udp_serve(
     ignore_errors: bool = False,
     backend: dns.asyncbackend.Backend | None = None,
 ) -> None:
-    if not backend:
-        backend = dns.asyncbackend.get_default_backend()
-    af = dns.inet.af_for_address(host)
-    addr = (host, port)
-    sock = await backend.make_socket(af, socket.SOCK_DGRAM, 0, addr)
-    while True:
+    async def handle_udp(sock: dns.asyncbackend.DatagramSocket):
         try:
             (m, _, from_address) = await dns.asyncquery.receive_udp(
                 sock,
@@ -44,6 +39,13 @@ async def udp_serve(
         except:
             pass
 
+    if not backend:
+        backend = dns.asyncbackend.get_default_backend()
+    af = dns.inet.af_for_address(host)
+    addr = (host, port)
+    server = await backend.make_server(handle_udp, af, socket.SOCK_DGRAM, addr)
+    await server.serve_forever()
+
 
 async def tcp_serve(
     cb: Callable[[dns.message.Message, str], Awaitable[dns.message.Message]],