From 2944ea6c6b42c19b49c7a2e13441c4cc28e76529 Mon Sep 17 00:00:00 2001
From: cjihrig <cjihrig@gmail.com>
Date: Mon, 31 Jul 2017 20:38:06 -0400
Subject: [PATCH] dgram: add custom lookup function in sockets

This commit adds support for custom DNS lookup functions in
dgram sockets. This is similar to the existing feature in net
sockets.

Refs: https://github.com/nodejs/node/issues/6189
PR-URL: https://github.com/nodejs/node/pull/14560
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Wyatt Preul <wpreul@gmail.com>
Reviewed-By: Evan Lucas <evanlucas@me.com>
---
 doc/api/dgram.md                          | 38 ++++++++++--------
 lib/dgram.js                              | 32 ++++++++-------
 test/parallel/test-dgram-custom-lookup.js | 47 +++++++++++++++++++++++
 3 files changed, 86 insertions(+), 31 deletions(-)
 create mode 100644 test/parallel/test-dgram-custom-lookup.js

diff --git a/doc/api/dgram.md b/doc/api/dgram.md
index f2968438b183c6..3a9864ab844655 100644
--- a/doc/api/dgram.md
+++ b/doc/api/dgram.md
@@ -535,27 +535,30 @@ s.bind(1234, () => {
 ### dgram.createSocket(options[, callback])
 <!-- YAML
 added: v0.11.13
+changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/14560
+    description: The `lookup` option is supported.
 -->
 
-* `options` {Object}
-* `callback` {Function} Attached as a listener to `'message'` events.
+* `options` {Object} Available options are:
+  * `type` {string} The family of socket. Must be either `'udp4'` or `'udp6'`.
+    Required.
+  * `reuseAddr` {boolean} When `true` [`socket.bind()`][] will reuse the
+    address, even if another process has already bound a socket on it. Optional.
+    Defaults to `false`.
+  * `lookup` {Function} Custom lookup function. Defaults to [`dns.lookup()`][].
+    Optional.
+* `callback` {Function} Attached as a listener for `'message'` events. Optional.
 * Returns: {dgram.Socket}
 
-Creates a `dgram.Socket` object. The `options` argument is an object that
-should contain a `type` field of either `udp4` or `udp6` and an optional
-boolean `reuseAddr` field.
-
-When `reuseAddr` is `true` [`socket.bind()`][] will reuse the address, even if
-another process has already bound a socket on it. `reuseAddr` defaults to
-`false`. The optional `callback` function is added as a listener for `'message'`
-events.
-
-Once the socket is created, calling [`socket.bind()`][] will instruct the
-socket to begin listening for datagram messages. When `address` and `port` are
-not passed to  [`socket.bind()`][] the method will bind the socket to the "all
-interfaces" address on a random port (it does the right thing for both `udp4`
-and `udp6` sockets). The bound address and port can be retrieved using
-[`socket.address().address`][] and [`socket.address().port`][].
+Creates a `dgram.Socket` object. Once the socket is created, calling
+[`socket.bind()`][] will instruct the socket to begin listening for datagram
+messages. When `address` and `port` are not passed to  [`socket.bind()`][] the
+method will bind the socket to the "all interfaces" address on a random port
+(it does the right thing for both `udp4` and `udp6` sockets). The bound address
+and port can be retrieved using [`socket.address().address`][] and
+[`socket.address().port`][].
 
 ### dgram.createSocket(type[, callback])
 <!-- YAML
@@ -585,6 +588,7 @@ and `udp6` sockets). The bound address and port can be retrieved using
 [`cluster`]: cluster.html
 [`dgram.Socket#bind()`]: #dgram_socket_bind_options_callback
 [`dgram.createSocket()`]: #dgram_dgram_createsocket_options_callback
+[`dns.lookup()`]: dns.html#dns_dns_lookup_hostname_options_callback
 [`socket.address().address`]: #dgram_socket_address
 [`socket.address().port`]: #dgram_socket_address
 [`socket.bind()`]: #dgram_socket_bind_port_address_callback
diff --git a/lib/dgram.js b/lib/dgram.js
index 3221784d533c3d..7c73c713e21abe 100644
--- a/lib/dgram.js
+++ b/lib/dgram.js
@@ -46,31 +46,32 @@ var cluster = null;
 const errnoException = util._errnoException;
 const exceptionWithHostPort = util._exceptionWithHostPort;
 
-function lookup(address, family, callback) {
-  return dns.lookup(address, family, callback);
-}
-
 
-function lookup4(address, callback) {
+function lookup4(lookup, address, callback) {
   return lookup(address || '127.0.0.1', 4, callback);
 }
 
 
-function lookup6(address, callback) {
+function lookup6(lookup, address, callback) {
   return lookup(address || '::1', 6, callback);
 }
 
 
-function newHandle(type) {
+function newHandle(type, lookup) {
+  if (lookup === undefined)
+    lookup = dns.lookup;
+  else if (typeof lookup !== 'function')
+    throw new errors.TypeError('ERR_INVALID_ARG_TYPE', 'lookup', 'function');
+
   if (type === 'udp4') {
     const handle = new UDP();
-    handle.lookup = lookup4;
+    handle.lookup = lookup4.bind(handle, lookup);
     return handle;
   }
 
   if (type === 'udp6') {
     const handle = new UDP();
-    handle.lookup = lookup6;
+    handle.lookup = lookup6.bind(handle, lookup);
     handle.bind = handle.bind6;
     handle.send = handle.send6;
     return handle;
@@ -100,13 +101,15 @@ function _createSocketHandle(address, port, addressType, fd, flags) {
 
 function Socket(type, listener) {
   EventEmitter.call(this);
+  var lookup;
 
   if (type !== null && typeof type === 'object') {
     var options = type;
     type = options.type;
+    lookup = options.lookup;
   }
 
-  var handle = newHandle(type);
+  var handle = newHandle(type, lookup);
   handle.owner = this;
 
   this._handle = handle;
@@ -186,10 +189,11 @@ Socket.prototype.bind = function(port_, address_ /*, callback*/) {
   }
 
   // defaulting address for bind to all interfaces
-  if (!address && this._handle.lookup === lookup4) {
-    address = '0.0.0.0';
-  } else if (!address && this._handle.lookup === lookup6) {
-    address = '::';
+  if (!address) {
+    if (this.type === 'udp4')
+      address = '0.0.0.0';
+    else
+      address = '::';
   }
 
   // resolve address first
diff --git a/test/parallel/test-dgram-custom-lookup.js b/test/parallel/test-dgram-custom-lookup.js
new file mode 100644
index 00000000000000..ec6cef9e2586af
--- /dev/null
+++ b/test/parallel/test-dgram-custom-lookup.js
@@ -0,0 +1,47 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const dgram = require('dgram');
+const dns = require('dns');
+
+{
+  // Verify that the provided lookup function is called.
+  const lookup = common.mustCall((host, family, callback) => {
+    dns.lookup(host, family, callback);
+  });
+
+  const socket = dgram.createSocket({ type: 'udp4', lookup });
+
+  socket.bind(common.mustCall(() => {
+    socket.close();
+  }));
+}
+
+{
+  // Verify that lookup defaults to dns.lookup().
+  const originalLookup = dns.lookup;
+
+  dns.lookup = common.mustCall((host, family, callback) => {
+    dns.lookup = originalLookup;
+    originalLookup(host, family, callback);
+  });
+
+  const socket = dgram.createSocket({ type: 'udp4' });
+
+  socket.bind(common.mustCall(() => {
+    socket.close();
+  }));
+}
+
+{
+  // Verify that non-functions throw.
+  [null, true, false, 0, 1, NaN, '', 'foo', {}, Symbol()].forEach((value) => {
+    assert.throws(() => {
+      dgram.createSocket({ type: 'udp4', lookup: value });
+    }, common.expectsError({
+      code: 'ERR_INVALID_ARG_TYPE',
+      type: TypeError,
+      message: 'The "lookup" argument must be of type function'
+    }));
+  });
+}