Add findOneAndUpdate and without $set support

lib/mongoose-field-encryption.js

@@ -124,29 +124,38 @@ const fieldEncryption = function(schema, options) {
     }
   });
 
-  schema.pre("update", function(_next) {
+  function updateHook(_next) {
     const next = getCompatitibleNextFunc(_next);
     for (let field of fieldsToEncrypt) {
-      let encryptedFieldName = encryptedFieldNamePrefix + field;
-      let encryptedFieldValue = this._update.$set[encryptedFieldName];
-      let plainTextValue = this._update.$set[field];
+      const encryptedFieldName = encryptedFieldNamePrefix + field;
+      this._update.$set = this._update.$set || {};
+      const plainTextValue = this._update.$set[field] || this._update[field];
+      const encryptedFieldValue = this._update.$set[encryptedFieldName] || this._update[encryptedFieldName];
 
-      if (encryptedFieldValue === false && plainTextValue) {
+      if (!encryptedFieldValue && plainTextValue) {
+        let updateObj = {};
         if (typeof plainTextValue === "string" || plainTextValue instanceof String) {
-          let updateObj = { $set: {} };
-          updateObj.$set[field] = encrypt(plainTextValue, secret);
-          updateObj.$set[encryptedFieldName] = true;
-          this.update({}, updateObj);
+          const encryptedData = encrypt(plainTextValue, secret);
+
+          updateObj[field] = encryptedData;
+          updateObj[encryptedFieldName] = true;
         } else {
-          return next(
-            new Error("Cannot apply mongoose-field-encryption plugin on update to encrypt non string fields")
-          );
+          const encryptedFieldData = encryptedFieldName + encryptedFieldDataSuffix;
+
+          updateObj[field] = undefined;
+          updateObj[encryptedFieldData] = encrypt(JSON.stringify(plainTextValue), secret);
+          updateObj[encryptedFieldName] = true;
         }
+        this.update({}, Object.keys(this._update.$set) > 0 ? { $set: updateObj } : updateObj);
       }
     }
 
     next();
-  });
+  }
+
+  schema.pre("findOneAndUpdate", updateHook);
+
+  schema.pre("update", updateHook);
 
   schema.methods.stripEncryptionFieldMarkers = function() {
     for (let field of fieldsToEncrypt) {

test/test-db.js

@@ -174,7 +174,57 @@ describe("mongoose-field-encryption plugin db", function() {
       });
   });
 
-  it("should not encrypt non string fields on update", () => {
+  it("should encrypt string fields on update without $set", () => {
+    // given
+    let sut = getSut();
+
+    // when
+    return sut
+      .save()
+      .then(() => {
+        expectEncryptionValues(sut);
+
+        return NestedFieldEncryption.update(
+          { _id: sut._id },
+          { toEncryptString: "snoop" }
+        );
+      })
+      .then(() => {
+        return NestedFieldEncryption.findById(sut._id);
+      })
+      .then(found => {
+        // then
+        expect(found.__enc_toEncryptString).to.be.false;
+        expect(found.toEncryptString).to.equal("snoop");
+      });
+  });
+
+  it("should encrypt string fields on fineOneAndUpdate", () => {
+    // given
+    let sut = getSut();
+
+    // when
+    return sut
+      .save()
+      .then(() => {
+        expectEncryptionValues(sut);
+
+        return NestedFieldEncryption.findOneAndUpdate(
+          { _id: sut._id },
+          { $set: { toEncryptString: "snoop", __enc_toEncryptString: false } }
+        );
+      })
+      .then(() => {
+        return NestedFieldEncryption.findById(sut._id);
+      })
+      .then(found => {
+        // then
+        expect(found.__enc_toEncryptString).to.be.false;
+        expect(found.toEncryptString).to.equal("snoop");
+      });
+  })
+
+  it("should encrypt non string fields on update", () => {
     // given
     let sut = getSut();
 
@@ -197,12 +247,67 @@ describe("mongoose-field-encryption plugin db", function() {
         );
       })
       .then(() => {
-        expect.fail("should not have updated");
+        return NestedFieldEncryption.findById(sut._id);
       })
-      .catch(err => {
-        // then
-        // TODO: this is a mongoose cast error
-        expect(err).to.not.be.null;
+      .then(found => {
+        expect(found.toEncryptObject.nested).to.eql("snoop");
       });
   });
+
+  it("should encrypt non string fields on fineOneAndUpdate without $set", () => {
+    // given
+    let sut = getSut();
+
+    // when
+    return sut
+      .save()
+      .then(() => {
+        expectEncryptionValues(sut);
+
+        return NestedFieldEncryption.findOneAndUpdate(
+          {
+            _id: sut._id
+          },
+          {
+            toEncryptObject: { nested: "snoop" }
+          }
+        );
+      })
+      .then(() => {
+        return NestedFieldEncryption.findById(sut._id);
+      })
+      .then(found => {
+        expect(found.toEncryptObject.nested).to.eql("snoop");
+      });
+  });
+
+  it ("should not encrypt already encrypted fields", () => {
+    // given
+    let sut = getSut();
+
+    // when
+    return sut
+      .save()
+      .then(() => {
+        expectEncryptionValues(sut);
+
+        return NestedFieldEncryption.update(
+          {
+            _id: sut._id
+          },
+          {
+            $set: {
+              toEncryptString: "already encrypted string",
+              __enc_toEncryptObject: true
+            }
+          }
+        );
+      })
+      .then(() => {
+        return NestedFieldEncryption.findById(sut._id);
+      })
+      .then(found => {
+        expect(found.toEncryptString).to.eql("already encrypted string");
+      });
+  })
 });