feat: add support for loadBalanced TXT option MONGOSH-819 (#1)

Author: rose-m

README.md

@@ -1,6 +1,7 @@
 # resolve-mongodb-srv
 
-Resolve mongodb+srv:// URLs to mongodb:// URLs
+Resolve `mongodb+srv://` URLs to `mongodb://` URLs as specified in the
+[Initial DNS Seedlist Discovery Specification](https://github.com/mongodb/specifications/blob/master/source/initial-dns-seedlist-discovery/initial-dns-seedlist-discovery.rst).
 
 ```js
 import resolveMongodbSrv from 'resolve-mongodb-srv';

src/index.ts

@@ -1,21 +1,78 @@
+import { promisify } from 'util';
 import { URL, URLSearchParams } from 'whatwg-url';
 
 class MongoParseError extends Error {}
 
+type Address = { name: string; port: number };
 type Options = {
   dns?: {
-    resolveSrv(hostname: string, cb: (err: Error | undefined | null, addresses: { name: string, port: number }[] | undefined) => void): void;
+    resolveSrv(hostname: string, cb: (err: Error | undefined | null, addresses: Address[] | undefined) => void): void;
     resolveTxt(hostname: string, cb: (err: Error | undefined | null, addresses: string[][] | undefined) => void): void;
   };
 };
 
+const ALLOWED_TXT_OPTIONS: Readonly<string[]> = ['authSource', 'replicaSet', 'loadBalanced'];
+
 function matchesParentDomain (srvAddress: string, parentDomain: string): boolean {
   const regex = /^.*?\./;
   const srv = `.${srvAddress.replace(regex, '')}`;
   const parent = `.${parentDomain.replace(regex, '')}`;
   return srv.endsWith(parent);
 }
 
+async function resolveDnsSrvRecord (dns: NonNullable<Options['dns']>, lookupAddress: string): Promise<string[]> {
+  const addresses = await promisify(dns.resolveSrv)(`_mongodb._tcp.${lookupAddress}`);
+  if (!addresses?.length) {
+    throw new MongoParseError('No addresses found at host');
+  }
+
+  for (const { name } of addresses) {
+    if (!matchesParentDomain(name, lookupAddress)) {
+      throw new MongoParseError('Server record does not share hostname with parent URI');
+    }
+  }
+
+  return addresses.map(r => r.name + ((r.port ?? 27017) === 27017 ? '' : `:${r.port}`));
+}
+
+async function resolveDnsTxtRecord (dns: NonNullable<Options['dns']>, lookupAddress: string): Promise<URLSearchParams> {
+  let records: string[][] | undefined;
+  try {
+    records = await promisify(dns.resolveTxt)(lookupAddress);
+  } catch (err) {
+    if (err.code && (err.code !== 'ENODATA' && err.code !== 'ENOTFOUND')) {
+      throw err;
+    }
+  }
+
+  let txtRecord: string;
+  if (records && records.length > 1) {
+    throw new MongoParseError('Multiple text records not allowed');
+  } else {
+    txtRecord = records?.[0]?.join('') ?? '';
+  }
+
+  const txtRecordOptions = new URLSearchParams(txtRecord);
+  const txtRecordOptionKeys = [...txtRecordOptions.keys()];
+  if (txtRecordOptionKeys.some(key => !ALLOWED_TXT_OPTIONS.includes(key))) {
+    throw new MongoParseError(`Text record must only set ${ALLOWED_TXT_OPTIONS.join(', ')}`);
+  }
+
+  const source = txtRecordOptions.get('authSource') ?? undefined;
+  const replicaSet = txtRecordOptions.get('replicaSet') ?? undefined;
+  const loadBalanced = txtRecordOptions.get('loadBalanced') ?? undefined;
+
+  if (source === '' || replicaSet === '' || loadBalanced === '') {
+    throw new MongoParseError('Cannot have empty URI params in DNS TXT Record');
+  }
+
+  if (loadBalanced !== undefined && loadBalanced !== 'true' && loadBalanced !== 'false') {
+    throw new MongoParseError(`DNS TXT Record contains invalid value ${loadBalanced} for loadBalanced option (allowed: true, false)`);
+  }
+
+  return txtRecordOptions;
+}
+
 async function resolveMongodbSrv (input: string, options?: Options): Promise<string> {
   const dns = options?.dns ?? require('dns');
 
@@ -33,60 +90,8 @@ async function resolveMongodbSrv (input: string, options?: Options): Promise<str
 
   const lookupAddress = url.hostname;
   const [srvResult, txtResult] = await Promise.all([
-    (async () => {
-      const addresses = await new Promise<{name: string, port: number}[]>((resolve, reject) => {
-        dns.resolveSrv(`_mongodb._tcp.${lookupAddress}`,
-          (err: Error | null, addresses: { name: string, port: number }[]) => {
-            if (err) return reject(err);
-            return resolve(addresses);
-          });
-      });
-      if (addresses.length === 0) {
-        throw new MongoParseError('No addresses found at host');
-      }
-
-      for (const { name } of addresses) {
-        if (!matchesParentDomain(name, lookupAddress)) {
-          throw new MongoParseError('Server record does not share hostname with parent URI');
-        }
-      }
-
-      return addresses.map(r => r.name + ((r.port ?? 27017) === 27017 ? '' : `:${r.port}`));
-    })(),
-    (async () => {
-      const txtRecord = await new Promise<string>((resolve, reject) => {
-        dns.resolveTxt(lookupAddress, (err: (Error & { code: string }) | null, record: string[][]) => {
-          if (err) {
-            if (err.code && (err.code !== 'ENODATA' && err.code !== 'ENOTFOUND')) {
-              reject(err);
-            } else {
-              resolve('');
-            }
-          } else {
-            if (record.length > 1) {
-              reject(new MongoParseError('Multiple text records not allowed'));
-            } else {
-              resolve(record[0]?.join('') ?? '');
-            }
-          }
-        });
-      });
-
-      const txtRecordOptions = new URLSearchParams(txtRecord);
-      const txtRecordOptionKeys = [...txtRecordOptions.keys()];
-      if (txtRecordOptionKeys.some(key => key !== 'authSource' && key !== 'replicaSet')) {
-        throw new MongoParseError('Text record must only set `authSource` or `replicaSet`');
-      }
-
-      const source = txtRecordOptions.get('authSource') ?? undefined;
-      const replicaSet = txtRecordOptions.get('replicaSet') ?? undefined;
-
-      if (source === '' || replicaSet === '') {
-        throw new MongoParseError('Cannot have empty URI params in DNS TXT Record');
-      }
-
-      return txtRecordOptions;
-    })()
+    resolveDnsSrvRecord(dns, lookupAddress),
+    resolveDnsTxtRecord(dns, lookupAddress)
   ]);
 
   url.protocol = 'mongodb:';

test/index.ts

@@ -162,6 +162,34 @@ describe('resolveMongodbSrv', () => {
         await resolveMongodbSrv('mongodb+srv://server.example.com/?tls=false', { dns }),
         'mongodb://asdf.example.com/?tls=false');
     });
+
+    it('accepts TXT lookup loadBalanced', async () => {
+      srvResult = [{ name: 'asdf.example.com', port: 27017 }];
+      txtResult = [['loadBalanced=true']];
+      assert.strictEqual(
+        await resolveMongodbSrv('mongodb+srv://server.example.com', { dns }),
+        'mongodb://asdf.example.com/?loadBalanced=true&tls=true');
+    });
+
+    it('rejects empty TXT lookup loadBalanced', async () => {
+      srvResult = [{ name: 'asdf.example.com', port: 27017 }];
+      txtResult = [['loadBalanced=']];
+      assert.rejects(resolveMongodbSrv('mongodb+srv://server.example.com', { dns }));
+    });
+
+    it('rejects non true/false TXT lookup loadBalanced', async () => {
+      srvResult = [{ name: 'asdf.example.com', port: 27017 }];
+      txtResult = [['loadBalanced=bla']];
+      assert.rejects(resolveMongodbSrv('mongodb+srv://server.example.com', { dns }));
+    });
+
+    it('prioritizes URL-provided over TXT lookup loadBalanced', async () => {
+      srvResult = [{ name: 'asdf.example.com', port: 27017 }];
+      txtResult = [['loadBalanced=false']];
+      assert.strictEqual(
+        await resolveMongodbSrv('mongodb+srv://server.example.com/?loadBalanced=true', { dns }),
+        'mongodb://asdf.example.com/?loadBalanced=true&tls=true');
+    });
   });
 
   for (const [name, dnsProvider] of [