Add support for encrypted models and discriminators

Author: baileympearson

docs/field-level-encryption.md

@@ -151,3 +151,13 @@ To declare a field as encrypted, you must:
 2. Choose an encryption type for the schema and configure the schema for the encryption type
 
 Not all schematypes are supported for CSFLE and QE.  For an overview of valid schema types, refer to MongoDB's documentation.
+
+### Registering Models
+
+Encrypted schemas must be registered on a connection, not the Mongoose global:
+
+```javascript
+
+const connection = mongoose.createConnection();
+const UserModel = connection.model('User', encryptedUserSchema);
+```

lib/drivers/node-mongodb-native/connection.js

@@ -12,6 +12,7 @@ const pkg = require('../../../package.json');
 const processConnectionOptions = require('../../helpers/processConnectionOptions');
 const setTimeout = require('../../helpers/timers').setTimeout;
 const utils = require('../../utils');
+const Schema = require('../../schema');
 
 /**
  * A [node-mongodb-native](https://github.com/mongodb/node-mongodb-native) connection implementation.
@@ -320,6 +321,16 @@ NativeConnection.prototype.createClient = async function createClient(uri, optio
     };
   }
 
+  const { schemaMap, encryptedFieldsMap } = this._buildEncryptionSchemas();
+
+  if (Object.keys(schemaMap).length > 0) {
+    options.autoEncryption.schemaMap = schemaMap;
+  }
+
+  if (Object.keys(encryptedFieldsMap).length > 0) {
+    options.autoEncryption.encryptedFieldsMap = encryptedFieldsMap;
+  }
+
   this.readyState = STATES.connecting;
   this._connectionString = uri;
 
@@ -343,6 +354,55 @@ NativeConnection.prototype.createClient = async function createClient(uri, optio
   return this;
 };
 
+/**
+ * Given a connection, which may or may not have encrypted models, build
+ * a schemaMap and/or an encryptedFieldsMap for the connection, combining all models
+ * into a single schemaMap and encryptedFields map.
+ *
+ * @returns a copy of the options object with a schemaMap and/or an encryptedFieldsMap added to the options' autoEncryption
+ * options.
+  */
+NativeConnection.prototype._buildEncryptionSchemas = function() {
+  const qeMappings = {};
+  const csfleMappings = {};
+
+  // If discriminators are configured for the collection, there might be multiple models
+  // pointing to the same namespace.  For this scenario, we merge all the schemas for each namespace
+  // into a single schema.
+  // Notably, this doesn't allow for discriminators to declare multiple values on the same fields.
+  for (const model of Object.values(this.models)) {
+    const { schema, collection: { collectionName } } = model;
+    const namespace = `${this.$dbName}.${collectionName}`;
+    if (schema.encryptionType() === 'csfle') {
+      csfleMappings[namespace] ??= new Schema({}, { encryptionType: 'csfle' });
+      csfleMappings[namespace].add(schema);
+    } else if (schema.encryptionType() === 'queryableEncryption') {
+      qeMappings[namespace] ??= new Schema({}, { encryptionType: 'queryableEncryption' });
+      qeMappings[namespace].add(schema);
+    }
+  }
+
+  const schemaMap = Object.entries(csfleMappings).reduce(
+    (schemaMap, [namespace, schema]) => {
+      schemaMap[namespace] = schema._buildSchemaMap();
+      return schemaMap;
+    },
+    {}
+  );
+
+  const encryptedFieldsMap = Object.entries(qeMappings).reduce(
+    (encryptedFieldsMap, [namespace, schema]) => {
+      encryptedFieldsMap[namespace] = schema._buildEncryptedFields();
+      return encryptedFieldsMap;
+    },
+    {}
+  );
+
+  return {
+    schemaMap, encryptedFieldsMap
+  };
+};
+
 /*!
  * ignore
  */

lib/schema.js

@@ -914,6 +914,63 @@ Schema.prototype._hasEncryptedFields = function _hasEncryptedFields() {
   return Object.keys(this.encryptedFields).length > 0;
 };
 
+Schema.prototype._buildEncryptedFields = function() {
+  const fields = Object.entries(this.encryptedFields).map(
+    ([path, config]) => {
+      const bsonType = this.path(path).autoEncryptionType();
+      // { path, bsonType, keyId, queries? }
+      return { path, bsonType, ...config };
+    });
+
+  return { fields };
+};
+
+Schema.prototype._buildSchemaMap = function() {
+  /**
+   * `schemaMap`s are JSON schemas, which use the following structure to represent objects:
+   *    { field: { bsonType: 'object', properties: { ... } } }
+   *
+   * for example, a schema that looks like this `{ a: { b: int32 } }` would be encoded as
+   * `{ a: { bsonType: 'object', properties: { b: < encryption configuration > } } }`
+   *
+   * This function takes an array of path segments, an output object (that gets mutated) and
+   * a value to associated with the full path, and constructs a valid CSFLE JSON schema path for
+   * the object.  This works for deeply nested properties as well.
+   *
+   * @param {string[]} path array of path components
+   * @param {object} object the object in which to build a JSON schema of `path`'s properties
+   * @param {object} value the value to associate with the path in object
+   */
+  function buildNestedPath(path, object, value) {
+    let i = 0, component = path[i];
+    for (; i < path.length - 1; ++i, component = path[i]) {
+      object[component] = object[component] == null ? {
+        bsonType: 'object',
+        properties: {}
+      } : object[component];
+      object = object[component].properties;
+    }
+    object[component] = value;
+  }
+
+  const schemaMapPropertyReducer = (accum, [path, propertyConfig]) => {
+    const bsonType = this.path(path).autoEncryptionType();
+    const pathComponents = path.split('.');
+    const configuration = { encrypt: { ...propertyConfig, bsonType } };
+    buildNestedPath(pathComponents, accum, configuration);
+    return accum;
+  };
+
+  const properties = Object.entries(this.encryptedFields).reduce(
+    schemaMapPropertyReducer,
+    {});
+
+  return {
+    bsonType: 'object',
+    properties
+  };
+};
+
 /**
  * Add an alias for `path`. This means getting or setting the `alias`
  * is equivalent to getting or setting the `path`.

lib/schema/bigint.js

@@ -255,7 +255,7 @@ SchemaBigInt.prototype.toJSONSchema = function toJSONSchema(options) {
 };
 
 SchemaBigInt.prototype.autoEncryptionType = function autoEncryptionType() {
-  return 'int64';
+  return 'long';
 };
 
 /*!

lib/schema/boolean.js

@@ -305,7 +305,7 @@ SchemaBoolean.prototype.toJSONSchema = function toJSONSchema(options) {
 };
 
 SchemaBoolean.prototype.autoEncryptionType = function autoEncryptionType() {
-  return 'boolean';
+  return 'bool';
 };
 
 /*!

lib/schema/buffer.js

@@ -315,7 +315,7 @@ SchemaBuffer.prototype.toJSONSchema = function toJSONSchema(options) {
 };
 
 SchemaBuffer.prototype.autoEncryptionType = function autoEncryptionType() {
-  return 'binary';
+  return 'binData';
 };
 
 /*!

lib/schema/decimal128.js

@@ -236,7 +236,7 @@ SchemaDecimal128.prototype.toJSONSchema = function toJSONSchema(options) {
 };
 
 SchemaDecimal128.prototype.autoEncryptionType = function autoEncryptionType() {
-  return 'decimal128';
+  return 'decimal';
 };
 
 /*!

lib/schema/int32.js

@@ -261,7 +261,7 @@ SchemaInt32.prototype.toJSONSchema = function toJSONSchema(options) {
 };
 
 SchemaInt32.prototype.autoEncryptionType = function autoEncryptionType() {
-  return 'int32';
+  return 'int';
 };
 
 

lib/schema/objectId.js

@@ -305,7 +305,7 @@ SchemaObjectId.prototype.toJSONSchema = function toJSONSchema(options) {
 };
 
 SchemaObjectId.prototype.autoEncryptionType = function autoEncryptionType() {
-  return 'objectid';
+  return 'objectId';
 };
 
 /*!

lib/utils.js

@@ -1029,3 +1029,13 @@ exports.injectTimestampsOption = function injectTimestampsOption(writeOperation,
   }
   writeOperation.timestamps = timestampsOption;
 };
+
+exports.print = function(...args) {
+  const { inspect } = require('util');
+  console.error(
+    inspect(
+      ...args,
+      { depth: Infinity }
+    )
+  );
+};