SERVER-64055 fle find rewrite on mongod

10 files changed, 71 insertions, 34 deletions

diff --git a/buildscripts/resmokeconfig/suites/fle2.yml b/buildscripts/resmokeconfig/suites/fle2.yml
index 4f75d04ee90..afce3dd15f9 100644
--- a/buildscripts/resmokeconfig/suites/fle2.yml
+++ b/buildscripts/resmokeconfig/suites/fle2.yml
@@ -2,8 +2,9 @@ test_kind: js_test
 selector:
   roots:
   - jstests/fle2/**/*.js
-  - src/mongo/db/modules/*/jstests/fle2/*.js
-
+  - src/mongo/db/modules/*/jstests/fle2/**/*.js
+  exclude_files:
+  - src/mongo/db/modules/*/jstests/fle2/query/aggregate*.js # TODO: SERVER-63312
 executor:
   archive:
     hooks:
diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp
index 60910d9df44..9f596d34a4f 100644
--- a/src/mongo/db/commands/find_cmd.cpp
+++ b/src/mongo/db/commands/find_cmd.cpp
@@ -42,6 +42,7 @@
 #include "mongo/db/cursor_manager.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/fle_crud.h"
 #include "mongo/db/matcher/extensions_callback_real.h"
 #include "mongo/db/pipeline/aggregation_request_helper.h"
 #include "mongo/db/pipeline/variables.h"
@@ -119,6 +120,11 @@ std::unique_ptr<FindCommandRequest> parseCmdObjectToFindCommandRequest(Operation
         std::move(nss),
         APIParameters::get(opCtx).getAPIStrict().value_or(false));
 
+    // Rewrite any FLE find payloads that exist in the query if this is a FLE 2 query.
+    if (shouldDoFLERewrite(findCommand)) {
+        processFLEFindD(opCtx, findCommand.get());
+    }
+
     return translateNtoReturnToLimitOrBatchSize(std::move(findCommand));
 }
 
diff --git a/src/mongo/db/fle_crud.cpp b/src/mongo/db/fle_crud.cpp
index a3f46bd6dac..34a1c4eb6ec 100644
--- a/src/mongo/db/fle_crud.cpp
+++ b/src/mongo/db/fle_crud.cpp
@@ -1209,4 +1209,8 @@ write_ops::FindAndModifyCommandReply FLEQueryInterfaceImpl::findAndModify(
     return write_ops::FindAndModifyCommandReply::parse(IDLParserErrorContext("reply"), response);
 }
 
+void processFLEFindS(OperationContext* opCtx, FindCommandRequest* findCommand) {
+    fle::processFindCommand(opCtx, findCommand, &getTransactionWithRetriesForMongoS);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/fle_crud.h b/src/mongo/db/fle_crud.h
index 3ec46fcf0be..299b1410d42 100644
--- a/src/mongo/db/fle_crud.h
+++ b/src/mongo/db/fle_crud.h
@@ -115,6 +115,28 @@ FLEBatchResult processFLEFindAndModify(OperationContext* opCtx,
 write_ops::FindAndModifyCommandReply processFLEFindAndModify(
     OperationContext* opCtx, const write_ops::FindAndModifyCommandRequest& findAndModifyRequest);
 
+/**
+ * Process a find command from mongos.
+ */
+void processFLEFindS(OperationContext* opCtx, FindCommandRequest* findCommand);
+
+/**
+ * Process a find command from a replica set.
+ */
+void processFLEFindD(OperationContext* opCtx, FindCommandRequest* findCommand);
+
+/**
+ * Helper function to determine if an IDL object with encryption information should be rewritten.
+ */
+template <typename T>
+bool shouldDoFLERewrite(const std::unique_ptr<T>& cmd) {
+    return gFeatureFlagFLE2.isEnabledAndIgnoreFCV() && cmd->getEncryptionInformation();
+}
+
+template <typename T>
+bool shouldDoFLERewrite(const T& cmd) {
+    return gFeatureFlagFLE2.isEnabledAndIgnoreFCV() && cmd.getEncryptionInformation();
+}
 
 /**
  * Abstraction layer for FLE
diff --git a/src/mongo/db/fle_crud_mongod.cpp b/src/mongo/db/fle_crud_mongod.cpp
index 635e94c851a..b15fda8cbb1 100644
--- a/src/mongo/db/fle_crud_mongod.cpp
+++ b/src/mongo/db/fle_crud_mongod.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/ops/write_ops_gen.h"
 #include "mongo/db/ops/write_ops_parsers.h"
 #include "mongo/db/query/find_command_gen.h"
+#include "mongo/db/query/fle/server_rewrite.h"
 #include "mongo/db/repl/repl_client_info.h"
 #include "mongo/db/session.h"
 #include "mongo/db/session_catalog.h"
@@ -232,4 +233,8 @@ write_ops::UpdateCommandReply processFLEUpdate(
     return updateReply;
 }
 
+void processFLEFindD(OperationContext* opCtx, FindCommandRequest* findCommand) {
+    fle::processFindCommand(opCtx, findCommand, &getTransactionWithRetriesForMongoD);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index e36aa6911b3..4d958aad302 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -134,9 +134,6 @@ StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize(
 
     std::unique_ptr<MatchExpression> me = std::move(statusWithMatcher.getValue());
 
-    // TODO: SERVER-64055 if encryptionInformation is present, rewrite MatchExpression FLE find
-    // payloads.
-
     Status initStatus =
         cq->init(opCtx,
                  std::move(newExpCtx),
diff --git a/src/mongo/db/query/fle/server_rewrite.cpp b/src/mongo/db/query/fle/server_rewrite.cpp
index 36ad96a1d4e..63e5e2fa87e 100644
--- a/src/mongo/db/query/fle/server_rewrite.cpp
+++ b/src/mongo/db/query/fle/server_rewrite.cpp
@@ -29,6 +29,9 @@
 
 
 #include "mongo/db/query/fle/server_rewrite.h"
+
+#include <memory>
+
 #include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/bson/bsontypes.h"
 #include "mongo/crypto/encryption_fields_gen.h"
@@ -39,7 +42,6 @@
 #include "mongo/db/pipeline/document_source_graph_lookup.h"
 #include "mongo/db/pipeline/document_source_match.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
-#include "mongo/db/transaction_api.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/transaction_router_resource_yielder.h"
 #include "mongo/util/assert_util.h"
@@ -153,11 +155,10 @@ public:
 // executor, and so we can't pass data by reference into the lambda. The provided rewriter should
 // hold all the data we need to do the rewriting inside the lambda, and is passed in a more
 // threadsafe shared_ptr. The result of applying the rewrites can be accessed in the RewriteBase.
-void doFLERewriteInTxn(OperationContext* opCtx, std::shared_ptr<RewriteBase> sharedBlock) {
-    auto txn = std::make_shared<txn_api::TransactionWithRetries>(
-        opCtx,
-        Grid::get(opCtx)->getExecutorPool()->getFixedExecutor(),
-        TransactionRouterResourceYielder::makeForLocalHandoff());
+void doFLERewriteInTxn(OperationContext* opCtx,
+                       std::shared_ptr<RewriteBase> sharedBlock,
+                       GetTxnCallback getTxn) {
+    auto txn = getTxn(opCtx);
 
     auto swCommitResult = txn->runSyncNoThrow(
         opCtx, [sharedBlock](const txn_api::TransactionClient& txnClient, auto txnExec) {
@@ -202,16 +203,17 @@ BSONObj rewriteEncryptedFilterInsideTxn(FLEQueryInterface* queryImpl,
 
 
 void processFindCommand(OperationContext* opCtx,
-                        NamespaceString nss,
-                        FindCommandRequest* findCommand) {
+                        FindCommandRequest* findCommand,
+                        GetTxnCallback getTransaction) {
+    invariant(findCommand->getNamespaceOrUUID().nss());
     invariant(findCommand->getEncryptionInformation());
 
     auto sharedBlock =
         std::make_shared<FilterRewrite>(makeExpCtx(opCtx, findCommand),
-                                        nss,
+                                        findCommand->getNamespaceOrUUID().nss().get(),
                                         findCommand->getEncryptionInformation().get(),
                                         findCommand->getFilter().getOwned());
-    doFLERewriteInTxn(opCtx, sharedBlock);
+    doFLERewriteInTxn(opCtx, sharedBlock, getTransaction);
 
     auto rewrittenFilter = sharedBlock->rewrittenFilter.getOwned();
     findCommand->setFilter(std::move(rewrittenFilter));
@@ -225,7 +227,13 @@ std::unique_ptr<Pipeline, PipelineDeleter> processPipeline(
     std::unique_ptr<Pipeline, PipelineDeleter> toRewrite) {
 
     auto sharedBlock = std::make_shared<PipelineRewrite>(nss, encryptInfo, std::move(toRewrite));
-    doFLERewriteInTxn(opCtx, sharedBlock);
+    doFLERewriteInTxn(opCtx, sharedBlock, [](auto* opCtx) {
+        // TODO: SERVER-63312 pass in the right transaction callback for mongod.
+        return std::make_shared<txn_api::TransactionWithRetries>(
+            opCtx,
+            Grid::get(opCtx)->getExecutorPool()->getFixedExecutor(),
+            TransactionRouterResourceYielder::makeForLocalHandoff());
+    });
 
     return sharedBlock->getPipeline();
 }
diff --git a/src/mongo/db/query/fle/server_rewrite.h b/src/mongo/db/query/fle/server_rewrite.h
index f0c3d304866..8d1da86009b 100644
--- a/src/mongo/db/query/fle/server_rewrite.h
+++ b/src/mongo/db/query/fle/server_rewrite.h
@@ -35,29 +35,23 @@
 
 #include "mongo/bson/bsonobj.h"
 #include "mongo/crypto/fle_crypto.h"
+#include "mongo/db/fle_crud.h"
 #include "mongo/db/matcher/expression_parser.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/transaction_api.h"
 
 namespace mongo {
 class FLEQueryInterface;
 namespace fle {
 
 /**
- * Helper function to determine if an IDL object with encryption information should be rewritten.
- */
-template <typename T>
-bool shouldRewrite(const T& cmd) {
-    return gFeatureFlagFLE2.isEnabledAndIgnoreFCV() && cmd->getEncryptionInformation();
-}
-
-/**
  * Process a find command with encryptionInformation in-place, rewriting the filter condition so
  * that any query on an encrypted field will properly query the underlying tags array.
  */
 void processFindCommand(OperationContext* opCtx,
-                        NamespaceString nss,
-                        FindCommandRequest* findCommand);
+                        FindCommandRequest* findCommand,
+                        GetTxnCallback txn);
 
 /**
  * Process a pipeline with encryptionInformation by rewriting the pipeline to query against the
diff --git a/src/mongo/s/commands/cluster_find_cmd.h b/src/mongo/s/commands/cluster_find_cmd.h
index 7d0ca61a0d3..888da770530 100644
--- a/src/mongo/s/commands/cluster_find_cmd.h
+++ b/src/mongo/s/commands/cluster_find_cmd.h
@@ -35,9 +35,9 @@
 #include "mongo/db/auth/authorization_checks.h"
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/commands.h"
+#include "mongo/db/fle_crud.h"
 #include "mongo/db/matcher/extensions_callback_noop.h"
 #include "mongo/db/query/cursor_response.h"
-#include "mongo/db/query/fle/server_rewrite.h"
 #include "mongo/db/stats/counters.h"
 #include "mongo/db/views/resolved_view.h"
 #include "mongo/rpc/get_status_from_command_result.h"
@@ -128,9 +128,6 @@ public:
                      rpc::ReplyBuilderInterface* result) override {
             // Parse the command BSON to a FindCommandRequest.
             auto findCommand = _parseCmdObjectToFindCommandRequest(opCtx, ns(), _request.body);
-            if (fle::shouldRewrite(findCommand)) {
-                fle::processFindCommand(opCtx, ns(), findCommand.get());
-            }
 
             try {
                 const auto explainCmd =
@@ -206,9 +203,6 @@ public:
             });
 
             auto findCommand = _parseCmdObjectToFindCommandRequest(opCtx, ns(), _request.body);
-            if (fle::shouldRewrite(findCommand)) {
-                fle::processFindCommand(opCtx, ns(), findCommand.get());
-            }
 
             const boost::intrusive_ptr<ExpressionContext> expCtx;
             auto cq = uassertStatusOK(
@@ -293,6 +287,11 @@ public:
             uassert(5746101,
                     "Cannot specify ntoreturn in a find command against mongos",
                     findCommand->getNtoreturn() == boost::none);
+
+            if (shouldDoFLERewrite(findCommand)) {
+                processFLEFindS(opCtx, findCommand.get());
+            }
+
             return findCommand;
         }
 
diff --git a/src/mongo/s/query/cluster_aggregate.cpp b/src/mongo/s/query/cluster_aggregate.cpp
index 3c754850628..bd291ed6f70 100644
--- a/src/mongo/s/query/cluster_aggregate.cpp
+++ b/src/mongo/s/query/cluster_aggregate.cpp
@@ -40,6 +40,7 @@
 #include "mongo/db/client.h"
 #include "mongo/db/commands.h"
 #include "mongo/db/curop.h"
+#include "mongo/db/fle_crud.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/pipeline/document_source_change_stream.h"
 #include "mongo/db/pipeline/document_source_internal_unpack_bucket.h"
@@ -299,7 +300,7 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
         opCtx, isSharded, request.getExplain(), serverGlobalParams.enableMajorityReadConcern);
     auto hasChangeStream = liteParsedPipeline.hasChangeStream();
     auto involvedNamespaces = liteParsedPipeline.getInvolvedNamespaces();
-    auto shouldDoFLERewrite = fle::shouldRewrite(&request);
+    auto shouldDoFLERewrite = ::mongo::shouldDoFLERewrite(request);
 
     uassert(6256300,
             str::stream() << "On mongos, " << AggregateCommandRequest::kCollectionUUIDFieldName