From 8715462f471eac5916e6cfe33f54ae8f5613b503 Mon Sep 17 00:00:00 2001
From: Szymon Marczak <36894700+szmarczak@users.noreply.github.com>
Date: Wed, 11 Nov 2020 22:22:41 +0100
Subject: [PATCH] dns: add a cancel() method to the promise Resolver
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR-URL: https://github.com/nodejs/node/pull/33099
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: Colin Ihrig <cjihrig@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Juan José Arboleda <soyjuanarbol@gmail.com>
Reviewed-By: Antoine du Hamel <duhamelantoine1995@gmail.com>
---
 doc/api/dns.md                                |  8 +++
 lib/internal/dns/promises.js                  |  1 +
 .../test-dns-channel-cancel-promise.js        | 65 +++++++++++++++++++
 test/parallel/test-dns-channel-cancel.js      | 49 ++++++++++----
 4 files changed, 110 insertions(+), 13 deletions(-)
 create mode 100644 test/parallel/test-dns-channel-cancel-promise.js

diff --git a/doc/api/dns.md b/doc/api/dns.md
index 3037f5fcf8a793..8aa8662d6bdc16 100644
--- a/doc/api/dns.md
+++ b/doc/api/dns.md
@@ -716,6 +716,14 @@ The following methods from the `dnsPromises` API are available:
 * [`resolver.reverse()`][`dnsPromises.reverse()`]
 * [`resolver.setServers()`][`dnsPromises.setServers()`]
 
+### `resolver.cancel()`
+<!-- YAML
+added: REPLACEME
+-->
+
+Cancel all outstanding DNS queries made by this resolver. The corresponding
+promises will be rejected with an error with code `ECANCELLED`.
+
 ### `dnsPromises.getServers()`
 <!-- YAML
 added: v10.6.0
diff --git a/lib/internal/dns/promises.js b/lib/internal/dns/promises.js
index 5d7ecaab1c42b8..42ba76dda4f0ac 100644
--- a/lib/internal/dns/promises.js
+++ b/lib/internal/dns/promises.js
@@ -217,6 +217,7 @@ class Resolver {
 
 Resolver.prototype.getServers = CallbackResolver.prototype.getServers;
 Resolver.prototype.setServers = CallbackResolver.prototype.setServers;
+Resolver.prototype.cancel = CallbackResolver.prototype.cancel;
 Resolver.prototype.setLocalAddress = CallbackResolver.prototype.setLocalAddress;
 Resolver.prototype.resolveAny = resolveMap.ANY = resolver('queryAny');
 Resolver.prototype.resolve4 = resolveMap.A = resolver('queryA');
diff --git a/test/parallel/test-dns-channel-cancel-promise.js b/test/parallel/test-dns-channel-cancel-promise.js
new file mode 100644
index 00000000000000..382ac3dd508d68
--- /dev/null
+++ b/test/parallel/test-dns-channel-cancel-promise.js
@@ -0,0 +1,65 @@
+'use strict';
+const common = require('../common');
+const { promises: dnsPromises } = require('dns');
+const assert = require('assert');
+const dgram = require('dgram');
+
+const server = dgram.createSocket('udp4');
+const resolver = new dnsPromises.Resolver();
+
+const addMessageListener = () => {
+  server.removeAllListeners('message');
+
+  server.once('message', () => {
+    server.once('message', common.mustNotCall);
+
+    resolver.cancel();
+  });
+};
+
+server.bind(0, common.mustCall(async () => {
+  resolver.setServers([`127.0.0.1:${server.address().port}`]);
+
+  addMessageListener();
+
+  // Single promise
+  {
+    const hostname = 'example0.org';
+
+    await assert.rejects(
+      resolver.resolve4(hostname),
+      {
+        code: 'ECANCELLED',
+        syscall: 'queryA',
+        hostname
+      }
+    );
+  }
+
+  addMessageListener();
+
+  // Multiple promises
+  {
+    const assertions = [];
+    const assertionCount = 10;
+
+    for (let i = 1; i <= assertionCount; i++) {
+      const hostname = `example${i}.org`;
+
+      assertions.push(
+        assert.rejects(
+          resolver.resolve4(hostname),
+          {
+            code: 'ECANCELLED',
+            syscall: 'queryA',
+            hostname: hostname
+          }
+        )
+      );
+    }
+
+    await Promise.all(assertions);
+  }
+
+  server.close();
+}));
diff --git a/test/parallel/test-dns-channel-cancel.js b/test/parallel/test-dns-channel-cancel.js
index 0c59ee53161a5e..f92fb2e30a4d12 100644
--- a/test/parallel/test-dns-channel-cancel.js
+++ b/test/parallel/test-dns-channel-cancel.js
@@ -1,6 +1,5 @@
 'use strict';
 const common = require('../common');
-const dnstools = require('../common/dns');
 const { Resolver } = require('dns');
 const assert = require('assert');
 const dgram = require('dgram');
@@ -8,21 +7,45 @@ const dgram = require('dgram');
 const server = dgram.createSocket('udp4');
 const resolver = new Resolver();
 
-server.bind(0, common.mustCall(() => {
+const desiredQueries = 11;
+let finishedQueries = 0;
+
+const addMessageListener = () => {
+  server.removeAllListeners('message');
+
+  server.once('message', () => {
+    server.once('message', common.mustNotCall);
+
+    resolver.cancel();
+  });
+};
+
+server.bind(0, common.mustCall(async () => {
   resolver.setServers([`127.0.0.1:${server.address().port}`]);
-  resolver.resolve4('example.org', common.mustCall((err, res) => {
+
+  const callback = common.mustCall((err, res) => {
     assert.strictEqual(err.code, 'ECANCELLED');
     assert.strictEqual(err.syscall, 'queryA');
-    assert.strictEqual(err.hostname, 'example.org');
-    server.close();
-  }));
-}));
+    assert.strictEqual(err.hostname, `example${finishedQueries}.org`);
+
+    finishedQueries++;
+    if (finishedQueries === desiredQueries) {
+      server.close();
+    }
+  }, desiredQueries);
+
+  const next = (...args) => {
+    callback(...args);
+
+    addMessageListener();
 
-server.on('message', common.mustCall((msg, { address, port }) => {
-  const parsed = dnstools.parseDNSPacket(msg);
-  const domain = parsed.questions[0].domain;
-  assert.strictEqual(domain, 'example.org');
+    // Multiple queries
+    for (let i = 1; i < desiredQueries; i++) {
+      resolver.resolve4(`example${i}.org`, callback);
+    }
+  };
 
-  // Do not send a reply.
-  resolver.cancel();
+  // Single query
+  addMessageListener();
+  resolver.resolve4('example0.org', next);
 }));