Enable requesting TCP connections to be closed

In response to client queries, AsyncDnsServer users can currently only
make the server either send a reply or silently ignore the query.  In
the case of TCP queries, neither of these actions causes the client's
connection to be closed - the onus of doing that is on the client.
However, in some cases the server may be required to close the
connection on its own, so AsyncDnsServer users need to have some way of
requesting such an action.

Add a new ResponseAction subclass, ResponseDropAndCloseConnection, which
enables AsyncDnsServer users to conveniently request TCP connections to
be closed.  Instead of returning the response to send,
ResponseDropAndCloseConnection raises a custom exception that
AsyncDnsServer._handle_tcp() handles accordingly.

diff --git a/bin/tests/system/isctest/asyncserver.py b/bin/tests/system/isctest/asyncserver.py
index 784ee8efd226305007cbc65d944a5e6424d96749..1baeec697b0713e845884664cab48cba29e97b8b 100644 (file)
--- a/bin/tests/system/isctest/asyncserver.py
+++ b/bin/tests/system/isctest/asyncserver.py
@@ -373,6 +373,28 @@ 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 ResponseHandler(abc.ABC):
     """
     Base class for generic response handlers.
@@ -690,6 +712,8 @@ class AsyncDnsServer(AsyncServer):
                 if not wire:
                     break
                 await self._send_tcp_response(writer, peer, wire)
+            except _ConnectionTeardownRequested:
+                break
             except ConnectionResetError:
                 logging.error("TCP connection from %s reset by peer", peer)
                 return