Add support for encrypted discriminators

Author: baileympearson

lib/helpers/model/discriminator.js

@@ -17,6 +17,60 @@ const CUSTOMIZABLE_DISCRIMINATOR_OPTIONS = {
   methods: true
 };
 
+/**
+ * Validate fields declared on the child schema when either schema is configured for encryption.  Specifically, this function ensures that:
+ *
+ * - any encrypted fields are declared on exactly one of the schemas (not both)
+ * - encrypted fields cannot be declared on either the parent or child schema, where the other schema declares the same field without encryption.
+ *
+ * @param {Schema} parentSchema
+ * @param {Schema} childSchema
+ */
+function validateDiscriminatorSchemasForEncryption(parentSchema, childSchema) {
+  /**
+   * @param schema { Schema }
+   *
+   * Given a schema, yields **all** paths to values inside that schema, recursively iterating over any nested schemas.  This is
+   * intended for use with encryption, so nested arrays are not considered because encryption on values inside arrays is not supported.
+   *
+   * @returns { Iterable<string> }
+   */
+  function* allPaths(schema, prefix) {
+    for (const path of Object.keys(schema.paths)) {
+      const fullPath = prefix != null ? `${prefix}.${path}` : path;
+      if (schema.path(path).instance === 'Embedded') {
+        yield* allPaths(schema.path(path).schema, fullPath);
+      } else {
+        yield fullPath;
+      }
+    }
+  }
+
+  /**
+   * @param {Iterable<T>} i1
+   * @param {Iterable<T>} i2
+   *
+   * @returns {Generator<T>}
+   */
+  function* setIntersection(i1, i2) {
+    const s1 = new Set(i1);
+    for (const item of i2) {
+      if (s1.has(item)) {
+        yield item;
+      }
+    }
+  }
+
+  for (const path of setIntersection(allPaths(parentSchema), allPaths(childSchema))) {
+    if (parentSchema._hasEncryptedField(path) && childSchema._hasEncryptedField(path)) {
+      throw new Error(`encrypted fields cannot be declared on both the base schema and the child schema in a discriminator. path=${path}`);
+    }
+
+    if (parentSchema._hasEncryptedField(path) || childSchema._hasEncryptedField(path)) {
+      throw new Error(`encrypted fields cannot have the same path as a non-encrypted field for discriminators. path=${path}`);
+    }
+  }
+}
 /*!
  * ignore
  */
@@ -80,6 +134,8 @@ module.exports = function discriminator(model, name, schema, tiedValue, applyPlu
     value = tiedValue;
   }
 
+  validateDiscriminatorSchemasForEncryption(model.schema, schema);
+
   function merge(schema, baseSchema) {
     // Retain original schema before merging base schema
     schema._baseSchema = baseSchema;
@@ -95,13 +151,6 @@ module.exports = function discriminator(model, name, schema, tiedValue, applyPlu
     const baseSchemaPaths = Object.keys(baseSchema.paths);
     const conflictingPaths = [];
 
-
-    baseSchema.eachPath((pathname) => {
-      if (schema._hasEncryptedField(pathname)) {
-        throw new Error(`cannot declare an encrypted field on child schema overriding base schema. key=${pathname}`);
-      }
-    });
-
     for (const path of baseSchemaPaths) {
       if (schema.nested[path]) {
         conflictingPaths.push(path);

test/encryption/encryption.test.js

@@ -1024,7 +1024,7 @@ describe('encryption integration tests', () => {
             }, {
               encryptionType: 'csfle'
             })),
-            /cannot declare an encrypted field on child schema overriding base schema\. key=name/
+            /encrypted fields cannot be declared on both the base schema and the child schema in a discriminator\. path=name/
             );
           });
         });
@@ -1085,12 +1085,45 @@ describe('encryption integration tests', () => {
             }, {
               encryptionType: 'queryableEncryption'
             })),
-            /cannot declare an encrypted field on child schema overriding base schema\. key=name/
+            /encrypted fields cannot be declared on both the base schema and the child schema in a discriminator\. path=name/
             );
           });
         });
       });
 
+      describe('Nested Schema overrides nested path', function() {
+        beforeEach(async function() {
+          connection = createConnection();
+        });
+
+        it('nested objects throw an error', async function() {
+          model = connection.model('Schema', new Schema({
+            name: {
+              first: { type: String, encrypt: { keyId: [keyId], algorithm } }
+            }
+          }, { encryptionType: 'csfle' }));
+
+          assert.throws(() => {
+            model.discriminator('Test', new Schema({
+              name: { first: Number } // Different type, no encryption, stored as same field in MDB
+            }));
+          });
+        });
+
+        it('nested schemas throw an error', async function() {
+          model = connection.model('Schema', new Schema({
+            name: {
+              first: { type: String, encrypt: { keyId: [keyId], algorithm } }
+            }
+          }, { encryptionType: 'csfle' }));
+
+          assert.throws(() => {
+            model.discriminator('Test', new Schema({
+              name: new Schema({ first: Number }) // Different type, no encryption, stored as same field in MDB
+            }));
+          });
+        });
+      });
     });
   });
 });