From a214f2eb9692040cf45e20739a7dd319f47ff8c8 Mon Sep 17 00:00:00 2001
From: Norman Maurer <norman_maurer@apple.com>
Date: Tue, 26 Jun 2018 07:36:23 +0200
Subject: [PATCH] Remove id from DnsQueryContextManager whenever the promise is
 fullfilled.

Motivation:

We did not handle the case when the query was cancelled which could lead to an exhausted id space. Beside this we did not not cancel the timeout on failed promises.

Modifications:

- Do the removal of the id from the manager in a FutureListener so its handled in all cases.
- Cancel the timeout whenever the original promise was full-filled.

Result:

Fixes https://github.com/netty/netty/issues/8013.
---
 .../netty/resolver/dns/DnsQueryContext.java   | 48 ++++++++++---------
 1 file changed, 26 insertions(+), 22 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 3aee7686d3..1805b85c4b 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
@@ -28,6 +28,7 @@ import io.netty.handler.codec.dns.DnsRecord;
 import io.netty.handler.codec.dns.DnsResponse;
 import io.netty.handler.codec.dns.DnsSection;
 import io.netty.util.concurrent.Future;
+import io.netty.util.concurrent.FutureListener;
 import io.netty.util.concurrent.GenericFutureListener;
 import io.netty.util.concurrent.Promise;
 import io.netty.util.concurrent.ScheduledFuture;
@@ -39,7 +40,7 @@ import java.util.concurrent.TimeUnit;
 
 import static io.netty.util.internal.ObjectUtil.checkNotNull;
 
-final class DnsQueryContext {
+final class DnsQueryContext implements FutureListener<AddressedEnvelope<DnsResponse, InetSocketAddress>> {
 
     private static final InternalLogger logger = InternalLoggerFactory.getInstance(DnsQueryContext.class);
 
@@ -68,6 +69,9 @@ final class DnsQueryContext {
         recursionDesired = parent.isRecursionDesired();
         id = parent.queryContextManager.add(this);
 
+        // Ensure we remove the id from the QueryContextManager once the query completes.
+        promise.addListener(this);
+
         if (parent.isOptResourceEnabled()) {
             optResource = new AbstractDnsOptPseudoRrRecord(parent.maxPayloadSize(), 0, 0) {
                 // We may want to remove this in the future and let the user just specify the opt record in the query.
@@ -119,9 +123,6 @@ final class DnsQueryContext {
                     if (future.isSuccess()) {
                         writeQuery(query, writePromise);
                     } else {
-                        // Remove the id from the manager as we fail the query.
-                        parent.queryContextManager.remove(nameServerAddr(), id);
-
                         Throwable cause = future.cause();
                         promise.tryFailure(cause);
                         writePromise.setFailure(cause);
@@ -138,7 +139,7 @@ final class DnsQueryContext {
         } else {
             writeFuture.addListener(new ChannelFutureListener() {
                 @Override
-                public void operationComplete(ChannelFuture future) throws Exception {
+                public void operationComplete(ChannelFuture future) {
                     onQueryWriteCompletion(writeFuture);
                 }
             });
@@ -184,29 +185,18 @@ final class DnsQueryContext {
     }
 
     private void setSuccess(AddressedEnvelope<? extends DnsResponse, InetSocketAddress> envelope) {
-        parent.queryContextManager.remove(nameServerAddr(), id);
-
-        // Cancel the timeout task.
-        final ScheduledFuture<?> timeoutFuture = this.timeoutFuture;
-        if (timeoutFuture != null) {
-            timeoutFuture.cancel(false);
-        }
-
         Promise<AddressedEnvelope<DnsResponse, InetSocketAddress>> promise = this.promise;
-        if (promise.setUncancellable()) {
-            @SuppressWarnings("unchecked")
-            AddressedEnvelope<DnsResponse, InetSocketAddress> castResponse =
-                    (AddressedEnvelope<DnsResponse, InetSocketAddress>) envelope.retain();
-            if (!promise.trySuccess(castResponse)) {
-                // We failed to notify the promise as it was failed before, thus we need to release the envelope
-                envelope.release();
-            }
+        @SuppressWarnings("unchecked")
+        AddressedEnvelope<DnsResponse, InetSocketAddress> castResponse =
+                (AddressedEnvelope<DnsResponse, InetSocketAddress>) envelope.retain();
+        if (!promise.trySuccess(castResponse)) {
+            // We failed to notify the promise as it was failed before, thus we need to release the envelope
+            envelope.release();
         }
     }
 
     private void setFailure(String message, Throwable cause) {
         final InetSocketAddress nameServerAddr = nameServerAddr();
-        parent.queryContextManager.remove(nameServerAddr, id);
 
         final StringBuilder buf = new StringBuilder(message.length() + 64);
         buf.append('[')
@@ -225,4 +215,18 @@ final class DnsQueryContext {
         }
         promise.tryFailure(e);
     }
+
+    @Override
+    public void operationComplete(Future<AddressedEnvelope<DnsResponse, InetSocketAddress>> future) {
+        // Cancel the timeout task.
+        final ScheduledFuture<?> timeoutFuture = this.timeoutFuture;
+        if (timeoutFuture != null) {
+            this.timeoutFuture = null;
+            timeoutFuture.cancel(false);
+        }
+
+        // Remove the id from the manager as soon as the query completes. This may be because of success, failure or
+        // cancellation
+        parent.queryContextManager.remove(nameServerAddr, id);
+    }
 }