From 11e4235b1ab0920b1888d09dd8bb9249a82df8ff Mon Sep 17 00:00:00 2001
From: Bob Halley <halley@dnspython.org>
Date: Sat, 5 Jan 2019 10:38:32 -0800
Subject: [PATCH] When decoding from wire format, if a message has TC set,
 raise a Truncated exception.  [Issue #297]

---
 dns/message.py        |  8 ++++++++
 dns/resolver.py       | 12 +++++++-----
 tests/test_message.py | 16 ++++++++++++++++
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/dns/message.py b/dns/message.py
index a5ed5239..59405e71 100644
--- 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 dc40efa7..cdf77574 100644
--- 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 d60e711c..9de11c9f 100644
--- 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()
-- 
2.47.3