SERVER-67263 Reject InsertUpdatePayload with mismatched IndexKeyId

1 files changed, 71 insertions, 16 deletions

diff --git a/src/mongo/db/fle_crud.cpp b/src/mongo/db/fle_crud.cpp
index 2d7e59abb2e..8739ab7894d 100644
--- a/src/mongo/db/fle_crud.cpp
+++ b/src/mongo/db/fle_crud.cpp
@@ -187,6 +187,31 @@ boost::intrusive_ptr<ExpressionContext> makeExpCtx(OperationContext* opCtx,
 
 }  // namespace
 
+/**
+ * Checks that all encrypted payloads correspond to an encrypted field,
+ * and that the encryption keyId used was appropriate for that field.
+ */
+void validateInsertUpdatePayloads(const std::vector<EncryptedField>& fields,
+                                  const std::vector<EDCServerPayloadInfo>& payload) {
+    std::map<StringData, UUID> pathToKeyIdMap;
+    for (const auto& field : fields) {
+        pathToKeyIdMap.insert({field.getPath(), field.getKeyId()});
+    }
+
+    for (const auto& field : payload) {
+        auto fieldPath = field.fieldPathName;
+        auto expect = pathToKeyIdMap.find(fieldPath);
+        uassert(6726300,
+                str::stream() << "Field '" << fieldPath << "' is unexpectedly encrypted",
+                expect != pathToKeyIdMap.end());
+        auto indexKeyId = field.payload.getIndexKeyId();
+        uassert(6726301,
+                str::stream() << "Mismatched keyId for field '" << fieldPath << "' expected "
+                              << expect->second << ", found " << indexKeyId,
+                indexKeyId == expect->second);
+    }
+}
+
 std::pair<FLEBatchResult, write_ops::InsertCommandReply> processInsert(
     OperationContext* opCtx,
     const write_ops::InsertCommandRequest& insertRequest,
@@ -218,6 +243,8 @@ std::pair<FLEBatchResult, write_ops::InsertCommandReply> processInsert(
             FLEBatchResult::kNotProcessed, write_ops::InsertCommandReply()};
     }
 
+    validateInsertUpdatePayloads(efc.getFields(), *serverPayload);
+
     auto reply = std::make_shared<write_ops::InsertCommandReply>();
 
     uint32_t stmtId = getStmtIdForWriteAt(insertRequest, 0);
@@ -633,39 +660,65 @@ std::shared_ptr<write_ops::FindAndModifyCommandRequest> constructDefaultReply()
     return std::make_shared<write_ops::FindAndModifyCommandRequest>(NamespaceString());
 }
 
-}  // namespace
-
-template <typename ReplyType>
-StatusWith<std::pair<ReplyType, OpMsgRequest>> processFindAndModifyRequest(
-    OperationContext* opCtx,
-    const write_ops::FindAndModifyCommandRequest& findAndModifyRequest,
-    GetTxnCallback getTxns,
-    ProcessFindAndModifyCallback<ReplyType> processCallback) {
-
-    // Is this a delete
-    bool isDelete = findAndModifyRequest.getRemove().value_or(false);
+/**
+ * Extracts update payloads from a {findAndModify: nss, ...} request,
+ * and proxies to `validateInsertUpdatePayload()`.
+ */
+void validateFindAndModifyRequest(const write_ops::FindAndModifyCommandRequest& request) {
+    // Is this a delete?
+    const bool isDelete = request.getRemove().value_or(false);
 
     // User can only specify either remove = true or update != {}
     uassert(6371401,
             "Must specify either update or remove to findAndModify, not both",
-            !(findAndModifyRequest.getUpdate().has_value() && isDelete));
+            !(request.getUpdate().has_value() && isDelete));
 
     uassert(6371402,
             "findAndModify with encryption only supports new: false",
-            findAndModifyRequest.getNew().value_or(false) == false);
+            request.getNew().value_or(false) == false);
 
     uassert(6371408,
             "findAndModify fields must be empty",
-            findAndModifyRequest.getFields().value_or(BSONObj()).isEmpty());
+            request.getFields().value_or(BSONObj()).isEmpty());
 
     // pipeline - is agg specific, delta is oplog, transform is internal (timeseries)
-    auto updateModicationType =
-        findAndModifyRequest.getUpdate().value_or(write_ops::UpdateModification()).type();
+    auto updateMod = request.getUpdate().get_value_or({});
+    const auto updateModicationType = updateMod.type();
+
     uassert(6439901,
             "FLE only supports modifier and replacement style updates",
             updateModicationType == write_ops::UpdateModification::Type::kModifier ||
                 updateModicationType == write_ops::UpdateModification::Type::kReplacement);
 
+    auto nss = request.getNamespace();
+    auto ei = request.getEncryptionInformation().get();
+    auto efc = EncryptionInformationHelpers::getAndValidateSchema(nss, ei);
+
+    BSONObj update;
+    if (updateMod.type() == write_ops::UpdateModification::Type::kReplacement) {
+        update = updateMod.getUpdateReplacement();
+    } else {
+        invariant(updateMod.type() == write_ops::UpdateModification::Type::kModifier);
+        update = updateMod.getUpdateModifier().getObjectField("$set"_sd);
+    }
+
+    if (!update.firstElement().eoo()) {
+        auto serverPayload = EDCServerCollection::getEncryptedFieldInfo(update);
+        validateInsertUpdatePayloads(efc.getFields(), serverPayload);
+    }
+}
+
+}  // namespace
+
+template <typename ReplyType>
+StatusWith<std::pair<ReplyType, OpMsgRequest>> processFindAndModifyRequest(
+    OperationContext* opCtx,
+    const write_ops::FindAndModifyCommandRequest& findAndModifyRequest,
+    GetTxnCallback getTxns,
+    ProcessFindAndModifyCallback<ReplyType> processCallback) {
+
+    validateFindAndModifyRequest(findAndModifyRequest);
+
     std::shared_ptr<txn_api::SyncTransactionWithRetries> trun = getTxns(opCtx);
 
     // The function that handles the transaction may outlive this function so we need to use
@@ -837,6 +890,7 @@ write_ops::UpdateCommandReply processUpdate(FLEQueryInterface* queryImpl,
         auto setObject = updateModifier.getObjectField("$set");
         EDCServerCollection::validateEncryptedFieldInfo(setObject, efc, bypassDocumentValidation);
         serverPayload = EDCServerCollection::getEncryptedFieldInfo(setObject);
+        validateInsertUpdatePayloads(efc.getFields(), serverPayload);
 
         processFieldsForInsert(
             queryImpl, edcNss, serverPayload, efc, &stmtId, bypassDocumentValidation);
@@ -851,6 +905,7 @@ write_ops::UpdateCommandReply processUpdate(FLEQueryInterface* queryImpl,
         EDCServerCollection::validateEncryptedFieldInfo(
             replacementDocument, efc, bypassDocumentValidation);
         serverPayload = EDCServerCollection::getEncryptedFieldInfo(replacementDocument);
+        validateInsertUpdatePayloads(efc.getFields(), serverPayload);
 
         processFieldsForInsert(
             queryImpl, edcNss, serverPayload, efc, &stmtId, bypassDocumentValidation);