When decoding from wire format, if a message has TC set, raise a Truncated

exception.  [Issue #297]

diff --git a/dns/message.py b/dns/message.py
index a5ed5239e20d9fd424cd4c4464c94476980b6a57..59405e7102ba71bb7cb9f293bb29798b2ea7897a 100644 (file)
--- a/dns/message.py
+++ b/dns/message.py
@@ -66,6 +66,10 @@ class UnknownTSIGKey(dns.exception.DNSException):
     """A TSIG with an unknown key was received."""
 
 
+class Truncated(dns.exception.DNSException):
+    """The truncated flag is set."""
+
+
 #: The question section number
 QUESTION = 0
 
@@ -741,6 +745,8 @@ class _WireReader(object):
         self.current = 12
         if dns.opcode.is_update(self.message.flags):
             self.updating = True
+        if self.message.flags & dns.flags.TC:
+            raise Truncated
         self._get_question(qcount)
         if self.question_only:
             return
@@ -804,6 +810,8 @@ def from_wire(wire, keyring=None, request_mac=b'', xfr=False, origin=None,
     Raises ``dns.message.BadTSIG`` if a TSIG record was not the last
     record of the additional data section.
 
+    Raises ``dns.message.Truncated`` if the TC flag is set.
+
     Returns a ``dns.message.Message``.
     """
 

diff --git a/dns/resolver.py b/dns/resolver.py
index dc40efa797a75c18852d1bb094d9955ea2b54468..cdf775745684c98630edc92a28a8c1235cdcb6a5 100644 (file)
--- a/dns/resolver.py
+++ b/dns/resolver.py
@@ -905,11 +905,13 @@ class Resolver(object):
                                                      source=source,
                                                      source_port=source_port)
                         else:
-                            response = dns.query.udp(request, nameserver,
-                                                     timeout, port,
-                                                     source=source,
-                                                     source_port=source_port)
-                            if response.flags & dns.flags.TC:
+                            try:
+                                response = dns.query.udp(request, nameserver,
+                                                         timeout, port,
+                                                         source=source,
+                                                         source_port=\
+                                                         source_port)
+                            except dns.message.Truncated:
                                 # Response truncated; retry with TCP.
                                 tcp_attempt = True
                                 timeout = self._compute_timeout(start, lifetime)

diff --git a/tests/test_message.py b/tests/test_message.py
index d60e711c5b0fc80e7a6de69a14d09009745522c8..9de11c9fac63ab560fa764078378c2c1cd75d710 100644 (file)
--- a/tests/test_message.py
+++ b/tests/test_message.py
@@ -208,5 +208,21 @@ class MessageTestCase(unittest.TestCase):
                             dns.rdatatype.SOA)
         self.failUnless(rrs1 == rrs2)
 
+    def test_CleanTruncated(self):
+        def bad():
+            a = dns.message.from_text(answer_text)
+            a.flags |= dns.flags.TC
+            wire = a.to_wire(want_shuffle=False)
+            dns.message.from_wire(wire)
+        self.failUnlessRaises(dns.message.Truncated, bad)
+
+    def test_MessyTruncated(self):
+        def bad():
+            a = dns.message.from_text(answer_text)
+            a.flags |= dns.flags.TC
+            wire = a.to_wire(want_shuffle=False)
+            dns.message.from_wire(wire[:-3])
+        self.failUnlessRaises(dns.message.Truncated, bad)
+
 if __name__ == '__main__':
     unittest.main()