From 06b0800df88d3575efff64f30104db3a053fb02c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20K=C4=99pie=C5=84?= <michal@isc.org>
Date: Tue, 15 Jul 2025 12:16:32 +0200
Subject: [PATCH] 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.
---
 bin/tests/system/isctest/asyncserver.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/bin/tests/system/isctest/asyncserver.py b/bin/tests/system/isctest/asyncserver.py
index 784ee8efd2..1baeec697b 100644
--- 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