SERVER-67263 Reject InsertUpdatePayload with mismatched IndexKeyId

3 files changed, 117 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);
diff --git a/src/mongo/db/fle_crud.h b/src/mongo/db/fle_crud.h
index c91f7c1eee4..9175f2ff3f9 100644
--- a/src/mongo/db/fle_crud.h
+++ b/src/mongo/db/fle_crud.h
@@ -488,4 +488,8 @@ processFindAndModifyRequest<write_ops::FindAndModifyCommandRequest>(
 write_ops::UpdateCommandReply processUpdate(OperationContext* opCtx,
                                             const write_ops::UpdateCommandRequest& updateRequest,
                                             GetTxnCallback getTxns);
+
+void validateInsertUpdatePayloads(const std::vector<EncryptedField>& fields,
+                                  const std::vector<EDCServerPayloadInfo>& payload);
+
 }  // namespace mongo
diff --git a/src/mongo/db/fle_crud_test.cpp b/src/mongo/db/fle_crud_test.cpp
index 4b277d60833..764d8846ea2 100644
--- a/src/mongo/db/fle_crud_test.cpp
+++ b/src/mongo/db/fle_crud_test.cpp
@@ -1105,6 +1105,48 @@ TEST_F(FleCrudTest, FindAndModify_SetSafeContent) {
     ASSERT_THROWS_CODE(doFindAndModify(req), DBException, 6666200);
 }
 
+BSONObj makeInsertUpdatePayload(StringData path, const UUID& uuid) {
+    // Actual values don't matter for these tests (apart from indexKeyId).
+    auto bson = FLE2InsertUpdatePayload({}, {}, {}, {}, uuid, BSONType::String, {}, {}).toBSON();
+    std::vector<std::uint8_t> bindata;
+    bindata.resize(bson.objsize() + 1);
+    bindata[0] = static_cast<std::uint8_t>(EncryptedBinDataType::kFLE2InsertUpdatePayload);
+    memcpy(bindata.data() + 1, bson.objdata(), bson.objsize());
+
+    BSONObjBuilder bob;
+    bob.appendBinData(path, bindata.size(), BinDataType::Encrypt, bindata.data());
+    return bob.obj();
+}
+
+TEST(FleCrudTest, validateIndexKeyValid) {
+    // This test assumes we have at least one field in EFC.
+    auto fields = getTestEncryptedFieldConfig().getFields();
+    ASSERT_GTE(fields.size(), 1);
+    auto field = fields[0];
+
+    auto validInsert = makeInsertUpdatePayload(field.getPath(), field.getKeyId());
+    auto validPayload = EDCServerCollection::getEncryptedFieldInfo(validInsert);
+    validateInsertUpdatePayloads(fields, validPayload);
+}
+
+TEST(FleCrudTest, validateIndexKeyInvalid) {
+    // This test assumes we have at least one field in EFC.
+    auto fields = getTestEncryptedFieldConfig().getFields();
+    ASSERT_GTE(fields.size(), 1);
+    auto field = fields[0];
+
+    auto invalidInsert = makeInsertUpdatePayload(field.getPath(), UUID::gen());
+    auto invalidPayload = EDCServerCollection::getEncryptedFieldInfo(invalidInsert);
+    ASSERT_THROWS_WITH_CHECK(validateInsertUpdatePayloads(fields, invalidPayload),
+                             DBException,
+                             [&](const DBException& ex) {
+                                 ASSERT_STRING_CONTAINS(ex.what(),
+                                                        str::stream()
+                                                            << "Mismatched keyId for field '"
+                                                            << field.getPath() << "'");
+                             });
+}
+
 TEST_F(FleTagsTest, InsertOne) {
     auto doc = BSON("encrypted"
                     << "a");