fix: allow updating unindexed encrypted fields COMPASS-5809 (#3262)

* chore(data-service): track equality-searchable encrypted fields

In addition to tracking all encrypted fields for a given collection’s
schema sources, also track equality-searchable encrypted fields.
This is helpful for knowing whether a field can be included in
the query part of a Compass CRUD update operation that includes
original values of fields in the query portion of that update.

As part of this, the list of encrypted fields reported by
`knownSchemaForCollection()` and used internally in the encrypted
field tracker is changed from a simple array of strings to a
pair of field path lists, one for all encrypted and one for
equality-searchable encrypted fields in the schema. Using
field paths instead of `.`-joined strings has no direct effect
here, but aligns this part of Compass with the recent changes
to nested field editing/dots and dollars support.

* fix(compass-crud): exclude unindexed encrypted fields in update query COMPASS-5809

When updating unindexed encrypted fields, exclude them from the
find part of the update operation, since they are by definition
not queryable.

This modifies the HadronDocument query building methods to take
options arguments instead of key lists, and updates these options
to be passed as field path lists rather than a plain array of strings
(where `.` was ambiguous between literal field names and nested
fields, which was not a problem previously because the field paths
were always coming from the context of shard keys, where `.` is
unambiguous but now is with FLE1 nested fields).

Author: addaleax

packages/compass-crud/src/stores/crud-store.js

@@ -456,7 +456,12 @@ const configureStore = (options = {}) => {
         // required for updated documents in sharded collections.
         const { query, updateDoc } =
           doc.generateUpdateUnlessChangedInBackgroundQuery(
-            Object.keys(this.state.shardKeys)
+            // '.' in shard keys means nested doc
+            {
+              alwaysIncludeKeys: Object.keys(this.state.shardKeys || {}).map(
+                (key) => key.split('.')
+              ),
+            }
           );
         debug('Performing findOneAndUpdate', { query, updateDoc });
 
@@ -513,38 +518,58 @@ const configureStore = (options = {}) => {
       track('Document Updated', { mode: this.modeForTelemetry() });
       try {
         doc.emit('update-start');
-        const object = doc.generateObject();
-        const query = doc.getQueryForOriginalKeysAndValuesForSpecifiedKeys([
-          '_id',
-          ...Object.keys(this.state.shardKeys || {}),
-        ]);
-        debug('Performing findOneAndReplace', { query, object });
 
         if (!(await this._verifyUpdateAllowed(this.state.ns, doc))) {
           // _verifyUpdateAllowed emitted update-error
           return;
         }
 
+        const object = doc.generateObject();
+        const queryKeyInclusionOptions = {
+          alwaysIncludeKeys: [
+            ['_id'],
+            // '.' in shard keys means nested doc
+            ...Object.keys(this.state.shardKeys || {}).map((key) =>
+              key.split('.')
+            ),
+          ],
+        };
+
         if (
           this.dataService.getCSFLEMode &&
-          this.dataService.getCSFLEMode() === 'enabled' &&
-          object.__safeContent__ &&
-          isEqual(
-            object.__safeContent__,
-            doc.generateOriginalObject().__safeContent__
-          ) &&
-          (
-            await this.dataService
-              .getCSFLECollectionTracker()
-              .knownSchemaForCollection(this.state.ns)
-          ).hasSchema
+          this.dataService.getCSFLEMode() === 'enabled'
         ) {
-          // SERVER-66662 blocks writes of __safeContent__ for queryable-encryption-enabled
-          // collections. We remove it unless it was edited, in which case we assume that the
-          // user really knows what they are doing.
-          delete object.__safeContent__;
+          const knownSchemaForCollection = await this.dataService
+            .getCSFLECollectionTracker()
+            .knownSchemaForCollection(this.state.ns);
+
+          // The find/query portion will typically exclude encrypted fields,
+          // because those cannot be queried to make sure that the original
+          // value matches the current one; however, if we know that the
+          // field is equality-searchable, we can (and should) still include it.
+          queryKeyInclusionOptions.includableEncryptedKeys =
+            knownSchemaForCollection.encryptedFields.equalityQueryableEncryptedFields;
+
+          if (
+            object.__safeContent__ &&
+            isEqual(
+              object.__safeContent__,
+              doc.generateOriginalObject().__safeContent__
+            ) &&
+            knownSchemaForCollection.hasSchema
+          ) {
+            // SERVER-66662 blocks writes of __safeContent__ for queryable-encryption-enabled
+            // collections. We remove it unless it was edited, in which case we assume that the
+            // user really knows what they are doing.
+            delete object.__safeContent__;
+          }
         }
 
+        const query = doc.getQueryForOriginalKeysAndValuesForSpecifiedKeys(
+          queryKeyInclusionOptions
+        );
+        debug('Performing findOneAndReplace', { query, object });
+
         // eslint-disable-next-line no-shadow
         const [error, d] = await findAndModifyWithFLEFallback(
           this.dataService,
@@ -734,10 +759,20 @@ const configureStore = (options = {}) => {
         // does not have a schema.
         const csfleCollectionTracker =
           this.dataService.getCSFLECollectionTracker();
-        const { hasSchema, encryptedFields } =
-          await csfleCollectionTracker.knownSchemaForCollection(this.state.ns);
-        if (encryptedFields) {
-          csfleState.encryptedFields = encryptedFields;
+        const {
+          hasSchema,
+          encryptedFields: { encryptedFields },
+        } = await csfleCollectionTracker.knownSchemaForCollection(
+          this.state.ns
+        );
+        if (encryptedFields.length > 0) {
+          // This is for displaying encrypted fields to the user. We do not really
+          // need to worry about the distinction between '.' as a nested-field
+          // indicator and '.' as a literal part of a field name here, esp. since
+          // automatic Queryable Encryption does not support '.' in field names at all.
+          csfleState.encryptedFields = encryptedFields.map((field) =>
+            field.join('.')
+          );
         }
         if (!hasSchema) {
           csfleState.state = 'no-known-schema';

packages/compass-crud/src/stores/crud-store.spec.js

@@ -1367,7 +1367,7 @@ describe('store', function () {
         });
 
         getCSFLEMode.returns('enabled');
-        knownSchemaForCollection.resolves({ hasSchema: false });
+        knownSchemaForCollection.resolves({ hasSchema: false, encryptedFields: { encryptedFields: [] } });
 
         store.openInsertDocumentDialog(doc, false);
 
@@ -1391,7 +1391,7 @@ describe('store', function () {
         getCSFLEMode.returns('enabled');
         knownSchemaForCollection.resolves({
           hasSchema: true,
-          encryptedFields: ['x'],
+          encryptedFields: { encryptedFields: [['x']] },
         });
         isUpdateAllowed.resolves(false);
 
@@ -1417,7 +1417,7 @@ describe('store', function () {
         getCSFLEMode.returns('enabled');
         knownSchemaForCollection.resolves({
           hasSchema: true,
-          encryptedFields: ['x'],
+          encryptedFields: { encryptedFields: [['x']] },
         });
         isUpdateAllowed.resolves(true);
 
@@ -1442,7 +1442,7 @@ describe('store', function () {
         getCSFLEMode.returns('disabled');
         knownSchemaForCollection.resolves({
           hasSchema: true,
-          encryptedFields: ['x'],
+          encryptedFields: { encryptedFields: [['x']] },
         });
         isUpdateAllowed.resolves(true);
 

packages/compass-e2e-tests/tests/in-use-encryption.test.ts

@@ -209,6 +209,7 @@ describe('FLE2', function () {
     describe('when fleEncryptedFieldsMap is specified while connecting', function () {
       const databaseName = 'fle-test';
       const collectionName = 'my-another-collection';
+      const collectionNameUnindexed = 'my-another-collection2';
       let compass: Compass;
       let browser: CompassBrowser;
       let plainMongo: MongoClient;
@@ -233,6 +234,15 @@ describe('FLE2', function () {
                   queries: { queryType: 'equality' }
                 }
               ]
+            },
+            '${databaseName}.${collectionNameUnindexed}': {
+              fields: [
+                {
+                  path: 'phoneNumber',
+                  keyId: UUID("28bbc608-524e-4717-9246-33633361788e"),
+                  bsonType: 'string'
+                }
+              ]
             }
           }`,
         });
@@ -382,50 +392,66 @@ describe('FLE2', function () {
         expect(isDocumentPhoneNumberDecryptedIconExisting).to.be.equal(true);
       });
 
-      it('can edit and query the encrypted field in the CRUD view', async function () {
-        await browser.shellEval(`db.createCollection('${collectionName}')`);
-        await browser.shellEval(
-          `db[${JSON.stringify(
-            collectionName
-          )}].insertOne({ "phoneNumber": "30303030", "name": "Person X" })`
-        );
+      for (const [mode, coll] of [
+        ['indexed', collectionName],
+        ['unindexed', collectionNameUnindexed],
+      ]) {
+        it(`can edit and query the ${mode} encrypted field in the CRUD view`, async function () {
+          await browser.shellEval(`db.createCollection('${coll}')`);
+          await browser.shellEval(
+            `db[${JSON.stringify(
+              coll
+            )}].insertOne({ "phoneNumber": "30303030", "name": "Person X" })`
+          );
 
-        await browser.navigateToCollectionTab(
-          databaseName,
-          collectionName,
-          'Documents'
-        );
+          await browser.navigateToCollectionTab(
+            databaseName,
+            coll,
+            'Documents'
+          );
 
-        const result = await getFirstListDocument(browser);
-        expect(result.phoneNumber).to.be.equal('"30303030"');
+          const result = await getFirstListDocument(browser);
+          expect(result.phoneNumber).to.be.equal('"30303030"');
 
-        const document = await browser.$(Selectors.DocumentListEntry);
-        const value = await document.$(
-          `${Selectors.HadronDocumentElement}[data-field="phoneNumber"] ${Selectors.HadronDocumentClickableValue}`
-        );
-        await value.doubleClick();
+          const document = await browser.$(Selectors.DocumentListEntry);
+          const value = await document.$(
+            `${Selectors.HadronDocumentElement}[data-field="phoneNumber"] ${Selectors.HadronDocumentClickableValue}`
+          );
+          await value.doubleClick();
 
-        const input = await document.$(
-          `${Selectors.HadronDocumentElement}[data-field="phoneNumber"] ${Selectors.HadronDocumentValueEditor}`
-        );
-        await input.setValue('10101010');
+          const input = await document.$(
+            `${Selectors.HadronDocumentElement}[data-field="phoneNumber"] ${Selectors.HadronDocumentValueEditor}`
+          );
+          await input.setValue('10101010');
 
-        const footer = await document.$(Selectors.DocumentFooterMessage);
-        expect(await footer.getText()).to.equal('Document modified.');
+          const footer = await document.$(Selectors.DocumentFooterMessage);
+          expect(await footer.getText()).to.equal('Document modified.');
 
-        const button = await document.$(Selectors.UpdateDocumentButton);
-        await button.click();
-        await footer.waitForDisplayed({ reverse: true });
+          const button = await document.$(Selectors.UpdateDocumentButton);
+          await button.click();
+          try {
+            await footer.waitForDisplayed({ reverse: true, timeout: 1000 });
+          } catch (err) {
+            if (
+              mode === 'unindexed' &&
+              (await footer.getText()) ===
+                'Found indexed encrypted fields but could not find __safeContent__'
+            ) {
+              return; // This version is affected by SERVER-68065 where updates on unindexed fields in QE are buggy.
+            }
+          }
+          await footer.waitForDisplayed({ reverse: true });
 
-        await browser.runFindOperation(
-          'Documents',
-          "{ phoneNumber: '10101010' }"
-        );
+          await browser.runFindOperation(
+            'Documents',
+            "{ phoneNumber: '10101010' }"
+          );
 
-        const modifiedResult = await getFirstListDocument(browser);
-        expect(modifiedResult.phoneNumber).to.be.equal('"10101010"');
-        expect(modifiedResult._id).to.be.equal(result._id);
-      });
+          const modifiedResult = await getFirstListDocument(browser);
+          expect(modifiedResult.phoneNumber).to.be.equal('"10101010"');
+          expect(modifiedResult._id).to.be.equal(result._id);
+        });
+      }
 
       it('can edit and query the encrypted field in the JSON view', async function () {
         await browser.shellEval(`db.createCollection('${collectionName}')`);