Add TransactionSetup and TransactionLimiter.

TransactionSetup allows a setup function to be called when making
transactions.  This allows transaction hooks to be installed.

TransactionLimiter provides a TransactionSetup that can be used
to limit the number of changes in a zone when loading via a file
or zone transfer.

Thanks to github users SaidBySolo and Mohammed Adnan for reporting
issues with unbounded zone size if the zone source was not trusted.

diff --git a/dns/transaction.py b/dns/transaction.py
index 222b6757b1537d5618cf18058776f6f870f8e73a..538daa4ebe3d774832494c84949334235502db39 100644 (file)
--- a/dns/transaction.py
+++ b/dns/transaction.py
@@ -651,3 +651,73 @@ class Transaction:
     def _origin_information(self):
         # This is only used by _add()
         return self.manager.origin_information()
+
+
+class TransactionSetup:
+    """Abstract base class for additional transaction setup.
+
+    In code which supports it, the setup() method is called on writable transactions
+    just after the transaction is created and before it is used on zone data.  This
+    allows the caller to set additional transaction attributes, e.g. checking.
+    """
+
+    def __init__(self):
+        pass
+
+    def setup(self, txn: Transaction):
+        pass
+
+
+class TooManyChanges(dns.exception.DNSException):
+    """Too many changes"""
+
+
+class TransactionLimiter(TransactionSetup):
+    """Transaction setup that enforces a maximum number of changes in the transaction.
+
+    Each rdataset put or deleted counts as a change.  Each whole name deletion counts
+    as a change.
+
+    The limiter is meant to ensure that reading a zonefile from an untrusted source
+    cannot cause too many changes, e.g. via large $GENERATE statements.  It is not
+    meant to limit the total number of records in a zone or the memory used by the
+    zone.
+    """
+
+    def __init__(self, limit: int):
+        super().__init__()
+        self.limit = limit
+        self.changes: int = 0
+
+    def _check_limit(self):
+        if self.changes >= self.limit:
+            raise TooManyChanges(f"limit is {self.limit} changes")
+
+    def _check_put_rdataset(
+        self, txn: Transaction, name: dns.name.Name, rdataset: dns.rdataset.Rdataset
+    ):
+        self.changes += len(rdataset)
+        self._check_limit()
+
+    def _check_delete_rdataset(
+        self,
+        txn: Transaction,
+        name: dns.name.Name,
+        type: dns.rdatatype.RdataType,
+        covers: dns.rdatatype.RdataType | None,
+    ):
+        self.changes += 1
+        self._check_limit()
+
+    def _check_delete_name(
+        self,
+        txn: Transaction,
+        name: dns.name.Name,
+    ):
+        self.changes += 1
+        self._check_limit()
+
+    def setup(self, txn: Transaction):
+        txn.check_put_rdataset(self._check_put_rdataset)
+        txn.check_delete_rdataset(self._check_delete_rdataset)
+        txn.check_delete_name(self._check_delete_name)

diff --git a/dns/xfr.py b/dns/xfr.py
index 4550fa6fa9b2e1fb8ab328744e9b0df012a5d41f..711f280b46f941951a9fc56f07a718c9fd717371 100644 (file)
--- a/dns/xfr.py
+++ b/dns/xfr.py
@@ -61,6 +61,7 @@ class Inbound:
         rdtype: dns.rdatatype.RdataType = dns.rdatatype.AXFR,
         serial: int | None = None,
         is_udp: bool = False,
+        transaction_setup: dns.transaction.TransactionSetup | None = None,
     ):
         """Initialize an inbound zone transfer.
 
@@ -73,6 +74,12 @@ class Inbound:
 
         :param is_udp: Whether UDP is being used for this XFR.
         :type is_udp: bool
+
+        :param transaction_setup: If not ``None``, call the object's setup()
+            method with the just-created writer transaction.  This lets the caller
+            alter the transaction's configuration before it is used, for example adding
+            checking policies.
+        :type transaction_setup: None or dns.transaction.TransactionSetup
         """
         self.txn_manager = txn_manager
         self.txn: dns.transaction.Transaction | None = None
@@ -97,6 +104,7 @@ class Inbound:
         self.done = False
         self.expecting_SOA = False
         self.delete_mode = False
+        self.transaction_setup = transaction_setup
 
     def process_message(self, message: dns.message.Message) -> bool:
         """Process one message in the transfer.
@@ -109,6 +117,8 @@ class Inbound:
         """
         if self.txn is None:
             self.txn = self.txn_manager.writer(not self.incremental)
+            if self.transaction_setup is not None:
+                self.transaction_setup.setup(self.txn)
         rcode = message.rcode()
         if rcode != dns.rcode.NOERROR:
             raise TransferError(rcode)
@@ -228,6 +238,8 @@ class Inbound:
                 self.delete_mode = False
                 self.txn.rollback()
                 self.txn = self.txn_manager.writer(True)
+                if self.transaction_setup is not None:
+                    self.transaction_setup.setup(self.txn)
                 #
                 # Note we are falling through into the code below
                 # so whatever rdataset this was gets written.

diff --git a/dns/zone.py b/dns/zone.py
index 638b04c345868f7c064b0b032f1c7fe7284e2aca..744ecc014f388eb7da14a07629561442bb5e3384 100644 (file)
--- a/dns/zone.py
+++ b/dns/zone.py
@@ -1265,6 +1265,7 @@ def _from_text(
     check_origin: bool = True,
     idna_codec: dns.name.IDNACodec | None = None,
     allow_directives: bool | Iterable[str] = True,
+    transaction_setup: dns.transaction.TransactionSetup | None = None,
 ) -> Zone:
     # See the comments for the public APIs from_text() and from_file() for
     # details.
@@ -1277,6 +1278,8 @@ def _from_text(
         filename = "<string>"
     zone = zone_factory(origin, rdclass, relativize=relativize)
     with zone.writer(True) as txn:
+        if transaction_setup is not None:
+            transaction_setup.setup(txn)
         tok = dns.tokenizer.Tokenizer(text, filename, idna_codec=idna_codec)
         reader = dns.zonefile.Reader(
             tok,
@@ -1308,6 +1311,7 @@ def from_text(
     check_origin: bool = True,
     idna_codec: dns.name.IDNACodec | None = None,
     allow_directives: bool | Iterable[str] = True,
+    transaction_setup: dns.transaction.TransactionSetup | None = None,
 ) -> Zone:
     """Build a zone object from a zone file format string.
 
@@ -1340,6 +1344,11 @@ def from_text(
         non-empty iterable, only the listed directives (including the ``$``)
         are allowed.
     :type allow_directives: bool or Iterable[str]
+    :param transaction_setup: If not ``None``, call the object's setup()
+        method with the just-created writer transaction.  This lets the caller
+        alter the transaction's configuration before it is used, for example adding
+        checking policies.
+    :type transaction_setup: None or dns.transaction.TransactionSetup
     :raises dns.zone.NoSOA: if there is no SOA RRset.
     :raises dns.zone.NoNS: if there is no NS RRset.
     :raises KeyError: if there is no origin node.
@@ -1356,6 +1365,7 @@ def from_text(
         check_origin,
         idna_codec,
         allow_directives,
+        transaction_setup,
     )
 
 
@@ -1370,6 +1380,7 @@ def from_file(
     check_origin: bool = True,
     idna_codec: dns.name.IDNACodec | None = None,
     allow_directives: bool | Iterable[str] = True,
+    transaction_setup: dns.transaction.TransactionSetup | None = None,
 ) -> Zone:
     """Read a zone file and build a zone object.
 
@@ -1403,6 +1414,11 @@ def from_file(
         non-empty iterable, only the listed directives (including the ``$``)
         are allowed.
     :type allow_directives: bool or Iterable[str]
+    :param transaction_setup: If not ``None``, call the object's setup()
+        method with the just-created writer transaction.  This lets the caller
+        alter the transaction's configuration before it is used, for example adding
+        checking policies.
+    :type transaction_setup: None or dns.transaction.TransactionSetup
     :raises dns.zone.NoSOA: if there is no SOA RRset.
     :raises dns.zone.NoNS: if there is no NS RRset.
     :raises KeyError: if there is no origin node.
@@ -1427,6 +1443,7 @@ def from_file(
             check_origin,
             idna_codec,
             allow_directives,
+            transaction_setup,
         )
     assert False  # make mypy happy  lgtm[py/unreachable-statement]
 

diff --git a/tests/test_zone.py b/tests/test_zone.py
index c5dfa331a0e22b0dd3e4f14ec06d947546331440..9de721450104852df74daf0fb34e146af5558052 100644 (file)
--- a/tests/test_zone.py
+++ b/tests/test_zone.py
@@ -33,6 +33,7 @@ import dns.rdataclass
 import dns.rdataset
 import dns.rdatatype
 import dns.rrset
+import dns.transaction
 import dns.versioned
 import dns.zone
 from tests.util import here
@@ -1280,6 +1281,17 @@ class ZoneTestCase(unittest.TestCase):
         print(example_unicode_justified)
         self.assertEqual(t1, example_unicode_justified)
 
+    def testFromFileHittingLimit(self):
+        limiter = dns.transaction.TransactionLimiter(20)
+        with self.assertRaises(dns.transaction.TooManyChanges):
+            z = dns.zone.from_file(
+                here("example"), "example", transaction_setup=limiter
+            )
+
+    def testFromFileLimitOk(self):
+        limiter = dns.transaction.TransactionLimiter(1000)
+        z = dns.zone.from_file(here("example"), "example", transaction_setup=limiter)
+
 
 class VersionedZoneTestCase(unittest.TestCase):
     zone_factory = dns.versioned.Zone