refactor: allow passing in a custom dns lookup function

tediousjs · May 28, 2020 · 4e07f7c · 4e07f7c
1 parent d1b3ad9
commit 4e07f7c
Show file tree

Hide file tree

Showing 6 changed files with 543 additions and 574 deletions.
diff --git a/src/connector.ts b/src/connector.ts
@@ -45,9 +45,11 @@ export class ParallelConnectionStrategy {
     }
 
     for (let i = 0, len = addresses.length; i < len; i++) {
-      const socket = sockets[i] = net.connect(Object.create(this.options, {
-        host: { value: addresses[i].address }
-      }));
+      const socket = sockets[i] = net.connect({
+        ...this.options,
+        host: addresses[i].address,
+        family: addresses[i].family
+      });
 
       socket.on('error', onError);
       socket.on('connect', onConnect);
@@ -70,9 +72,11 @@ export class SequentialConnectionStrategy {
       return callback(new Error('Could not connect (sequence)'));
     }
 
-    const socket = net.connect(Object.create(this.options, {
-      host: { value: next.address }
-    }));
+    const socket = net.connect({
+      ...this.options,
+      host: next.address,
+      family: next.family
+    });
 
     const onError = (_err: Error) => {
       socket.removeListener('error', onError);
@@ -95,48 +99,21 @@ export class SequentialConnectionStrategy {
   }
 }
 
+type LookupFunction = (hostname: string, options: dns.LookupAllOptions, callback: (err: NodeJS.ErrnoException | null, addresses: dns.LookupAddress[]) => void) => void;
+
 export class Connector {
   options: { port: number, host: string, localAddress?: string };
   multiSubnetFailover: boolean;
+  lookup: LookupFunction;
 
-  constructor(options: { port: number, host: string, localAddress?: string }, multiSubnetFailover: boolean) {
+  constructor(options: { port: number, host: string, localAddress?: string, lookup?: LookupFunction }, multiSubnetFailover: boolean) {
     this.options = options;
+    this.lookup = options.lookup ?? dns.lookup;
     this.multiSubnetFailover = multiSubnetFailover;
   }
 
   execute(cb: (err: Error | null, socket?: net.Socket) => void) {
-    if (net.isIP(this.options.host)) {
-      this.executeForIP(cb);
-    } else {
-      this.executeForHostname(cb);
-    }
-  }
-
-  executeForIP(cb: (err: Error | null, socket?: net.Socket) => void) {
-    const socket = net.connect(this.options);
-
-    const onError = (err: Error) => {
-      socket.removeListener('error', onError);
-      socket.removeListener('connect', onConnect);
-
-      socket.destroy();
-
-      cb(err);
-    };
-
-    const onConnect = () => {
-      socket.removeListener('error', onError);
-      socket.removeListener('connect', onConnect);
-
-      cb(null, socket);
-    };
-
-    socket.on('error', onError);
-    socket.on('connect', onConnect);
-  }
-
-  executeForHostname(cb: (err: Error | null, socket?: net.Socket) => void) {
-    dns.lookup(punycode.toASCII(this.options.host), { all: true }, (err, addresses) => {
+    this.lookupAllAddresses(this.options.host, (err, addresses) => {
       if (err) {
         return cb(err);
       }
@@ -148,4 +125,14 @@ export class Connector {
       }
     });
   }
+
+  lookupAllAddresses(host: string, callback: (err: NodeJS.ErrnoException | null, addresses: dns.LookupAddress[]) => void) {
+    if (net.isIPv6(host)) {
+      process.nextTick(callback, null, [{ address: host, family: 6 }]);
+    } else if (net.isIPv4(host)) {
+      process.nextTick(callback, null, [{ address: host, family: 4 }]);
+    } else {
+      this.lookup.call(null, punycode.toASCII(host), { all: true }, callback);
+    }
+  }
 }
diff --git a/src/instance-lookup.ts b/src/instance-lookup.ts
@@ -1,19 +1,22 @@
 import { Sender } from './sender';
+import dns from 'dns';
 
 const SQL_SERVER_BROWSER_PORT = 1434;
 const TIMEOUT = 2 * 1000;
 const RETRIES = 3;
 // There are three bytes at the start of the response, whose purpose is unknown.
 const MYSTERY_HEADER_LENGTH = 3;
 
+type LookupFunction = (hostname: string, options: dns.LookupAllOptions, callback: (err: NodeJS.ErrnoException | null, addresses: dns.LookupAddress[]) => void) => void;
+
 // Most of the functionality has been determined from from jTDS's MSSqlServerInfo class.
 export class InstanceLookup {
   // Wrapper allows for stubbing Sender when unit testing instance-lookup.
-  createSender(host: string, port: number, request: Buffer) {
-    return new Sender(host, port, request);
+  createSender(host: string, port: number, lookup: LookupFunction, request: Buffer) {
+    return new Sender(host, port, lookup, request);
   }
 
-  instanceLookup(options: { server: string, instanceName: string, timeout?: number, retries?: number }, callback: (message: string | undefined, port?: number) => void) {
+  instanceLookup(options: { server: string, instanceName: string, timeout?: number, retries?: number, port?: number, lookup?: LookupFunction }, callback: (message: string | undefined, port?: number) => void) {
     const server = options.server;
     if (typeof server !== 'string') {
       throw new TypeError('Invalid arguments: "server" must be a string');
@@ -34,25 +37,37 @@ export class InstanceLookup {
       throw new TypeError('Invalid arguments: "retries" must be a number');
     }
 
+    if (options.lookup !== undefined && typeof options.lookup !== 'function') {
+      throw new TypeError('Invalid arguments: "lookup" must be a function');
+    }
+    const lookup = options.lookup ?? dns.lookup;
+
+    if (options.port !== undefined && typeof options.port !== 'number') {
+      throw new TypeError('Invalid arguments: "port" must be a number');
+    }
+    const port = options.port ?? SQL_SERVER_BROWSER_PORT;
+
     if (typeof callback !== 'function') {
       throw new TypeError('Invalid arguments: "callback" must be a function');
     }
 
-    let sender: Sender;
-    let timer: NodeJS.Timeout;
     let retriesLeft = retries;
 
-    const onTimeout = () => {
-      sender.cancel();
-      makeAttempt();
-    };
-
     const makeAttempt = () => {
+      let sender: Sender;
+      let timer: NodeJS.Timeout;
+
+      const onTimeout = () => {
+        sender.cancel();
+        makeAttempt();
+      };
+
       if (retriesLeft > 0) {
         retriesLeft--;
 
         const request = Buffer.from([0x02]);
-        sender = this.createSender(options.server, SQL_SERVER_BROWSER_PORT, request);
+        sender = this.createSender(options.server, port, lookup, request);
+        timer = setTimeout(onTimeout, timeout);
         sender.execute((err, response) => {
           clearTimeout(timer);
           if (err) {
@@ -69,8 +84,6 @@ export class InstanceLookup {
             }
           }
         });
-
-        timer = setTimeout(onTimeout, timeout);
       } else {
         callback('Failed to get response from SQL Server Browser on ' + server);
       }

diff --git a/src/sender.ts b/src/sender.ts
@@ -3,6 +3,8 @@ import dns from 'dns';
 import net from 'net';
 import * as punycode from 'punycode';
 
+type LookupFunction = (hostname: string, options: dns.LookupAllOptions, callback: (err: NodeJS.ErrnoException | null, addresses: dns.LookupAddress[]) => void) => void;
+
 export class ParallelSendStrategy {
   addresses: dns.LookupAddress[];
   port: number;
@@ -106,11 +108,13 @@ export class Sender {
   port: number;
   request: Buffer;
   parallelSendStrategy: ParallelSendStrategy | null;
+  lookup: LookupFunction;
 
-  constructor(host: string, port: number, request: Buffer) {
+  constructor(host: string, port: number, lookup: LookupFunction, request: Buffer) {
     this.host = host;
     this.port = port;
     this.request = request;
+    this.lookup = lookup;
 
     this.parallelSendStrategy = null;
   }
@@ -129,7 +133,7 @@ export class Sender {
 
   // Wrapper for stubbing. Sinon does not have support for stubbing module functions.
   invokeLookupAll(host: string, cb: (error: Error | null, addresses?: dns.LookupAddress[]) => void) {
-    dns.lookup(punycode.toASCII(host), { all: true }, cb);
+    this.lookup.call(null, punycode.toASCII(host), { all: true }, cb);
   }
 
   executeForHostname(cb: (error: Error | null, message?: Buffer) => void) {