feat: add support for srvMaxHosts/srvServiceName COMPASS-5324 (#2)

Author: addaleax

src/index.ts

@@ -20,8 +20,8 @@ function matchesParentDomain (srvAddress: string, parentDomain: string): boolean
   return srv.endsWith(parent);
 }
 
-async function resolveDnsSrvRecord (dns: NonNullable<Options['dns']>, lookupAddress: string): Promise<string[]> {
-  const addresses = await promisify(dns.resolveSrv)(`_mongodb._tcp.${lookupAddress}`);
+async function resolveDnsSrvRecord (dns: NonNullable<Options['dns']>, lookupAddress: string, srvServiceName: string): Promise<string[]> {
+  const addresses = await promisify(dns.resolveSrv)(`_${srvServiceName}._tcp.${lookupAddress}`);
   if (!addresses?.length) {
     throw new MongoParseError('No addresses found at host');
   }
@@ -39,8 +39,8 @@ async function resolveDnsTxtRecord (dns: NonNullable<Options['dns']>, lookupAddr
   let records: string[][] | undefined;
   try {
     records = await promisify(dns.resolveTxt)(lookupAddress);
-  } catch (err) {
-    if (err.code && (err.code !== 'ENODATA' && err.code !== 'ENOTFOUND')) {
+  } catch (err: any) {
+    if (err?.code && (err.code !== 'ENODATA' && err.code !== 'ENOTFOUND')) {
       throw err;
     }
   }
@@ -89,11 +89,19 @@ async function resolveMongodbSrv (input: string, options?: Options): Promise<str
   }
 
   const lookupAddress = url.hostname;
+  const srvServiceName = url.searchParams.get('srvServiceName') || 'mongodb';
+  const srvMaxHosts = +(url.searchParams.get('srvMaxHosts') || '0');
+
   const [srvResult, txtResult] = await Promise.all([
-    resolveDnsSrvRecord(dns, lookupAddress),
+    resolveDnsSrvRecord(dns, lookupAddress, srvServiceName),
     resolveDnsTxtRecord(dns, lookupAddress)
   ]);
 
+  if (srvMaxHosts && srvMaxHosts < srvResult.length) {
+    // Replace srvResult with shuffled + limited srvResult
+    srvResult.splice(0, srvResult.length, ...shuffle(srvResult, srvMaxHosts));
+  }
+
   url.protocol = 'mongodb:';
   url.hostname = '__DUMMY_HOSTNAME__';
   if (!url.pathname) {
@@ -107,8 +115,39 @@ async function resolveMongodbSrv (input: string, options?: Options): Promise<str
   if (!url.searchParams.has('tls') && !url.searchParams.has('ssl')) {
     url.searchParams.set('tls', 'true');
   }
+  url.searchParams.delete('srvServiceName');
+  url.searchParams.delete('srvMaxHosts');
 
   return url.toString().replace('__DUMMY_HOSTNAME__', srvResult.join(','));
 }
 
+/**
+ * Fisher–Yates Shuffle
+ * (shamelessly copied from https://github.com/mongodb/node-mongodb-native/blob/1f8b539cd3d60dd9f36baa22fd287241b5c65380/src/utils.ts#L1423-L1451)
+ *
+ * Reference: https://bost.ocks.org/mike/shuffle/
+ * @param sequence - items to be shuffled
+ * @param limit - Defaults to `0`. If nonzero shuffle will slice the randomized array e.g, `.slice(0, limit)` otherwise will return the entire randomized array.
+ */
+function shuffle<T> (sequence: Iterable<T>, limit = 0): Array<T> {
+  const items = Array.from(sequence); // shallow copy in order to never shuffle the input
+
+  limit = Math.min(limit, items.length);
+
+  let remainingItemsToShuffle = items.length;
+  const lowerBound = limit % items.length === 0 ? 1 : items.length - limit;
+  while (remainingItemsToShuffle > lowerBound) {
+    // Pick a remaining element
+    const randomIndex = Math.floor(Math.random() * remainingItemsToShuffle);
+    remainingItemsToShuffle -= 1;
+
+    // And swap it with the current element
+    const swapHold = items[remainingItemsToShuffle];
+    items[remainingItemsToShuffle] = items[randomIndex];
+    items[randomIndex] = swapHold;
+  }
+
+  return limit % items.length === 0 ? items : items.slice(lowerBound);
+}
+
 export = resolveMongodbSrv;

test/index.ts

@@ -190,6 +190,41 @@ describe('resolveMongodbSrv', () => {
         await resolveMongodbSrv('mongodb+srv://server.example.com/?loadBalanced=true', { dns }),
         'mongodb://asdf.example.com/?loadBalanced=true&tls=true');
     });
+
+    it('allows specifying a custom SRV service name', async () => {
+      srvResult = [{ name: 'asdf.example.com', port: 27017 }];
+      txtResult = [['loadBalanced=false']];
+      assert.strictEqual(
+        await resolveMongodbSrv('mongodb+srv://server.example.com/?loadBalanced=true&srvServiceName=custom', { dns }),
+        'mongodb://asdf.example.com/?loadBalanced=true&tls=true');
+      assert.deepStrictEqual(srvQueries, ['_custom._tcp.server.example.com']);
+    });
+
+    it('defaults to _mongodb._tcp as a SRV service name', 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');
+      assert.deepStrictEqual(srvQueries, ['_mongodb._tcp.server.example.com']);
+    });
+
+    it('allows limiting the SRV result to a specific number of hosts', async () => {
+      srvResult = ['host1', 'host2', 'host3'].map(name => ({ name: `${name}.example.com`, port: 27017 }));
+      txtResult = [];
+      assert.strictEqual(
+        await resolveMongodbSrv('mongodb+srv://server.example.com/?srvMaxHosts=0', { dns }),
+        'mongodb://host1.example.com,host2.example.com,host3.example.com/?tls=true');
+      assert.strictEqual(
+        await resolveMongodbSrv('mongodb+srv://server.example.com/?srvMaxHosts=3', { dns }),
+        'mongodb://host1.example.com,host2.example.com,host3.example.com/?tls=true');
+      assert.strictEqual(
+        await resolveMongodbSrv('mongodb+srv://server.example.com/?srvMaxHosts=6', { dns }),
+        'mongodb://host1.example.com,host2.example.com,host3.example.com/?tls=true');
+      assert.match(
+        await resolveMongodbSrv('mongodb+srv://server.example.com/?srvMaxHosts=1', { dns }),
+        /^mongodb:\/\/host[1-3]\.example\.com\/\?tls=true$/);
+    });
   });
 
   for (const [name, dnsProvider] of [