Implement new inbound xfr design.xfr

1 files changed, 127 insertions, 10 deletions

diff --git a/dns/query.py b/dns/query.py
index d4a3afa..01452ee 100644
--- a/dns/query.py
+++ b/dns/query.py
@@ -18,6 +18,7 @@
 """Talk to a DNS server."""
 
 import contextlib
+import enum
 import errno
 import os
 import selectors
@@ -35,6 +36,7 @@ import dns.rcode
 import dns.rdataclass
 import dns.rdatatype
 import dns.serial
+import dns.xfr
 
 try:
     import requests
@@ -73,20 +75,15 @@ class BadResponse(dns.exception.FormError):
     """A DNS query response does not respond to the question asked."""
 
 
-class TransferError(dns.exception.DNSException):
-    """A zone transfer response got a non-zero rcode."""
-
-    def __init__(self, rcode):
-        message = 'Zone transfer error: %s' % dns.rcode.to_text(rcode)
-        super().__init__(message)
-        self.rcode = rcode
-
-
 class NoDOH(dns.exception.DNSException):
     """DNS over HTTPS (DOH) was requested but the requests module is not
     available."""
 
 
+# for backwards compatibility
+TransferError = dns.xfr.TransferError
+
+
 def _compute_times(timeout):
     now = time.time()
     if timeout is None:
@@ -917,7 +914,7 @@ def xfr(where, zone, rdtype=dns.rdatatype.AXFR, rdclass=dns.rdataclass.IN,
                (expiration is not None and mexpiration > expiration):
                 mexpiration = expiration
             if use_udp:
-                (wire, _) = _udp_recv(s, 65535, expiration)
+                (wire, _) = _udp_recv(s, 65535, mexpiration)
             else:
                 ldata = _net_read(s, 2, mexpiration)
                 (l,) = struct.unpack("!H", ldata)
@@ -984,3 +981,123 @@ def xfr(where, zone, rdtype=dns.rdatatype.AXFR, rdclass=dns.rdataclass.IN,
             if done and q.keyring and not r.had_tsig:
                 raise dns.exception.FormError("missing TSIG")
             yield r
+
+
+class UDPMode(enum.IntEnum):
+    """How should UDP be used in an IXFR from :py:func:`inbound_xfr()`?
+
+    NEVER means "never use UDP; always use TCP"
+    TRY_FIRST means "try to use UDP but fall back to TCP if needed"
+    ONLY means "raise ``dns.xfr.UseTCP`` if trying UDP does not succeed"
+    """
+    NEVER = 0
+    TRY_FIRST = 1
+    ONLY = 2
+
+
+def inbound_xfr(where, txn_manager, query=None,
+                port=53, timeout=None, lifetime=None, source=None,
+                source_port=0, udp_mode=UDPMode.NEVER,
+                keyring=None, keyname=None,
+                keyalgorithm=dns.tsig.default_algorithm):
+    """Conduct an inbound transfer and apply it via a transaction from the
+    txn_manager.
+
+    *where*, a ``str`` containing an IPv4 or IPv6 address,  where
+    to send the message.
+
+    *txn_manager*, a ``dns.transaction.TransactionManager``, the txn_manager
+    for this transfer (typically a ``dns.zone.Zone``).
+
+    *query*, the query to send.  If not supplied, a default query is
+    constructed using information from the *txn_manager*.
+
+    *port*, an ``int``, the port send the message to.  The default is 53.
+
+    *timeout*, a ``float``, the number of seconds to wait for each
+    response message.  If None, the default, wait forever.
+
+    *lifetime*, a ``float``, the total number of seconds to spend
+    doing the transfer.  If ``None``, the default, then there is no
+    limit on the time the transfer may take.
+
+    *source*, a ``str`` containing an IPv4 or IPv6 address, specifying
+    the source address.  The default is the wildcard address.
+
+    *source_port*, an ``int``, the port from which to send the message.
+    The default is 0.
+
+    *udp_mode*, a ``dns.query.UDPMode``, determines how UDP is used
+    for IXFRs.  The default is ``dns.UDPMode.NEVER``, i.e. only use
+    TCP.  Other possibilites are ``dns.UDPMode.TRY_FIRST``, which
+    means "try UDP but fallback to TCP if needed", and
+    ``dns.UDPMode.ONLY``, which means "try UDP and raise
+    ``dns.xfr.UseTCP`` if it does not succeeed.
+
+    *keyring*, a ``dict``, the keyring to use for TSIG.
+
+    *keyname*, a ``dns.name.Name`` or ``str``, the name of the TSIG
+    key to use.
+
+    *keyalgorithm*, a ``dns.name.Name`` or ``str``, the TSIG algorithm to use.
+
+    Raises on errors.
+
+    """
+    if query is None:
+        (query, serial) = dns.xfr.make_query(txn_manager)
+    rdtype = query.question[0].rdtype
+    is_ixfr = rdtype == dns.rdatatype.IXFR
+    origin = txn_manager.from_wire_origin()
+    wire = query.to_wire()
+    (af, destination, source) = _destination_and_source(where, port,
+                                                        source, source_port)
+    (_, expiration) = _compute_times(lifetime)
+    retry = True
+    while retry:
+        retry = False
+        if is_ixfr and udp_mode != UDPMode.NEVER:
+            sock_type = socket.SOCK_DGRAM
+            is_udp = True
+        else:
+            sock_type = socket.SOCK_STREAM
+            is_udp = False
+        with _make_socket(af, sock_type, source) as s:
+            _connect(s, destination, expiration)
+            if is_udp:
+                _udp_send(s, wire, None, expiration)
+            else:
+                tcpmsg = struct.pack("!H", len(wire)) + wire
+                _net_write(s, tcpmsg, expiration)
+            with dns.xfr.Inbound(txn_manager, rdtype, serial) as inbound:
+                done = False
+                tsig_ctx = None
+                while not done:
+                    (_, mexpiration) = _compute_times(timeout)
+                    if mexpiration is None or \
+                       (expiration is not None and mexpiration > expiration):
+                        mexpiration = expiration
+                    if is_udp:
+                        (rwire, _) = _udp_recv(s, 65535, mexpiration)
+                    else:
+                        ldata = _net_read(s, 2, mexpiration)
+                        (l,) = struct.unpack("!H", ldata)
+                        rwire = _net_read(s, l, mexpiration)
+                    r = dns.message.from_wire(rwire, keyring=query.keyring,
+                                              request_mac=query.mac, xfr=True,
+                                              origin=origin, tsig_ctx=tsig_ctx,
+                                              multi=(not is_udp),
+                                              one_rr_per_rrset=is_ixfr)
+                    try:
+                        done = inbound.process_message(r, is_udp)
+                    except dns.xfr.UseTCP:
+                        assert is_udp  # should not happen if we used TCP!
+                        if udp_mode == UDPMode.ONLY:
+                            raise
+                        done = True
+                        retry = True
+                        udp_mode = UDPMode.NEVER
+                        continue
+                    tsig_ctx = r.tsig_ctx
+                if not retry and query.keyring and not r.had_tsig:
+                    raise dns.exception.FormError("missing TSIG")