--- /dev/null
+"""
+Copyright (C) Internet Systems Consortium, Inc. ("ISC")
+
+SPDX-License-Identifier: MPL-2.0
+
+This Source Code Form is subject to the terms of the Mozilla Public
+License, v. 2.0. If a copy of the MPL was not distributed with this
+file, you can obtain one at https://mozilla.org/MPL/2.0/.
+
+See the COPYRIGHT file distributed with this work for additional
+information regarding copyright ownership.
+"""
+
+from collections.abc import AsyncGenerator
+from typing import final
+
+import abc
+import asyncio
+
+import dns.flags
+import dns.name
+import dns.rcode
+import dns.rdataclass
+import dns.rdatatype
+import dns.rrset
+
+from isctest.asyncserver import (
+ ControlCommand,
+ ControllableAsyncDnsServer,
+ DnsResponseSend,
+ QnameHandler,
+ QueryContext,
+ ResponseAction,
+ ResponseHandler,
+)
+
+
+class ReclimitStateHandler(QnameHandler):
+ """
+ Handler for the "count." and "reset." queries that also holds the state
+ shared by all the handlers in one server.
+ """
+
+ qnames = ["count.", "reset."]
+
+ def __init__(self, indirect_send_response_default: bool = True) -> None:
+ self._indirect_send_response_default = indirect_send_response_default
+ self.count = 0
+ self.limit = 0
+ self.indirect_send_response = indirect_send_response_default
+ super().__init__()
+
+ async def get_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[DnsResponseSend, None]:
+ if f"{qctx.qname}" == "count.":
+ self.count += 1
+ qctx.response.answer.append(
+ dns.rrset.from_text(
+ "count.", 0, dns.rdataclass.IN, dns.rdatatype.TXT, f"{self.count}"
+ )
+ )
+ yield DnsResponseSend(qctx.response, authoritative=True)
+ elif f"{qctx.qname}" == "reset.":
+ self.reset()
+ yield DnsResponseSend(qctx.response, authoritative=False)
+
+ def reset(self) -> None:
+ self.count = 0
+ self.indirect_send_response = self._indirect_send_response_default
+
+
+class ReclimitHandler(ResponseHandler):
+ """
+ Base class for handlers in this test.
+
+ Increments the shared query counter on each query and delegates the actual
+ response generation to the `_get_counted_responses()` method.
+ """
+
+ def __init__(self, state_handler: ReclimitStateHandler) -> None:
+ self._state = state_handler
+ super().__init__()
+
+ @final
+ async def get_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[ResponseAction, None]:
+ self._state.count += 1
+ async for response in self._get_counted_responses(qctx):
+ yield response
+
+ @abc.abstractmethod
+ async def _get_counted_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[ResponseAction, None]:
+ yield DnsResponseSend(qctx.response)
+
+
+class LimitControlCommand(ControlCommand):
+ control_subdomain = "limit"
+
+ def __init__(self, state_handler: ReclimitStateHandler) -> None:
+ self._state_handler = state_handler
+ super().__init__()
+
+ def handle(
+ self, args: list[str], server: ControllableAsyncDnsServer, qctx: QueryContext
+ ) -> str | None:
+ if len(args) != 1:
+ return "Expected exactly one label"
+
+ try:
+ limit = int(args[0])
+ except ValueError:
+ return "Expected an integer"
+
+ self._state_handler.limit = limit
+ return f"Limit set to {limit}"
+
+
+def a(owner: str | dns.name.Name, ns_number: int) -> dns.rrset.RRset:
+ return dns.rrset.from_text(
+ f"{owner}", 3600, dns.rdataclass.IN, dns.rdatatype.A, f"10.53.0.{ns_number}"
+ )
+
+
+def ns(owner: str | dns.name.Name, target: str | dns.name.Name) -> dns.rrset.RRset:
+ return dns.rrset.from_text(
+ f"{owner}", 86400, dns.rdataclass.IN, dns.rdatatype.NS, f"{target}"
+ )
+
+
+class DirectExampleHandler(ReclimitHandler, QnameHandler):
+ qnames = ["direct.example.org", "direct.example.net"]
+
+ def __init__(
+ self, state_handler: ReclimitStateHandler, local_ns_number: int
+ ) -> None:
+ self._local_ns_number = local_ns_number
+ super().__init__(state_handler)
+
+ async def _get_counted_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[DnsResponseSend, None]:
+ if qctx.qtype == dns.rdatatype.A:
+ qctx.response.answer.append(a(qctx.qname, self._local_ns_number))
+ yield DnsResponseSend(qctx.response)
+
+
+class IndirectExampleOrgHandler(ReclimitHandler, QnameHandler):
+ qnames = [f"indirect{i}.example.org" for i in range(1, 9)]
+
+ def __init__(
+ self, state_handler: ReclimitStateHandler, local_ns_number: int
+ ) -> None:
+ self._local_ns_number = local_ns_number
+ super().__init__(state_handler)
+
+ async def _get_counted_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[DnsResponseSend, None]:
+ if not self._state.indirect_send_response:
+ qctx.response.authority.append(ns(f"{qctx.qname}", "ns1.1.example.org."))
+ qctx.response.flags &= ~dns.flags.AA
+ elif qctx.qtype == dns.rdatatype.A:
+ qctx.response.answer.append(a(qctx.qname, self._local_ns_number))
+ yield DnsResponseSend(qctx.response)
+
+
+def is_ns1_example(qname: dns.name.Name, tld: str) -> bool:
+ labels = qname.labels
+ return (
+ len(labels) == 5
+ and labels[3] == tld.encode()
+ and labels[2] == b"example"
+ and labels[1].isdigit()
+ and labels[0] == b"ns1"
+ )
+
+
+class Ns1ExampleOrgHandler(ReclimitHandler):
+ def __init__(self, state_handler: ReclimitStateHandler) -> None:
+ self._second_query_events: dict[dns.name.Name, asyncio.Event] = {}
+ super().__init__(state_handler)
+
+ def match(self, qctx: QueryContext) -> bool:
+ return is_ns1_example(qctx.qname, "org") and qctx.qtype in (
+ dns.rdatatype.A,
+ dns.rdatatype.AAAA,
+ )
+
+ async def _get_counted_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[ResponseAction, None]:
+ ns_number = int(qctx.qname.labels[1])
+ next_ns_number = ns_number + 1
+ if not self._state.limit or (
+ not self._state.indirect_send_response
+ and next_ns_number <= self._state.limit
+ ):
+ qctx.response.authority.append(
+ ns(f"{ns_number}.example.org.", f"ns1.{next_ns_number}.example.org.")
+ )
+ qctx.response.flags &= ~dns.flags.AA
+ else:
+ self._state.indirect_send_response = True
+ if qctx.qtype == dns.rdatatype.A:
+ qctx.response.answer.append(a(qctx.qname, 4))
+
+ second_query_event = self._second_query_events.get(qctx.qname)
+ if second_query_event is not None:
+ # Second query arrived, release the first response.
+ second_query_event.set()
+ await asyncio.sleep(0) # Yield to allow the first response to be sent.
+ yield DnsResponseSend(qctx.response)
+ else:
+ # Delay the response until the second query for the same QNAME
+ # arrives; give up waiting after 500 ms.
+ second_query_event = asyncio.Event()
+ self._second_query_events[qctx.qname] = second_query_event
+ try:
+ await asyncio.wait_for(second_query_event.wait(), timeout=0.5)
+ except asyncio.TimeoutError:
+ pass
+ finally:
+ del self._second_query_events[qctx.qname]
+ yield DnsResponseSend(qctx.response)
+
+
+class FallbackNxdomainHandler(ReclimitHandler):
+ async def _get_counted_responses(
+ self, qctx: QueryContext
+ ) -> AsyncGenerator[DnsResponseSend, None]:
+ qctx.response.set_rcode(dns.rcode.NXDOMAIN)
+ yield DnsResponseSend(qctx.response)
projects / thirdparty / bind9.git / commitdiff
