SERVER-63293 Construct binary tags from payload

8 files changed, 393 insertions, 30 deletions

diff --git a/src/mongo/crypto/SConscript b/src/mongo/crypto/SConscript
index 2e84f15e626..1b349cb440e 100644
--- a/src/mongo/crypto/SConscript
+++ b/src/mongo/crypto/SConscript
@@ -94,6 +94,7 @@ env.Library(
     source=[
         "encryption_fields_util.cpp",
         "fle_crypto.cpp",
+        "fle_tags.cpp",
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/base/secure_allocator',
diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index 138066b409c..5c9e74c8858 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -446,7 +446,7 @@ StatusWith<std::tuple<T1, T2>> decryptAndUnpack(ConstDataRange cdr, FLEToken<Tok
 
 
 template <typename collectionT, typename tagTokenT, typename valueTokenT>
-boost::optional<uint64_t> emuBinaryCommon(FLEStateCollectionReader* reader,
+boost::optional<uint64_t> emuBinaryCommon(const FLEStateCollectionReader& reader,
                                           tagTokenT tagToken,
                                           valueTokenT valueToken) {
 
@@ -458,7 +458,7 @@ boost::optional<uint64_t> emuBinaryCommon(FLEStateCollectionReader* reader,
     // Search for null record
     PrfBlock nullRecordId = collectionT::generateId(tagToken, boost::none);
 
-    BSONObj nullDoc = reader->getById(nullRecordId);
+    BSONObj nullDoc = reader.getById(nullRecordId);
 
     if (!nullDoc.isEmpty()) {
         auto swNullEscDoc = collectionT::decryptNullDocument(valueToken, nullDoc);
@@ -471,7 +471,7 @@ boost::optional<uint64_t> emuBinaryCommon(FLEStateCollectionReader* reader,
     }
 
     // step 4, 5: get document count
-    uint64_t rho = reader->getDocumentCount();
+    uint64_t rho = reader.getDocumentCount();
 
 #ifdef DEBUG_ENUM_BINARY
     std::cout << fmt::format("start: lambda: {}, i: {}, rho: {}", lambda, i, rho) << std::endl;
@@ -485,7 +485,7 @@ boost::optional<uint64_t> emuBinaryCommon(FLEStateCollectionReader* reader,
     // condition
     while (flag) {
         // 7 a
-        BSONObj doc = reader->getById(collectionT::generateId(tagToken, rho + lambda));
+        BSONObj doc = reader.getById(collectionT::generateId(tagToken, rho + lambda));
 
 #ifdef DEBUG_ENUM_BINARY
         std::cout << fmt::format("search1: rho: {},  doc: {}", rho, doc.toString()) << std::endl;
@@ -515,7 +515,7 @@ boost::optional<uint64_t> emuBinaryCommon(FLEStateCollectionReader* reader,
 
 
         // 9b
-        BSONObj doc = reader->getById(collectionT::generateId(tagToken, median + lambda));
+        BSONObj doc = reader.getById(collectionT::generateId(tagToken, median + lambda));
 
 #ifdef DEBUG_ENUM_BINARY
         std::cout << fmt::format("search_stat: min: {}, median: {}, max: {}, i: {}, doc: {}",
@@ -546,7 +546,7 @@ boost::optional<uint64_t> emuBinaryCommon(FLEStateCollectionReader* reader,
             // explicitly
             if (j == maxIterations && min == 1) {
                 // 9 d ii A
-                BSONObj doc = reader->getById(collectionT::generateId(tagToken, 1 + lambda));
+                BSONObj doc = reader.getById(collectionT::generateId(tagToken, 1 + lambda));
                 // 9 d ii B
                 if (!doc.isEmpty()) {
                     i = 1 + lambda;
@@ -1494,6 +1494,11 @@ BSONObj ESCCollection::generateCompactionPlaceholderDocument(ESCTwiceDerivedTagT
 
 StatusWith<ESCNullDocument> ESCCollection::decryptNullDocument(ESCTwiceDerivedValueToken valueToken,
                                                                BSONObj& doc) {
+    return ESCCollection::decryptNullDocument(valueToken, std::move(doc));
+}
+
+StatusWith<ESCNullDocument> ESCCollection::decryptNullDocument(ESCTwiceDerivedValueToken valueToken,
+                                                               BSONObj&& doc) {
     BSONElement encryptedValue;
     auto status = bsonExtractTypedField(doc, kValue, BinData, &encryptedValue);
     if (!status.isOK()) {
@@ -1511,9 +1516,13 @@ StatusWith<ESCNullDocument> ESCCollection::decryptNullDocument(ESCTwiceDerivedVa
     return ESCNullDocument{std::get<0>(value), std::get<1>(value)};
 }
 
-
 StatusWith<ESCDocument> ESCCollection::decryptDocument(ESCTwiceDerivedValueToken valueToken,
                                                        BSONObj& doc) {
+    return ESCCollection::decryptDocument(valueToken, std::move(doc));
+}
+
+StatusWith<ESCDocument> ESCCollection::decryptDocument(ESCTwiceDerivedValueToken valueToken,
+                                                       BSONObj&& doc) {
     BSONElement encryptedValue;
     auto status = bsonExtractTypedField(doc, kValue, BinData, &encryptedValue);
     if (!status.isOK()) {
@@ -1533,7 +1542,7 @@ StatusWith<ESCDocument> ESCCollection::decryptDocument(ESCTwiceDerivedValueToken
 }
 
 
-boost::optional<uint64_t> ESCCollection::emuBinary(FLEStateCollectionReader* reader,
+boost::optional<uint64_t> ESCCollection::emuBinary(const FLEStateCollectionReader& reader,
                                                    ESCTwiceDerivedTagToken tagToken,
                                                    ESCTwiceDerivedValueToken valueToken) {
     return emuBinaryCommon<ESCCollection, ESCTwiceDerivedTagToken, ESCTwiceDerivedValueToken>(
@@ -1663,7 +1672,7 @@ StatusWith<ECCDocument> ECCCollection::decryptDocument(ECCTwiceDerivedValueToken
                        std::get<1>(value)};
 }
 
-boost::optional<uint64_t> ECCCollection::emuBinary(FLEStateCollectionReader* reader,
+boost::optional<uint64_t> ECCCollection::emuBinary(const FLEStateCollectionReader& reader,
                                                    ECCTwiceDerivedTagToken tagToken,
                                                    ECCTwiceDerivedValueToken valueToken) {
     return emuBinaryCommon<ECCCollection, ECCTwiceDerivedTagToken, ECCTwiceDerivedValueToken>(
diff --git a/src/mongo/crypto/fle_crypto.h b/src/mongo/crypto/fle_crypto.h
index de2b9e93c13..e7452a5ee5a 100644
--- a/src/mongo/crypto/fle_crypto.h
+++ b/src/mongo/crypto/fle_crypto.h
@@ -445,12 +445,12 @@ public:
      *
      * TODO - how perfect does it need to be ? Is too high or too low ok if it is just an estimate?
      */
-    virtual uint64_t getDocumentCount() = 0;
+    virtual uint64_t getDocumentCount() const = 0;
 
     /**
      * Get a document by its _id.
      */
-    virtual BSONObj getById(PrfBlock block) = 0;
+    virtual BSONObj getById(PrfBlock block) const = 0;
 };
 
 class ESCCollection {
@@ -499,15 +499,27 @@ public:
                                                            BSONObj& doc);
 
     /**
+     * Decrypt the null document.
+     */
+    static StatusWith<ESCNullDocument> decryptNullDocument(ESCTwiceDerivedValueToken valueToken,
+                                                           BSONObj&& doc);
+
+    /**
      * Decrypt a regular document.
      */
     static StatusWith<ESCDocument> decryptDocument(ESCTwiceDerivedValueToken valueToken,
                                                    BSONObj& doc);
 
     /**
+     * Decrypt a regular document.
+     */
+    static StatusWith<ESCDocument> decryptDocument(ESCTwiceDerivedValueToken valueToken,
+                                                   BSONObj&& doc);
+
+    /**
      * Search for the highest document id for a given field/value pair based on the token.
      */
-    static boost::optional<uint64_t> emuBinary(FLEStateCollectionReader* reader,
+    static boost::optional<uint64_t> emuBinary(const FLEStateCollectionReader& reader,
                                                ESCTwiceDerivedTagToken tagToken,
                                                ESCTwiceDerivedValueToken valueToken);
 };
@@ -655,7 +667,7 @@ public:
     /**
      * Search for the highest document id for a given field/value pair based on the token.
      */
-    static boost::optional<uint64_t> emuBinary(FLEStateCollectionReader* reader,
+    static boost::optional<uint64_t> emuBinary(const FLEStateCollectionReader& reader,
                                                ECCTwiceDerivedTagToken tagToken,
                                                ECCTwiceDerivedValueToken valueToken);
 };
diff --git a/src/mongo/crypto/fle_crypto_test.cpp b/src/mongo/crypto/fle_crypto_test.cpp
index a92c704550b..a949ab75f9b 100644
--- a/src/mongo/crypto/fle_crypto_test.cpp
+++ b/src/mongo/crypto/fle_crypto_test.cpp
@@ -409,7 +409,7 @@ public:
         // std::sort(_docs.begin(), _docs.end());
     }
 
-    BSONObj getById(PrfBlock id) override {
+    BSONObj getById(PrfBlock id) const override {
         for (const auto& doc : _docs) {
             auto el = doc.firstElement();
             int len;
@@ -424,7 +424,7 @@ public:
         return BSONObj();
     }
 
-    uint64_t getDocumentCount() override {
+    uint64_t getDocumentCount() const override {
         return _docs.size();
     }
 
@@ -454,7 +454,7 @@ TEST(FLE_ESC, EmuBinary_Empty) {
         FLETwiceDerivedTokenGenerator::generateESCTwiceDerivedValueToken(escDerivedToken);
 
 
-    auto i = ESCCollection::emuBinary(&coll, escTwiceTag, escTwiceValue);
+    auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
 
     ASSERT_TRUE(i.has_value());
     ASSERT_EQ(i.value(), 0);
@@ -486,7 +486,7 @@ TEST(FLE_ESC, EmuBinary) {
         coll.insert(doc);
     }
 
-    auto i = ESCCollection::emuBinary(&coll, escTwiceTag, escTwiceValue);
+    auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
 
     ASSERT_TRUE(i.has_value());
     ASSERT_EQ(i.value(), 5);
@@ -537,12 +537,12 @@ TEST(FLE_ESC, EmuBinary2) {
         coll.insert(doc);
     }
 
-    auto i = ESCCollection::emuBinary(&coll, escTwiceTag, escTwiceValue);
+    auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
 
     ASSERT_TRUE(i.has_value());
     ASSERT_EQ(i.value(), 13);
 
-    i = ESCCollection::emuBinary(&coll, escTwiceTag2, escTwiceValue2);
+    i = ESCCollection::emuBinary(coll, escTwiceTag2, escTwiceValue2);
 
     ASSERT_TRUE(i.has_value());
     ASSERT_EQ(i.value(), 5);
@@ -572,7 +572,7 @@ TEST(FLE_ESC, EmuBinary_NullRecord) {
     BSONObj doc = ESCCollection::generateNullDocument(escTwiceTag, escTwiceValue, 7, 7);
     coll.insert(doc);
 
-    auto i = ESCCollection::emuBinary(&coll, escTwiceTag, escTwiceValue);
+    auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
 
     ASSERT_FALSE(i.has_value());
 }
diff --git a/src/mongo/crypto/fle_tags.cpp b/src/mongo/crypto/fle_tags.cpp
new file mode 100644
index 00000000000..4422182d5dc
--- /dev/null
+++ b/src/mongo/crypto/fle_tags.cpp
@@ -0,0 +1,152 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include <boost/optional.hpp>
+
+#include "mongo/crypto/fle_crypto.h"
+#include "mongo/crypto/fle_tags.h"
+#include "mongo/stdx/unordered_set.h"
+
+namespace mongo::fle {
+
+using DerivedToken = FLEDerivedFromDataTokenAndContentionFactorTokenGenerator;
+using TwiceDerived = FLETwiceDerivedTokenGenerator;
+
+// The algorithm for constructing a list of tags matching an equality predicate on an encrypted
+// field is as follows:
+//
+// (1) Query ESC to obtain the counter value (n) after the most recent insert.
+// (2) Set the initial candidate set of tags to 1..n inclusive.
+//      (a) We know there have been (n) inserts into the collection for this field/value pair,
+//          however there could have been deletes, so this represents a superset of the true set of
+//          tags.
+// (3) Query ECC for a null document.
+//      (a) A null document means there has been at least one compaction. Therefore the deletes that
+//          were present in ECC before the compaction, have now been reflected in ESC.
+//      (b) No null document means there has not been a compaction. Therefore we have to check all
+//          the tags from 1..n to see if they have been deleted.
+// (4) Return the surviving set of tags from 1..n (encrypted).
+//
+// Given:
+//      s: ESCDerivedFromDataToken
+//      c: ECCDerivedFromDataToken
+//      d: EDCDerivedFromDataToken
+// Do:
+//      n = ESC::emuBinary(s)
+//      tags = [1..n]
+//      pos = ECC::nullDocument(c) ? ECC::nullDocument(c).position : 1
+//      LOOP:
+//          doc = ECC::getDocument(c, pos)
+//          if doc {
+//              tags = tags \ [doc.start..doc.end]
+//          } else {
+//              break
+//          }
+//          pos++
+//      return [EDC::encrypt(i) | i in tags]
+//
+// Note, a positive contention factor (cm) means we must repeat the above process (cm) times.
+std::vector<PrfBlock> readTags(const FLEStateCollectionReader& esc,
+                               const FLEStateCollectionReader& ecc,
+                               ESCDerivedFromDataToken s,
+                               ECCDerivedFromDataToken c,
+                               EDCDerivedFromDataToken d,
+                               boost::optional<int64_t> cm) {
+
+    std::vector<PrfBlock> binaryTags;
+
+    for (auto i = 0; cm && i <= cm.get(); i++) {
+        auto escTok = DerivedToken::generateESCDerivedFromDataTokenAndContentionFactorToken(s, i);
+        auto escTag = TwiceDerived::generateESCTwiceDerivedTagToken(escTok);
+        auto escVal = TwiceDerived::generateESCTwiceDerivedValueToken(escTok);
+
+        auto eccTok = DerivedToken::generateECCDerivedFromDataTokenAndContentionFactorToken(c, i);
+        auto eccTag = TwiceDerived::generateECCTwiceDerivedTagToken(eccTok);
+        auto eccVal = TwiceDerived::generateECCTwiceDerivedValueToken(eccTok);
+
+        auto edcTok = DerivedToken::generateEDCDerivedFromDataTokenAndContentionFactorToken(d, i);
+        auto edcTag = TwiceDerived::generateEDCTwiceDerivedToken(edcTok);
+
+        // (1) Query ESC for the counter value after the most recent insert.
+        //     0 => 0 inserts for this field value pair.
+        //     n => n inserts for this field value pair.
+        //     none => compaction => query ESC for null document to find # of inserts.
+        auto insertCounter = ESCCollection::emuBinary(esc, escTag, escVal);
+        if (insertCounter && insertCounter.get() == 0) {
+            continue;
+        }
+
+        stdx::unordered_set<int64_t> tags;
+
+        auto numInserts = insertCounter
+            ? uassertStatusOK(
+                  ESCCollection::decryptDocument(
+                      escVal, esc.getById(ESCCollection::generateId(escTag, insertCounter))))
+                  .count
+            : uassertStatusOK(
+                  ESCCollection::decryptNullDocument(
+                      escVal, esc.getById(ESCCollection::generateId(escTag, boost::none))))
+                  .count;
+
+        // (2) Set the initial set of tags to 1..n inclusive - a superset of the true tag set.
+        for (uint64_t i = 1; i <= numInserts; i++) {
+            tags.insert(i);
+        }
+
+        // (3) Query ECC for a null document.
+        auto eccNullDoc = ecc.getById(ECCCollection::generateId(eccTag, boost::none));
+        auto pos = eccNullDoc.isEmpty()
+            ? 1
+            : uassertStatusOK(ECCCollection::decryptNullDocument(eccVal, eccNullDoc)).position + 2;
+
+        // (3) Search ECC for deleted tag(counter) values.
+        while (true) {
+            auto eccObj = ecc.getById(ECCCollection::generateId(eccTag, pos));
+            if (eccObj.isEmpty()) {
+                break;
+            }
+            auto eccDoc = uassertStatusOK(ECCCollection::decryptDocument(eccVal, eccObj));
+            // Compaction placeholders only present for positive contention factors (cm).
+            if (eccDoc.valueType == ECCValueType::kCompactionPlaceholder) {
+                break;
+            }
+            for (uint64_t i = eccDoc.start; i <= eccDoc.end; i++) {
+                tags.erase(i);
+            }
+            pos++;
+        }
+
+        for (auto counter : tags) {
+            // (4) Derive binary tag values (encrypted) from the set of counter values (tags).
+            binaryTags.emplace_back(EDCServerCollection::generateTag(edcTag, counter));
+        }
+    }
+    return binaryTags;
+}
+}  // namespace mongo::fle
diff --git a/src/mongo/crypto/fle_tags.h b/src/mongo/crypto/fle_tags.h
new file mode 100644
index 00000000000..08cd7b1e392
--- /dev/null
+++ b/src/mongo/crypto/fle_tags.h
@@ -0,0 +1,45 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include <boost/optional.hpp>
+
+#include "mongo/crypto/fle_crypto.h"
+
+namespace mongo::fle {
+
+// Read a list of binary tags given ESC, ECC, and EDC derived tokens.
+std::vector<PrfBlock> readTags(const FLEStateCollectionReader& esc,
+                               const FLEStateCollectionReader& ecc,
+                               ESCDerivedFromDataToken s,
+                               ECCDerivedFromDataToken c,
+                               EDCDerivedFromDataToken d,
+                               boost::optional<int64_t> cm);
+}  // namespace mongo::fle
diff --git a/src/mongo/db/fle_crud.cpp b/src/mongo/db/fle_crud.cpp
index b281920069b..bbfda99161e 100644
--- a/src/mongo/db/fle_crud.cpp
+++ b/src/mongo/db/fle_crud.cpp
@@ -285,11 +285,11 @@ public:
     TxnCollectionReader(uint64_t count, FLEQueryInterface* queryImpl, const NamespaceString& nss)
         : _count(count), _queryImpl(queryImpl), _nss(nss) {}
 
-    uint64_t getDocumentCount() override {
+    uint64_t getDocumentCount() const override {
         return _count;
     }
 
-    BSONObj getById(PrfBlock block) override {
+    BSONObj getById(PrfBlock block) const override {
         auto doc = BSON("v" << BSONBinData(block.data(), block.size(), BinDataGeneral));
         BSONElement element = doc.firstElement();
         return _queryImpl->getById(_nss, element);
@@ -503,7 +503,7 @@ void processFieldsForInsert(FLEQueryInterface* queryImpl,
 
         int position = 1;
         int count = 1;
-        auto alpha = ESCCollection::emuBinary(&reader, tagToken, valueToken);
+        auto alpha = ESCCollection::emuBinary(reader, tagToken, valueToken);
 
         if (alpha.has_value() && alpha.value() == 0) {
             position = 1;
@@ -594,7 +594,7 @@ void processRemovedFields(FLEQueryInterface* queryImpl,
         auto valueToken =
             FLETwiceDerivedTokenGenerator::generateECCTwiceDerivedValueToken(plainTextField.ecc);
 
-        auto alpha = ECCCollection::emuBinary(&reader, tagToken, valueToken);
+        auto alpha = ECCCollection::emuBinary(reader, tagToken, valueToken);
 
         uint64_t index = 0;
         if (alpha.has_value() && alpha.value() == 0) {
diff --git a/src/mongo/db/fle_crud_test.cpp b/src/mongo/db/fle_crud_test.cpp
index 768f7a6d009..2823749a0b5 100644
--- a/src/mongo/db/fle_crud_test.cpp
+++ b/src/mongo/db/fle_crud_test.cpp
@@ -47,6 +47,7 @@
 #include "mongo/crypto/encryption_fields_gen.h"
 #include "mongo/crypto/fle_crypto.h"
 #include "mongo/crypto/fle_field_schema_gen.h"
+#include "mongo/crypto/fle_tags.h"
 #include "mongo/db/catalog/collection_options.h"
 #include "mongo/db/fle_crud.h"
 #include "mongo/db/matcher/schema/encrypt_schema_gen.h"
@@ -264,11 +265,10 @@ std::string fieldNameFromInt(uint64_t i) {
 }
 
 class FleCrudTest : public ServiceContextMongoDTest {
-private:
-    void setUp() final;
-    void tearDown() final;
-
 protected:
+    void setUp();
+    void tearDown();
+
     void createCollection(const NamespaceString& ns);
 
     void assertDocumentCounts(uint64_t edc, uint64_t esc, uint64_t ecc, uint64_t ecoc);
@@ -290,6 +290,10 @@ protected:
 
     void validateDocument(int id, boost::optional<BSONObj> doc);
 
+    ESCDerivedFromDataToken getTestESCDataToken(BSONObj obj);
+    ECCDerivedFromDataToken getTestECCDataToken(BSONObj obj);
+    EDCDerivedFromDataToken getTestEDCDataToken(BSONObj obj);
+
     ESCTwiceDerivedTagToken getTestESCToken(BSONElement value);
     ESCTwiceDerivedTagToken getTestESCToken(BSONObj obj);
     ESCTwiceDerivedTagToken getTestESCToken(StringData name, StringData value);
@@ -358,11 +362,37 @@ ConstDataRange toCDR(BSONElement element) {
     return ConstDataRange(element.value(), element.value() + element.valuesize());
 }
 
-ESCTwiceDerivedTagToken FleCrudTest::getTestESCToken(BSONElement element) {
+ESCDerivedFromDataToken FleCrudTest::getTestESCDataToken(BSONObj obj) {
+    auto element = obj.firstElement();
     auto c1token = FLELevel1TokenGenerator::generateCollectionsLevel1Token(
         _keyVault.getIndexKeyById(indexKeyId).key);
     auto escToken = FLECollectionTokenGenerator::generateESCToken(c1token);
+    return FLEDerivedFromDataTokenGenerator::generateESCDerivedFromDataToken(escToken,
+                                                                             toCDR(element));
+}
+
+ECCDerivedFromDataToken FleCrudTest::getTestECCDataToken(BSONObj obj) {
+    auto element = obj.firstElement();
+    auto c1token = FLELevel1TokenGenerator::generateCollectionsLevel1Token(
+        _keyVault.getIndexKeyById(indexKeyId).key);
+    auto eccToken = FLECollectionTokenGenerator::generateECCToken(c1token);
+    return FLEDerivedFromDataTokenGenerator::generateECCDerivedFromDataToken(eccToken,
+                                                                             toCDR(element));
+}
 
+EDCDerivedFromDataToken FleCrudTest::getTestEDCDataToken(BSONObj obj) {
+    auto element = obj.firstElement();
+    auto c1token = FLELevel1TokenGenerator::generateCollectionsLevel1Token(
+        _keyVault.getIndexKeyById(indexKeyId).key);
+    auto edcToken = FLECollectionTokenGenerator::generateEDCToken(c1token);
+    return FLEDerivedFromDataTokenGenerator::generateEDCDerivedFromDataToken(edcToken,
+                                                                             toCDR(element));
+}
+
+ESCTwiceDerivedTagToken FleCrudTest::getTestESCToken(BSONElement element) {
+    auto c1token = FLELevel1TokenGenerator::generateCollectionsLevel1Token(
+        _keyVault.getIndexKeyById(indexKeyId).key);
+    auto escToken = FLECollectionTokenGenerator::generateESCToken(c1token);
     auto escDataToken =
         FLEDerivedFromDataTokenGenerator::generateESCDerivedFromDataToken(escToken, toCDR(element));
     auto escContentionToken = FLEDerivedFromDataTokenAndContentionFactorTokenGenerator::
@@ -398,7 +428,6 @@ ECCDerivedFromDataTokenAndContentionFactorToken FleCrudTest::getTestECCToken(BSO
     auto c1token = FLELevel1TokenGenerator::generateCollectionsLevel1Token(
         _keyVault.getIndexKeyById(indexKeyId).key);
     auto eccToken = FLECollectionTokenGenerator::generateECCToken(c1token);
-
     auto eccDataToken =
         FLEDerivedFromDataTokenGenerator::generateECCDerivedFromDataToken(eccToken, toCDR(element));
     return FLEDerivedFromDataTokenAndContentionFactorTokenGenerator::
@@ -624,6 +653,43 @@ void FleCrudTest::doFindAndModify(write_ops::FindAndModifyCommandRequest& reques
     processFindAndModify(_queryImpl.get(), request);
 }
 
+class CollectionReader : public FLEStateCollectionReader {
+public:
+    CollectionReader(std::string&& coll, FLEQueryTestImpl& queryImpl)
+        : _coll(NamespaceString(coll)), _queryImpl(queryImpl) {}
+
+    uint64_t getDocumentCount() const override {
+        return _queryImpl.countDocuments(_coll);
+    }
+
+    BSONObj getById(PrfBlock block) const override {
+        auto doc = BSON("v" << BSONBinData(block.data(), block.size(), BinDataGeneral));
+        return _queryImpl.getById(_coll, doc.firstElement());
+    }
+
+private:
+    NamespaceString _coll;
+    FLEQueryTestImpl& _queryImpl;
+};
+
+class FleTagsTest : public FleCrudTest {
+protected:
+    void setUp() {
+        FleCrudTest::setUp();
+    }
+    void tearDown() {
+        FleCrudTest::tearDown();
+    }
+    std::vector<PrfBlock> readTags(BSONObj obj) {
+        auto s = getTestESCDataToken(obj);
+        auto c = getTestECCDataToken(obj);
+        auto d = getTestEDCDataToken(obj);
+        auto esc = CollectionReader("test.esc", *_queryImpl);
+        auto ecc = CollectionReader("test.ecc", *_queryImpl);
+        return mongo::fle::readTags(esc, ecc, s, c, d, 0);
+    }
+};
+
 // Insert one document
 TEST_F(FleCrudTest, InsertOne) {
     auto doc = BSON("encrypted"
@@ -1013,6 +1079,84 @@ TEST_F(FleCrudTest, FindAndModify_SetSafeContent) {
     ASSERT_THROWS_CODE(doFindAndModify(req), DBException, 6371507);
 }
 
+TEST_F(FleTagsTest, InsertOne) {
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    doSingleInsert(1, doc);
+
+    ASSERT_EQ(1, readTags(doc).size());
+}
+
+TEST_F(FleTagsTest, InsertTwoSame) {
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    doSingleInsert(1, doc);
+    doSingleInsert(2, doc);
+
+    ASSERT_EQ(2, readTags(doc).size());
+}
+
+TEST_F(FleTagsTest, InsertTwoDifferent) {
+    auto doc1 = BSON("encrypted"
+                     << "a");
+    auto doc2 = BSON("encrypted"
+                     << "b");
+
+    doSingleInsert(1, doc1);
+    doSingleInsert(2, doc2);
+
+    ASSERT_EQ(1, readTags(doc1).size());
+    ASSERT_EQ(1, readTags(doc2).size());
+}
+
+TEST_F(FleTagsTest, InsertAndDeleteOne) {
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    doSingleInsert(1, doc);
+    doSingleDelete(1);
+
+    ASSERT_EQ(0, readTags(doc).size());
+}
+
+TEST_F(FleTagsTest, InsertTwoSameAndDeleteOne) {
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    doSingleInsert(1, doc);
+    doSingleInsert(2, doc);
+    doSingleDelete(2);
+
+    ASSERT_EQ(1, readTags(doc).size());
+}
 
+TEST_F(FleTagsTest, InsertTwoDifferentAndDeleteOne) {
+    auto doc1 = BSON("encrypted"
+                     << "a");
+    auto doc2 = BSON("encrypted"
+                     << "b");
+
+    doSingleInsert(1, doc1);
+    doSingleInsert(2, doc2);
+    doSingleDelete(1);
+
+    ASSERT_EQ(0, readTags(doc1).size());
+    ASSERT_EQ(1, readTags(doc2).size());
+}
+
+TEST_F(FleTagsTest, InsertAndUpdate) {
+    auto doc1 = BSON("encrypted"
+                     << "a");
+    auto doc2 = BSON("encrypted"
+                     << "b");
+
+    doSingleInsert(1, doc1);
+    doSingleUpdate(1, doc2);
+
+    ASSERT_EQ(0, readTags(doc1).size());
+    ASSERT_EQ(1, readTags(doc2).size());
+}
 }  // namespace
 }  // namespace mongo