Make it possible to attach CacheableLookup to an Agent

Fixes #4

index.d.ts

@@ -1,5 +1,6 @@
-import * as  Keyv from 'keyv';
+import Keyv = require('keyv');
 import {Resolver, LookupAddress} from 'dns';
+import {Agent} from 'http';
 
 type IPFamily = 4 | 6;
 
@@ -83,7 +84,7 @@ export default class CacheableLookup {
 	/**
 	 * The asynchronous version of `dns.lookup(…)`.
 	 */
-	lookupAsync(hostname: string, options: LookupOptions & {all: true, details: true}): Promise<ReadonlyArray<EntryObject & {ttl: number, expires: number}>>; 
+	lookupAsync(hostname: string, options: LookupOptions & {all: true, details: true}): Promise<ReadonlyArray<EntryObject & {ttl: number, expires: number}>>;
 	lookupAsync(hostname: string, options: LookupOptions & {all: true}): Promise<ReadonlyArray<EntryObject>>;
 	lookupAsync(hostname: string, options: LookupOptions & {details: true}): Promise<EntryObject & {ttl: number, expires: number}>;
 	lookupAsync(hostname: string, options: LookupOptions): Promise<EntryObject>;
@@ -97,4 +98,12 @@ export default class CacheableLookup {
 	 * An asynchronous function which makes a new DNS lookup query and updates the database. This is used by `query(hostname, family)` if no entry in the database is present. Returns an array of objects with `address`, `family`, `ttl` and `expires` properties.
 	 */
 	queryAndCache(hostname: string, family: IPFamily): Promise<ReadonlyArray<EntryObject>>;
+	/**
+	 * Attaches itself to an Agent instance.
+	 */
+	install(agent: Agent): void;
+	/**
+	 * Removes itself from an Agent instance.
+	 */
+	uninstall(agent: Agent): void;
 }

index.js

@@ -4,6 +4,15 @@ const {promisify} = require('util');
 const os = require('os');
 const Keyv = require('keyv');
 
+const kCacheableLookupData = Symbol('cacheableLookupData');
+const kCacheableLookupInstance = Symbol('cacheableLookupInstance');
+
+const verifyAgent = agent => {
+	if (!(agent && typeof agent.createConnection === 'function')) {
+		throw new Error('Expected an Agent instance as the first argument');
+	}
+};
+
 const map4to6 = entries => {
 	for (const entry of entries) {
 		entry.address = `::ffff:${entry.address}`;
@@ -170,6 +179,40 @@ class CacheableLookup {
 	_getEntry(entries) {
 		return entries[Math.floor(Math.random() * entries.length)];
 	}
+
+	install(agent) {
+		verifyAgent(agent);
+
+		if (kCacheableLookupData in agent) {
+			throw new Error('CacheableLookup has been already installed');
+		}
+
+		agent[kCacheableLookupData] = agent.createConnection;
+		agent[kCacheableLookupInstance] = this;
+
+		agent.createConnection = (options, callback) => {
+			if (!('lookup' in options)) {
+				options.lookup = this.lookup;
+			}
+
+			return agent[kCacheableLookupData](options, callback);
+		};
+	}
+
+	uninstall(agent) {
+		verifyAgent(agent);
+
+		if (agent[kCacheableLookupData]) {
+			if (agent[kCacheableLookupInstance] !== this) {
+				throw new Error('The agent is not owned by this CacheableLookup instance');
+			}
+
+			agent.createConnection = agent[kCacheableLookupData];
+
+			delete agent[kCacheableLookupData];
+			delete agent[kCacheableLookupInstance];
+		}
+	}
 }
 
 module.exports = CacheableLookup;

test.js

@@ -5,6 +5,10 @@ const test = require('ava');
 const proxyquire = require('proxyquire');
 const CacheableLookup = require('.');
 
+const makeRequest = options => new Promise((resolve, reject) => {
+	http.get(options, resolve).once('error', reject);
+});
+
 const sleep = ms => new Promise(resolve => setTimeout(resolve, ms));
 
 const mockedInterfaces = (options = {}) => {
@@ -426,17 +430,17 @@ test('entry with 0 ttl', async t => {
 test('http example', async t => {
 	const cacheable = new CacheableLookup({resolver});
 
-	const makeRequest = () => new Promise((resolve, reject) => {
-		http.get({
-			hostname: 'example',
-			port: 8080,
-			lookup: cacheable.lookup
-		}, response => {
-			resolve(response);
-		}).once('error', reject);
-	});
+	const options = {
+		hostname: 'example',
+		port: 8080,
+		lookup: cacheable.lookup
+	};
 
+<<<<<<< HEAD
 	await t.throwsAsync(() => makeRequest(), {message: 'connect ECONNREFUSED 127.0.0.127:8080'});
+=======
+	await t.throwsAsync(() => makeRequest(options), 'connect ECONNREFUSED 127.0.0.127:8080');
+>>>>>>> Make it possible to attach CacheableLookup to an Agent
 });
 
 test('.lookup() and .lookupAsync() are automatically bounded', async t => {
@@ -453,3 +457,79 @@ test('works (Internet connection)', async t => {
 	t.true(typeof address === 'string');
 	t.is(family, 4);
 });
+
+test.serial('install & uninstall', async t => {
+	const cacheable = new CacheableLookup({resolver});
+	cacheable.install(http.globalAgent);
+
+	const options = {
+		hostname: 'example',
+		port: 8080
+	};
+
+	await t.throwsAsync(() => makeRequest(options), 'connect ECONNREFUSED 127.0.0.127:8080');
+
+	cacheable.uninstall(http.globalAgent);
+
+	await t.throwsAsync(() => makeRequest(options), 'getaddrinfo ENOTFOUND example');
+});
+
+test('`.install()` throws if no Agent provided', t => {
+	const cacheable = new CacheableLookup();
+
+	t.throws(() => cacheable.install(), 'Expected an Agent instance as the first argument');
+	t.throws(() => cacheable.install(1), 'Expected an Agent instance as the first argument');
+});
+
+test('`.uninstall()` throws if no Agent provided', t => {
+	const cacheable = new CacheableLookup();
+
+	t.throws(() => cacheable.uninstall(), 'Expected an Agent instance as the first argument');
+	t.throws(() => cacheable.uninstall(1), 'Expected an Agent instance as the first argument');
+});
+
+test.serial('`.uninstall()` does not alter unmodified Agents', t => {
+	const cacheable = new CacheableLookup();
+	const {createConnection} = http.globalAgent;
+
+	cacheable.uninstall(http.globalAgent);
+
+	t.is(createConnection, http.globalAgent.createConnection);
+});
+
+test.serial('throws if double-installing CacheableLookup', t => {
+	const cacheable = new CacheableLookup();
+
+	cacheable.install(http.globalAgent);
+	t.throws(() => cacheable.install(http.globalAgent), 'CacheableLookup has been already installed');
+
+	cacheable.uninstall(http.globalAgent);
+});
+
+test.serial('install - providing custom lookup function anyway', async t => {
+	const a = new CacheableLookup();
+	const b = new CacheableLookup({resolver});
+
+	a.install(http.globalAgent);
+
+	const options = {
+		hostname: 'example',
+		port: 8080,
+		lookup: b.lookup
+	};
+
+	await t.throwsAsync(() => makeRequest(options), 'connect ECONNREFUSED 127.0.0.127:8080');
+
+	a.uninstall(http.globalAgent);
+});
+
+test.serial('throws when calling `.uninstall()` on the wrong instance', t => {
+	const a = new CacheableLookup();
+	const b = new CacheableLookup({resolver});
+
+	a.install(http.globalAgent);
+
+	t.throws(() => b.uninstall(http.globalAgent), 'The agent is not owned by this CacheableLookup instance');
+
+	a.uninstall(http.globalAgent);
+});