fix: preserve key order in encryption scenarios

Author: nbbeeken

src/cmap/connection.ts

@@ -594,11 +594,27 @@ export class CryptoConnection extends Connection {
       return;
     }
 
+    const sort: Map<string, number> | null = cmd.find || cmd.findAndModify ? cmd.sort : null;
+    const indexKeys: Map<string, number>[] | null = cmd.createIndexes
+      ? cmd.indexes.map((index: { key: Map<string, number> }) => index.key)
+      : null;
+
     autoEncrypter.encrypt(ns.toString(), cmd, options, (err, encrypted) => {
       if (err || encrypted == null) {
         callback(err, null);
         return;
       }
+
+      // Repair JS Map loss
+      if ((cmd.find || cmd.findAndModify) && sort != null) {
+        encrypted.sort = sort;
+      }
+      if (cmd.createIndexes && indexKeys != null) {
+        for (const [offset, index] of indexKeys.entries()) {
+          encrypted.indexes[offset].key = index;
+        }
+      }
+
       super.command(ns, encrypted, options, (err, response) => {
         if (err || response == null) {
           callback(err, response);

test/integration/client-side-encryption/client_side_encryption.prose.test.js

@@ -729,21 +729,6 @@ describe('Client Side Encryption Prose Tests', metadata, function () {
       });
     });
 
-    it('should maintain an ordered sort', async function () {
-      const coll = this.encryptedColl;
-      const events = [];
-      this.clientEncrypted.on('commandStarted', ev => events.push(ev));
-      const sort = new Map([
-        ['1', 1],
-        ['0', 1]
-      ]);
-      await coll.findOne({}, { sort });
-      expect(events).to.have.lengthOf(2);
-      expect(events[0]).to.have.property('commandName', 'listCollections');
-      expect(events[1]).to.have.property('commandName', 'find');
-      expect(events[1].command.sort).to.deep.equal(sort);
-    });
-
     function pruneEvents(events) {
       return events.map(event => {
         // We are pruning out the bunch of repeating As, mostly

test/integration/client-side-encryption/driver.test.js

@@ -222,4 +222,72 @@ describe('Client Side Encryption Functional', function () {
       });
     });
   });
+
+  describe('key order aware command properties', () => {
+    let client;
+    let collection;
+
+    beforeEach(async function () {
+      const encryptionOptions = {
+        monitorCommands: true,
+        autoEncryption: {
+          keyVaultNamespace,
+          kmsProviders: { local: { key: 'A'.repeat(128) } }
+        }
+      };
+      client = this.configuration.newClient({}, encryptionOptions);
+      collection = client.db(dataDbName).collection('keyOrder');
+    });
+
+    afterEach(async () => {
+      if (client) await client.close();
+    });
+
+    describe('find', () => {
+      it('should maintain ordered sort', async function () {
+        const events = [];
+        client.on('commandStarted', ev => events.push(ev));
+        const sort = new Map([
+          ['1', 1],
+          ['0', 1]
+        ]);
+        await collection.findOne({}, { sort });
+        const findEvent = events.find(event => !!event.command.find);
+        expect(findEvent).to.have.property('commandName', 'find');
+        expect(findEvent.command.sort).to.deep.equal(sort);
+      });
+    });
+
+    describe('findAndModify', () => {
+      it('should maintain ordered sort', async function () {
+        const events = [];
+        client.on('commandStarted', ev => events.push(ev));
+        const sort = new Map([
+          ['1', 1],
+          ['0', 1]
+        ]);
+        await collection.findOneAndUpdate({}, { $setOnInsert: { a: 1 } }, { sort });
+        const findAndModifyEvent = events.find(event => !!event.command.findAndModify);
+        expect(findAndModifyEvent).to.have.property('commandName', 'findAndModify');
+        expect(findAndModifyEvent.command.sort).to.deep.equal(sort);
+      });
+    });
+
+    describe('createIndexes', () => {
+      it('should maintain ordered index keys', async function () {
+        const events = [];
+        client.on('commandStarted', ev => events.push(ev));
+        const indexDescription = new Map([
+          ['1', 1],
+          ['0', 1]
+        ]);
+        await collection.createIndex(indexDescription, { name: 'myIndex' });
+        const createIndexEvent = events.find(event => !!event.command.createIndexes);
+        expect(createIndexEvent).to.have.property('commandName', 'createIndexes');
+        expect(createIndexEvent.command.indexes).to.have.lengthOf(1);
+        const index = createIndexEvent.command.indexes[0];
+        expect(index.key).to.deep.equal(indexDescription);
+      });
+    });
+  });
 });