From be8ec6ef183fae7223a9f0f9cb8b9283bca5ab6f Mon Sep 17 00:00:00 2001
From: Aram Sargsyan <aram@isc.org>
Date: Wed, 15 Jun 2022 12:57:14 +0000
Subject: [PATCH] Fix DiG UDP query retry and fail-over bug

When the `udp_ready()` callback function gets called with a failure
result code, DiG erroneously cancels the lookup.

Copy the logic behind `tcp_connected()` callback function into
`udp_ready()` so that DiG will now retry the failed query (if retries
are enabled) and then, if it fails again, it will fail-over to the next
server in the list, which synchronizes the behavior between TCP and UDP
modes.

Also, `udp_ready()` was calling `lookup_detach()` without calling
`lookup_attach()` first, but the issue was masked behind the fact
that `clear_current_lookup()` wasn't being called when needed, and
`lookup_detach()` was compensating for that. This also has been fixed.

(cherry picked from commit 3f3108552577c326b4dab6c3b631c51cf0040144)
---
 bin/dig/dighost.c | 48 ++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 5 deletions(-)

diff --git a/bin/dig/dighost.c b/bin/dig/dighost.c
index ca3a406232..a9a4b0f388 100644
--- a/bin/dig/dighost.c
+++ b/bin/dig/dighost.c
@@ -3112,6 +3112,9 @@ send_udp(dig_query_t *query) {
 static void
 udp_ready(isc_nmhandle_t *handle, isc_result_t eresult, void *arg) {
 	dig_query_t *query = (dig_query_t *)arg;
+	dig_query_t *next = NULL;
+	char sockstr[ISC_SOCKADDR_FORMATSIZE];
+	dig_lookup_t *l = NULL;
 	dig_query_t *readquery = NULL;
 	int local_timeout = timeout * 1000;
 
@@ -3132,30 +3135,63 @@ udp_ready(isc_nmhandle_t *handle, isc_result_t eresult, void *arg) {
 	debug("udp_ready(%p, %s, %p)", handle, isc_result_totext(eresult),
 	      query);
 
-	if (eresult == ISC_R_CANCELED || query->canceled) {
-		dig_lookup_t *l = query->lookup;
+	LOCK_LOOKUP;
+	lookup_attach(query->lookup, &l);
 
+	if (eresult == ISC_R_CANCELED || query->canceled) {
 		debug("in cancel handler");
 		if (!query->canceled) {
 			cancel_lookup(l);
 		}
 		query_detach(&query);
 		lookup_detach(&l);
+		clear_current_lookup();
+		UNLOCK_LOOKUP;
 		return;
 	} else if (eresult != ISC_R_SUCCESS) {
-		dig_lookup_t *l = query->lookup;
-
 		debug("udp setup failed: %s", isc_result_totext(eresult));
+		isc_sockaddr_format(&query->sockaddr, sockstr, sizeof(sockstr));
+		dighost_warning("UDP setup with %s(%s) for %s failed: %s.",
+				sockstr, query->servname, l->textname,
+				isc_result_totext(eresult));
 
 		if (exitcode < 9) {
 			exitcode = 9;
 		}
+
+		if (l->retries > 1) {
+			l->retries--;
+			debug("making new UDP request, %d tries left",
+			      l->retries);
+			requeue_lookup(l, true);
+			next = NULL;
+		} else if ((l->current_query != NULL) &&
+			   (ISC_LINK_LINKED(l->current_query, link)))
+		{
+			next = ISC_LIST_NEXT(l->current_query, link);
+		} else {
+			next = NULL;
+		}
+
 		query_detach(&query);
-		cancel_lookup(l);
+		if (next == NULL) {
+			cancel_lookup(l);
+		}
 		lookup_detach(&l);
+
+		if (next != NULL) {
+			start_udp(next);
+		} else {
+			clear_current_lookup();
+		}
+
+		check_if_done();
+		UNLOCK_LOOKUP;
 		return;
 	}
 
+	exitcode = 0;
+
 	query_attach(query, &readquery);
 
 	debug("recving with lookup=%p, query=%p, handle=%p", query->lookup,
@@ -3177,6 +3213,8 @@ udp_ready(isc_nmhandle_t *handle, isc_result_t eresult, void *arg) {
 	send_udp(readquery);
 
 	query_detach(&query);
+	lookup_detach(&l);
+	UNLOCK_LOOKUP;
 }
 
 /*%