From 36d7aaccbb6ebc4f53fd8a5fcae384c1a2017597 Mon Sep 17 00:00:00 2001
From: AJ ONeal <aj@daplie.com>
Date: Mon, 18 Sep 2017 18:06:27 -0600
Subject: [PATCH] udp6 and timeout fix

---
 bin/dig.js     |  9 ++++++++-
 dns-request.js | 37 ++++++++++++++++++++++---------------
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/bin/dig.js b/bin/dig.js
index b17356a..84fc620 100755
--- a/bin/dig.js
+++ b/bin/dig.js
@@ -93,6 +93,10 @@ cli.main(function (args, cli) {
     if (!cli.timeout) {
       cli.timeout = 3000;
     }
+  } else {
+    if (!cli.timeout) {
+      cli.timeout = 5000;
+    }
   }
 
   if (!cli.type) {
@@ -228,7 +232,7 @@ cli.main(function (args, cli) {
 
     if (cli.debug) {
       console.log('');
-      console.log('Bound and Listening:');
+      console.log('Bound and Listening:', server.type);
       console.log(server.address());
     }
 
@@ -256,6 +260,9 @@ cli.main(function (args, cli) {
         console.log('request sent to', res.nameserver);
       }
     }
+  , onTimeout: function (res) {
+      console.log("Timeout after " + res.timeout + "ms: No response received.");
+    }
   , mdns: cli.mdns
   , nameserver: cli.nameserver
   , port: cli.port
diff --git a/dns-request.js b/dns-request.js
index 878ab5e..636e2f3 100644
--- a/dns-request.js
+++ b/dns-request.js
@@ -6,17 +6,31 @@ var dgram = require('dgram');
 
 function request(queryAb, opts) {
   var handlers = {};
+  var nameservers;
+  var nameserver = opts.nameserver;
+  var index;
+  var udpType;
+  var receivedMessage;
+  if (!nameserver) {
+    nameservers = require('dns').getServers();
+    index = crypto.randomBytes(2).readUInt16BE(0) % nameservers.length;
+    nameserver = nameservers[index];
+  }
+  udpType = /:/.test(nameserver) ? 'udp6' : 'udp4';
   var server = dgram.createSocket({
-    type: 'udp4'
+    type: udpType
   , reuseAddr: true
   });
+  server.nameserver = nameserver;
 
   handlers.onError = function (err) {
     if (opts.onError) { opts.onError(err); }
     server.close();
   };
   handlers.onMessage = function (bin) {
+    receivedMessage = true;
     if (!opts.mdns) {
+      clearTimeout(server._timeoutToken);
       server.close();
     }
 
@@ -25,19 +39,10 @@ function request(queryAb, opts) {
   handlers.onListening = function () {
     /*jshint validthis:true*/
     var server = this;
-    var nameservers;
-    var index;
-    server.nameserver = opts.nameserver;
 
-    if (!server.nameserver) {
-      nameservers = require('dns').getServers();
-      index = crypto.randomBytes(2).readUInt16BE(0) % nameservers.length;
-      server.nameserver = nameservers[index];
-    }
-
-    if (opts.mdns || '224.0.0.251' === opts.nameserver) {
+    if (opts.mdns || '224.0.0.251' === server.nameserver) {
       server.setBroadcast(true);
-      server.addMembership(opts.nameserver);
+      server.addMembership(server.nameserver || '224.0.0.251');
     }
 
     if (opts.onListening) { opts.onListening.apply(server); }
@@ -56,13 +61,15 @@ function request(queryAb, opts) {
   // 5353 mdns
   if (opts.mdns) {
     server.bind(opts.port /*5353*/);
-    setTimeout(function () {
-      server.close();
-    }, opts.timeout || (5 * 1000));
   }
   else {
     server.bind(0);
   }
+  var ms = opts.timeout || (5 * 1000);
+  server._timeoutToken = setTimeout(function () {
+    if (!receivedMessage && opts.onTimeout) { opts.onTimeout({ timeout: ms }); }
+    server.close();
+  }, ms);
 }
 
 function requestJson(query, opts) {