From 2ab3d1628c9ae0545e225522b3b445c3478dc6ad Mon Sep 17 00:00:00 2001
From: Bob Halley <halley@dnspython.org>
Date: Sun, 18 Feb 2024 10:27:43 -0800
Subject: [PATCH] The Tudoor fix should not eat valid Truncated exceptions
 [#1053] (#1054)

* The Tudoor fix should not eat valid Truncated exceptions [##1053]

* Make logic more readable

Conflict: delete tests, because no function about mock
Reference:https://github.com/rthalley/dnspython/commit/2ab3d1628c9ae0545e225522b3b445c3478dc6ad
---
 dns/asyncquery.py   | 10 ++++++++
 dns/query.py        | 14 +++++++++++
 2 files changed, 24 insertions(+), 0 deletions(-)

diff --git a/dns/asyncquery.py b/dns/asyncquery.py
index 94cb2413..4d9ab9ae 100644
--- a/dns/asyncquery.py
+++ b/dns/asyncquery.py
@@ -151,6 +151,16 @@ async def receive_udp(
                 ignore_trailing=ignore_trailing,
                 raise_on_truncation=raise_on_truncation,
             )
+        except dns.message.Truncated as e:
+            # See the comment in query.py for details.
+            if (
+                ignore_errors
+                and query is not None
+                and not query.is_response(e.message())
+            ):
+                continue
+            else:
+                raise
         except Exception:
             if ignore_errors:
                 continue
diff --git a/dns/query.py b/dns/query.py
index 06d186c7..384bf31e 100644
--- a/dns/query.py
+++ b/dns/query.py
@@ -618,6 +618,20 @@ def receive_udp(
                 ignore_trailing=ignore_trailing,
                 raise_on_truncation=raise_on_truncation,
             )
+        except dns.message.Truncated as e:
+            # If we got Truncated and not FORMERR, we at least got the header with TC
+            # set, and very likely the question section, so we'll re-raise if the
+            # message seems to be a response as we need to know when truncation happens.
+            # We need to check that it seems to be a response as we don't want a random
+            # injected message with TC set to cause us to bail out.
+            if (
+                ignore_errors
+                and query is not None
+                and not query.is_response(e.message())
+            ):
+                continue
+            else:
+                raise
         except Exception:
             if ignore_errors:
                 continue