SERVER-64357 Rewrite filter over encrypted fields for findAndModify command

7 files changed, 45 insertions, 31 deletions

diff --git a/buildscripts/resmokeconfig/suites/fle2_sharding.yml b/buildscripts/resmokeconfig/suites/fle2_sharding.yml
index 4272fa547a7..1ea4ed27c78 100644
--- a/buildscripts/resmokeconfig/suites/fle2_sharding.yml
+++ b/buildscripts/resmokeconfig/suites/fle2_sharding.yml
@@ -27,6 +27,6 @@ executor:
     mongod_options:
       set_parameters:
         enableTestCommands: 1
-    num_rs_nodes_per_shard: 1
+    num_rs_nodes_per_shard: 2
     enable_sharding:
     - test
diff --git a/jstests/fle2/libs/encrypted_client_util.js b/jstests/fle2/libs/encrypted_client_util.js
index 196d760a1d1..4970b67e836 100644
--- a/jstests/fle2/libs/encrypted_client_util.js
+++ b/jstests/fle2/libs/encrypted_client_util.js
@@ -149,7 +149,7 @@ class EncryptedClient {
         let encryptedDoc = encryptedDocs[0];
         let unEncryptedDoc = unEncryptedDocs[0];
 
-        assert(encryptedDoc["__safeContent__"] !== undefined);
+        assert(encryptedDoc[kSafeContentField] !== undefined);
 
         for (let field in fields) {
             assert(encryptedDoc.hasOwnProperty(field),
@@ -218,25 +218,16 @@ class EncryptedClient {
     }
 
     /**
-     * Take a snapshot of a collection sorted by _id, run a operation, take a second snapshot.
-     *
-     * Ensure that the documents listed by index in unchangedDocumentIndexArray remain unchanged.
-     * Ensure that the documents listed by index in changedDocumentIndexArray are changed.
+     * Verify that the collection 'collName' contains exactly the documents 'docs'.
      *
      * @param {string} collName
-     * @param {Array} unchangedDocumentIndexArray
-     * @param {Array} changedDocumentIndexArray
-     * @param {Function} func
+     * @param {Array} docs
      * @returns
      */
     assertEncryptedCollectionDocuments(collName, docs) {
         let coll = this._edb.getCollection(collName);
 
-        let onDiskDocs = coll.find({}).sort({_id: 1}).toArray();
-
-        for (let doc of onDiskDocs) {
-            delete doc.__safeContent__;
-        }
+        let onDiskDocs = coll.find({}, {[kSafeContentField]: 0}).sort({_id: 1}).toArray();
 
         assert.docEq(onDiskDocs, docs);
     }
diff --git a/src/mongo/db/fle_crud.cpp b/src/mongo/db/fle_crud.cpp
index 4ae8c2fdbec..296e9c76cf9 100644
--- a/src/mongo/db/fle_crud.cpp
+++ b/src/mongo/db/fle_crud.cpp
@@ -44,6 +44,7 @@
 #include "mongo/db/query/collation/collator_factory_interface.h"
 #include "mongo/db/query/find_command_gen.h"
 #include "mongo/db/query/fle/server_rewrite.h"
+#include "mongo/db/transaction_api.h"
 #include "mongo/idl/idl_parser.h"
 #include "mongo/logv2/log.h"
 #include "mongo/s/grid.h"
@@ -583,10 +584,12 @@ StatusWith<write_ops::FindAndModifyCommandReply> processFindAndModifyRequest(
 
     std::shared_ptr<txn_api::TransactionWithRetries> trun = getTxns(opCtx);
 
+    auto expCtx = makeExpCtx(opCtx, findAndModifyRequest, findAndModifyRequest);
+
     // The function that handles the transaction may outlive this function so we need to use
     // shared_ptrs
     write_ops::FindAndModifyCommandReply reply;
-    auto findAndModifyBlock = std::tie(findAndModifyRequest, reply);
+    auto findAndModifyBlock = std::tie(findAndModifyRequest, reply, expCtx);
     auto sharedFindAndModifyBlock =
         std::make_shared<decltype(findAndModifyBlock)>(findAndModifyBlock);
 
@@ -596,10 +599,9 @@ StatusWith<write_ops::FindAndModifyCommandReply> processFindAndModifyRequest(
                                    ExecutorPtr txnExec) {
             FLEQueryInterfaceImpl queryImpl(txnClient);
 
-            auto [findAndModifyRequest2, reply2] = *sharedFindAndModifyBlock.get();
-
-            reply2 = processFindAndModify(&queryImpl, findAndModifyRequest2);
+            auto [findAndModifyRequest2, reply2, expCtx] = *sharedFindAndModifyBlock.get();
 
+            reply2 = processFindAndModify(expCtx, &queryImpl, findAndModifyRequest2);
 
             if (MONGO_unlikely(fleCrudHangFindAndModify.shouldFail())) {
                 LOGV2(6371900, "Hanging due to fleCrudHangFindAndModify fail point");
@@ -855,6 +857,7 @@ FLEBatchResult processFLEBatch(OperationContext* opCtx,
 
 // See processUpdate for algorithm overview
 write_ops::FindAndModifyCommandReply processFindAndModify(
+    boost::intrusive_ptr<ExpressionContext> expCtx,
     FLEQueryInterface* queryImpl,
     const write_ops::FindAndModifyCommandRequest& findAndModifyRequest) {
 
@@ -865,9 +868,19 @@ write_ops::FindAndModifyCommandReply processFindAndModify(
     auto tokenMap = EncryptionInformationHelpers::getDeleteTokens(edcNss, ei);
     int32_t stmtId = findAndModifyRequest.getStmtId().value_or(0);
 
+    auto newFindAndModifyRequest = findAndModifyRequest;
+
+    // Step 0 ----
+    // Rewrite filter
+    newFindAndModifyRequest.setQuery(fle::rewriteEncryptedFilterInsideTxn(
+        queryImpl, edcNss.db(), efc, expCtx, findAndModifyRequest.getQuery()));
+
+    // Make sure not to inherit the command's writeConcern, this should be set at the transaction
+    // level.
+    newFindAndModifyRequest.setWriteConcern(boost::none);
+
     // Step 1 ----
     // If we have an update object, we have to process for ESC
-    auto newFindAndModifyRequest = findAndModifyRequest;
     if (findAndModifyRequest.getUpdate().has_value()) {
 
         std::vector<EDCServerPayloadInfo> serverPayload;
@@ -973,7 +986,6 @@ write_ops::FindAndModifyCommandReply processFindAndModify(
 }
 
 FLEBatchResult processFLEFindAndModify(OperationContext* opCtx,
-                                       const std::string& dbName,
                                        const BSONObj& cmdObj,
                                        BSONObjBuilder& result) {
     // There is no findAndModify parsing in mongos so we need to first parse to decide if it is for
diff --git a/src/mongo/db/fle_crud.h b/src/mongo/db/fle_crud.h
index 54439d34273..303c8704015 100644
--- a/src/mongo/db/fle_crud.h
+++ b/src/mongo/db/fle_crud.h
@@ -106,7 +106,6 @@ write_ops::UpdateCommandReply processFLEUpdate(
  * Process a findAndModify request from mongos
  */
 FLEBatchResult processFLEFindAndModify(OperationContext* opCtx,
-                                       const std::string& dbName,
                                        const BSONObj& cmdObj,
                                        BSONObjBuilder& result);
 
@@ -356,6 +355,7 @@ write_ops::UpdateCommandReply processUpdate(FLEQueryInterface* queryImpl,
  * Used by unit tests.
  */
 write_ops::FindAndModifyCommandReply processFindAndModify(
+    boost::intrusive_ptr<ExpressionContext> expCtx,
     FLEQueryInterface* queryImpl,
     const write_ops::FindAndModifyCommandRequest& findAndModifyRequest);
 
diff --git a/src/mongo/db/fle_crud_test.cpp b/src/mongo/db/fle_crud_test.cpp
index d84aa78d394..8b89b500c3b 100644
--- a/src/mongo/db/fle_crud_test.cpp
+++ b/src/mongo/db/fle_crud_test.cpp
@@ -547,7 +547,13 @@ void FleCrudTest::doFindAndModify(write_ops::FindAndModifyCommandRequest& reques
 
     request.setEncryptionInformation(ei);
 
-    processFindAndModify(_queryImpl.get(), request);
+    std::unique_ptr<CollatorInterface> collator;
+    auto expCtx = make_intrusive<ExpressionContext>(_opCtx.get(),
+                                                    std::move(collator),
+                                                    request.getNamespace(),
+                                                    request.getLegacyRuntimeConstants(),
+                                                    request.getLet());
+    processFindAndModify(expCtx, _queryImpl.get(), request);
 }
 
 class CollectionReader : public FLEStateCollectionReader {
diff --git a/src/mongo/db/fle_query_interface_mock.cpp b/src/mongo/db/fle_query_interface_mock.cpp
index c6268fc05c2..2aeb39788dd 100644
--- a/src/mongo/db/fle_query_interface_mock.cpp
+++ b/src/mongo/db/fle_query_interface_mock.cpp
@@ -151,20 +151,24 @@ write_ops::FindAndModifyCommandReply FLEQueryInterfaceMock::findAndModify(
             "findAndModify 'new' field must be 'false'",
             findAndModifyRequest.getNew().get_value_or(false) == false);
 
-    BSONObj preimage = getById(nss, findAndModifyRequest.getQuery().firstElement());
+    // The query may be the short form {_id: 1} or the long form {_id: {$eq: 1}}.
+    auto idElt = [&]() {
+        auto id = findAndModifyRequest.getQuery().firstElement();
+        if (id.type() == BSONType::Object && id.Obj().hasField("$eq")) {
+            return id.Obj()["$eq"];
+        }
+        return id;
+    }();
+    BSONObj preimage = getById(nss, idElt);
 
     if (findAndModifyRequest.getRemove().get_value_or(false)) {
         // Remove
-        auto swDoc =
-            _storage->deleteById(_opCtx, nss, findAndModifyRequest.getQuery().firstElement());
+        auto swDoc = _storage->deleteById(_opCtx, nss, idElt);
         uassertStatusOK(swDoc);
 
     } else {
-        uassertStatusOK(
-            _storage->upsertById(_opCtx,
-                                 nss,
-                                 findAndModifyRequest.getQuery().firstElement(),
-                                 findAndModifyRequest.getUpdate()->getUpdateModifier()));
+        uassertStatusOK(_storage->upsertById(
+            _opCtx, nss, idElt, findAndModifyRequest.getUpdate()->getUpdateModifier()));
     }
 
     write_ops::FindAndModifyCommandReply reply;
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
index e374f3adb5b..1d6e8cafa0d 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -441,7 +441,7 @@ public:
              BSONObjBuilder& result) override {
         const NamespaceString nss(CommandHelpers::parseNsCollectionRequired(dbName, cmdObj));
 
-        if (processFLEFindAndModify(opCtx, dbName, cmdObj, result) == FLEBatchResult::kProcessed) {
+        if (processFLEFindAndModify(opCtx, cmdObj, result) == FLEBatchResult::kProcessed) {
             return true;
         }
 
@@ -499,6 +499,7 @@ private:
                             const BSONObj& cmdObj,
                             BSONObjBuilder* result) {
         bool isRetryableWrite = opCtx->getTxnNumber() && !TransactionRouter::get(opCtx);
+
         const auto response = [&] {
             std::vector<AsyncRequestsSender::Request> requests;
             BSONObj filteredCmdObj = CommandHelpers::filterCommandRequestForPassthrough(cmdObj);