chg: test: Use isctest.asyncserver in the "dispatch" test

Replace the custom DNS server used in the "dispatch" system test with
new code based on the isctest.asyncserver module.

Merge branch 'stepan/dispatch-asyncserver' into 'main'

See merge request isc-projects/bind9!10689

bin/tests/system/dispatch/ans3/ans.py

@@ -9,91 +9,37 @@
 # See the COPYRIGHT file distributed with this work for additional
 # information regarding copyright ownership.
 
-import os
+from typing import AsyncGenerator
-import select
-import signal
-import socket
-import sys
-import time
 
-import dns.flags
+import dns
-import dns.message
+
+from isctest.asyncserver import (
+    AsyncDnsServer,
+    ConnectionReset,
+    DnsProtocol,
+    DnsResponseSend,
+    QueryContext,
+    ResponseAction,
+    ResponseHandler,
+)
 
 
-def port():
+class TruncateOnUdpHandler(ResponseHandler):
-    env_port = os.getenv("PORT")
+    async def get_responses(
-    if env_port is None:
+        self, qctx: QueryContext
-        env_port = 5300
+    ) -> AsyncGenerator[ResponseAction, None]:
-    else:
+        assert qctx.protocol == DnsProtocol.UDP, "This server only supports UDP"
-        env_port = int(env_port)
+        qctx.response.set_rcode(dns.rcode.NOERROR)
-
+        qctx.response.flags |= dns.flags.TC
-    return env_port
+        yield DnsResponseSend(qctx.response)
 
 
-def udp_listen(port):
+def main() -> None:
-    udp = socket.socket(type=socket.SOCK_DGRAM)
+    server = AsyncDnsServer()
-    udp.bind(("10.53.0.3", port))
+    server.install_connection_handler(ConnectionReset(delay=1.0))
-
+    server.install_response_handler(TruncateOnUdpHandler())
-    return udp
+    server.run()
 
 
-def tcp_listen(port):
+if __name__ == "__main__":
-    tcp = socket.socket(type=socket.SOCK_STREAM)
+    main()
-    tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-    tcp.bind(("10.53.0.3", port))
-    tcp.listen(100)
-
-    return tcp
-
-
-def udp_tc_once(udp):
-    qrybytes, clientaddr = udp.recvfrom(65535)
-    qry = dns.message.from_wire(qrybytes)
-    answ = dns.message.make_response(qry)
-    answ.flags |= dns.flags.TC
-    answbytes = answ.to_wire()
-    udp.sendto(answbytes, clientaddr)
-
-
-def tcp_once(tcp):
-    csock, _clientaddr = tcp.accept()
-    time.sleep(5)
-    csock.close()
-
-
-def sigterm(signum, frame):
-    os.remove("ans.pid")
-    sys.exit(0)
-
-
-def write_pid():
-    with open("ans.pid", "w") as f:
-        pid = os.getpid()
-        f.write("{}".format(pid))
-
-
-signal.signal(signal.SIGTERM, sigterm)
-write_pid()
-
-udp = udp_listen(port())
-tcp = tcp_listen(port())
-
-input = [udp, tcp]
-
-while True:
-    try:
-        inputready, outputready, exceptready = select.select(input, [], [])
-    except select.error:
-        break
-    except socket.error:
-        break
-    except KeyboardInterrupt:
-        break
-
-    for s in inputready:
-        if s == udp:
-            udp_tc_once(udp)
-        if s == tcp:
-            tcp_once(tcp)
-
-sigterm(signal.SIGTERM, 0)

bin/tests/system/isctest/asyncserver.py

@@ -373,6 +373,118 @@ class ResponseDrop(ResponseAction):
         return None
 
 
+class _ConnectionTeardownRequested(Exception):
+    pass
+
+
+@dataclass
+class ResponseDropAndCloseConnection(ResponseAction):
+    """
+    Action which makes the server close the connection after the DNS query is
+    received by the server (TCP only).
+
+    The connection may be closed with a delay if requested.
+    """
+
+    delay: float = 0.0
+
+    async def perform(self) -> Optional[Union[dns.message.Message, bytes]]:
+        if self.delay > 0:
+            logging.info("Waiting %.1fs before closing TCP connection", self.delay)
+            await asyncio.sleep(self.delay)
+        raise _ConnectionTeardownRequested
+
+
+class ConnectionHandler(abc.ABC):
+    """
+    Base class for TCP connection handlers.
+
+    An installed connection handler is called when a new TCP connection is
+    established.  It may be used to perform arbitrary actions before
+    AsyncDnsServer processes DNS queries.
+    """
+
+    @abc.abstractmethod
+    async def handle(
+        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, peer: Peer
+    ) -> None:
+        """
+        Handle the connection with the provided reader and writer.
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class ConnectionReset(ConnectionHandler):
+    """
+    A connection handler that makes the server close the connection without
+    reading anything from the client socket.
+
+    The connection may be closed with a delay if requested.
+
+    The sole purpose of this handler is to trigger a connection reset, i.e. to
+    make the server send an RST segment; this happens when the server closes a
+    client's socket while there is still unread data in that socket's buffer.
+    If closing the connection _after_ the query is read by the server is enough
+    for a given use case, the ResponseDropAndCloseConnection response handler
+    should be used instead.
+    """
+
+    delay: float = 0.0
+
+    async def handle(
+        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, peer: Peer
+    ) -> None:
+        try:
+            # Python >= 3.7
+            loop = asyncio.get_running_loop()
+        except AttributeError:
+            # Python < 3.7
+            loop = asyncio.get_event_loop()
+
+        logging.info("Blocking reads from %s", peer)
+
+        # This is Michał's submission for the Ugliest Hack of the Year contest.
+        # (The alternative was implementing an asyncio transport from scratch.)
+        #
+        # In order to prevent the client socket from being read from, simply
+        # not calling `reader.read()` is not enough, because asyncio buffers
+        # incoming data itself on the transport level.  However, `StreamReader`
+        # does not expose the underlying transport as a property.  Therefore,
+        # cheat by extracting it from `StreamWriter` as it is the same
+        # bidirectional transport as for the read side (a `Transport`, which is
+        # a subclass of both `ReadTransport` and `WriteTransport`) and call
+        # `ReadTransport.pause_reading()` to remove the underlying socket from
+        # the set of descriptors monitored by the selector, thereby preventing
+        # any reads from happening on the client socket.  However...
+        loop.call_soon(writer.transport.pause_reading)  # type: ignore
+
+        # ...due to `AsyncDnsServer._handle_tcp()` being a coroutine, by the
+        # time it gets executed, asyncio transport code will already have added
+        # the client socket to the set of descriptors monitored by the
+        # selector.  Therefore, if the client starts sending data immediately,
+        # a read from the socket will have already been scheduled by the time
+        # this handler gets executed.  There is no way to prevent that from
+        # happening, so work around it by abusing the fact that the transport
+        # at hand is specifically an instance of `_SelectorSocketTransport`
+        # (from asyncio.selector_events) and set the size of its read buffer to
+        # just a single byte.  This does give asyncio enough time to read that
+        # single byte from the client socket's buffer before that socket is
+        # removed from the set of monitored descriptors, but prevents the
+        # one-off read from emptying the client socket buffer _entirely_, which
+        # is enough to trigger sending an RST segment when the connection is
+        # closed shortly afterwards.
+        writer.transport.max_size = 1  # type: ignore
+
+        if self.delay > 0:
+            logging.info(
+                "Waiting %.1fs before closing TCP connection from %s", self.delay, peer
+            )
+            await asyncio.sleep(self.delay)
+
+        raise _ConnectionTeardownRequested
+
+
 class ResponseHandler(abc.ABC):
     """
     Base class for generic response handlers.
@@ -605,6 +717,7 @@ class AsyncDnsServer(AsyncServer):
         super().__init__(self._handle_udp, self._handle_tcp, "ans.pid")
 
         self._zone_tree: _ZoneTree = _ZoneTree()
+        self._connection_handler: Optional[ConnectionHandler] = None
         self._response_handlers: List[ResponseHandler] = []
         self._acknowledge_manual_dname_handling = acknowledge_manual_dname_handling
         self._acknowledge_tsig_dnspython_hacks = acknowledge_tsig_dnspython_hacks
@@ -637,6 +750,15 @@ class AsyncDnsServer(AsyncServer):
         logging.info("Uninstalling response handler: %s", handler)
         self._response_handlers.remove(handler)
 
+    def install_connection_handler(self, handler: ConnectionHandler) -> None:
+        """
+        Install a connection handler that will be called when a new TCP
+        connection is established.
+        """
+        if self._connection_handler:
+            raise RuntimeError("Only one connection handler can be installed")
+        self._connection_handler = handler
+
     def _load_zones(self) -> None:
         for entry in os.scandir():
             entry_path = pathlib.Path(entry.path)
@@ -684,15 +806,19 @@ class AsyncDnsServer(AsyncServer):
         peer = Peer(peer_info[0], peer_info[1])
         logging.debug("Accepted TCP connection from %s", peer)
 
-        while True:
+        try:
-            try:
+            if self._connection_handler:
+                await self._connection_handler.handle(reader, writer, peer)
+            while True:
                 wire = await self._read_tcp_query(reader, peer)
                 if not wire:
                     break
                 await self._send_tcp_response(writer, peer, wire)
-            except ConnectionResetError:
+        except _ConnectionTeardownRequested:
-                logging.error("TCP connection from %s reset by peer", peer)
+            pass
-                return
+        except ConnectionResetError:
+            logging.error("TCP connection from %s reset by peer", peer)
+            return
 
         logging.debug("Closing TCP connection from %s", peer)
         writer.close()