From 60441afc56decf21808b8b2767d6769c787f82a5 Mon Sep 17 00:00:00 2001
From: Norman Maurer <norman_maurer@apple.com>
Date: Wed, 17 Jan 2024 13:12:19 +0100
Subject: [PATCH] DnsNameResolver: Fail query if id space is exhausted

Motivation:

When we try to execute a query we will try to select / generate an id that is used. When we are not able to find one we throw an IllegalStateException. In this case we also need to ensure we fail the original promise as otherwise the user might never be notified of the problem.

Modifications:

- Move throw code out of the DnsQueryContextManager to make it easier to reason about
- Fail query with IllegalStateException

Result:

Query will be correctly failed in the case of id space exhausting
---
 .../netty/resolver/dns/DnsQueryContext.java   | 13 ++++++---
 .../resolver/dns/DnsQueryContextManager.java  | 27 +++++++++++++++++--
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContext.java b/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContext.java
index da1091b1d160..348d15b30bd0 100644
--- a/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContext.java
+++ b/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContext.java
@@ -81,7 +81,7 @@ abstract class DnsQueryContext {
 
     private volatile Future<?> timeoutFuture;
 
-    private int id = -1;
+    private int id = Integer.MIN_VALUE;
 
     DnsQueryContext(Channel channel,
                     Future<? extends Channel> channelReadyFuture,
@@ -172,8 +172,15 @@ final DnsQuestion question() {
      * @return                      the {@link ChannelFuture} that is notified once once the write completes.
      */
     final ChannelFuture writeQuery(boolean flush) {
-        assert id == -1 : this.getClass().getSimpleName() + ".writeQuery(...) can only be executed once.";
-        id = queryContextManager.add(nameServerAddr, this);
+        assert id == Integer.MIN_VALUE : this.getClass().getSimpleName() +
+                ".writeQuery(...) can only be executed once.";
+
+        if ((id = queryContextManager.add(nameServerAddr, this)) == -1) {
+            // We did exhaust the id space, fail the query
+            IllegalStateException e = new IllegalStateException("query ID space exhausted: " + question());
+            finishFailure("failed to send a query via " + protocol(), e, false);
+            return channel.newFailedFuture(e);
+        }
 
         // Ensure we remove the id from the QueryContextManager once the query completes.
         promise.addListener(new FutureListener<AddressedEnvelope<DnsResponse, InetSocketAddress>>() {
diff --git a/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContextManager.java b/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContextManager.java
index f82648c43dd2..773eecca0aad 100644
--- a/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContextManager.java
+++ b/resolver-dns/src/main/java/io/netty/resolver/dns/DnsQueryContextManager.java
@@ -38,11 +38,27 @@ final class DnsQueryContextManager {
     private final Map<InetSocketAddress, DnsQueryContextMap> map =
             new HashMap<InetSocketAddress, DnsQueryContextMap>();
 
+    /**
+     * Add {@link DnsQueryContext} to the context manager and return the ID that should be used for the query.
+     * This method will return {@code -1} if an ID could not be generated and the context was not stored.
+     *
+     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver to query.
+     * @param qCtx              The {@link {@link DnsQueryContext} to store.
+     * @return                  the ID that should be used or {@code -1} if none could be generated.
+     */
     int add(InetSocketAddress nameServerAddr, DnsQueryContext qCtx) {
         final DnsQueryContextMap contexts = getOrCreateContextMap(nameServerAddr);
         return contexts.add(qCtx);
     }
 
+    /**
+     * Return the {@link DnsQueryContext} for the given {@link InetSocketAddress} and id or {@code null} if
+     * none could be found.
+     *
+     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver.
+     * @param id                The id that identifies the {@link DnsQueryContext} and was used for the query.
+     * @return                  The context or {@code null} if none could be found.
+     */
     DnsQueryContext get(InetSocketAddress nameServerAddr, int id) {
         final DnsQueryContextMap contexts = getContextMap(nameServerAddr);
         if (contexts == null) {
@@ -51,6 +67,14 @@ DnsQueryContext get(InetSocketAddress nameServerAddr, int id) {
         return contexts.get(id);
     }
 
+    /**
+     * Remove the {@link DnsQueryContext} for the given {@link InetSocketAddress} and id or {@code null} if
+     * none could be found.
+     *
+     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver.
+     * @param id                The id that identifies the {@link DnsQueryContext} and was used for the query.
+     * @return                  The context or {@code null} if none could be removed.
+     */
     DnsQueryContext remove(InetSocketAddress nameServerAddr, int id) {
         final DnsQueryContextMap contexts = getContextMap(nameServerAddr);
         if (contexts == null) {
@@ -150,8 +174,7 @@ synchronized int add(DnsQueryContext ctx) {
 
                 id = id + 1 & 0xFFFF;
                 if (++tries >= MAX_TRIES) {
-                    throw new IllegalStateException(
-                            "query ID space exhausted after " + MAX_TRIES + ": " + ctx.question());
+                    return -1;
                 }
             }
         }
