SERVER-72913 Implement generate/decrypt functions for new ESC document formats

3 files changed, 157 insertions, 0 deletions

diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index 381f453babc..10e1673c1a8 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -139,6 +139,10 @@ constexpr uint64_t kESCNonNullId = 1;
 constexpr uint64_t KESCInsertRecordValue = 0;
 constexpr uint64_t kESCompactionRecordValue = std::numeric_limits<uint64_t>::max();
 
+constexpr uint64_t kESCAnchorId = 0;
+constexpr uint64_t kESCNullAnchorPosition = 0;
+constexpr uint64_t kESCNonNullAnchorValuePrefix = 0;
+
 constexpr auto kId = "_id";
 constexpr auto kValue = "value";
 constexpr auto kFieldName = "fieldName";
@@ -2278,6 +2282,19 @@ PrfBlock ESCCollection::generateId(ESCTwiceDerivedTagToken tagToken,
     }
 }
 
+PrfBlock ESCCollection::generateNonAnchorId(const ESCTwiceDerivedTagToken& tagToken,
+                                            uint64_t cpos) {
+    return prf(tagToken.data, cpos);
+}
+
+PrfBlock ESCCollection::generateAnchorId(const ESCTwiceDerivedTagToken& tagToken, uint64_t apos) {
+    return prf(tagToken.data, kESCAnchorId, apos);
+}
+
+PrfBlock ESCCollection::generateNullAnchorId(const ESCTwiceDerivedTagToken& tagToken) {
+    return ESCCollection::generateAnchorId(tagToken, kESCNullAnchorPosition);
+}
+
 BSONObj ESCCollection::generateNullDocument(ESCTwiceDerivedTagToken tagToken,
                                             ESCTwiceDerivedValueToken valueToken,
                                             uint64_t pos,
@@ -2336,6 +2353,44 @@ BSONObj ESCCollection::generateCompactionPlaceholderDocument(ESCTwiceDerivedTagT
     return builder.obj();
 }
 
+BSONObj ESCCollection::generateNonAnchorDocument(const ESCTwiceDerivedTagToken& tagToken,
+                                                 uint64_t cpos) {
+    auto block = ESCCollection::generateNonAnchorId(tagToken, cpos);
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    return builder.obj();
+}
+
+BSONObj ESCCollection::generateAnchorDocument(const ESCTwiceDerivedTagToken& tagToken,
+                                              const ESCTwiceDerivedValueToken& valueToken,
+                                              uint64_t apos,
+                                              uint64_t cpos) {
+    auto block = ESCCollection::generateAnchorId(tagToken, apos);
+
+    auto swCipherText = packAndEncrypt(std::tie(kESCNonNullAnchorValuePrefix, cpos), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+    return builder.obj();
+}
+
+BSONObj ESCCollection::generateNullAnchorDocument(const ESCTwiceDerivedTagToken& tagToken,
+                                                  const ESCTwiceDerivedValueToken& valueToken,
+                                                  uint64_t apos,
+                                                  uint64_t cpos) {
+    auto block = ESCCollection::generateNullAnchorId(tagToken);
+
+    auto swCipherText = packAndEncrypt(std::tie(apos, cpos), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+    return builder.obj();
+}
+
 StatusWith<ESCNullDocument> ESCCollection::decryptNullDocument(ESCTwiceDerivedValueToken valueToken,
                                                                BSONObj& doc) {
     return ESCCollection::decryptNullDocument(valueToken, std::move(doc));
@@ -2385,6 +2440,10 @@ StatusWith<ESCDocument> ESCCollection::decryptDocument(ESCTwiceDerivedValueToken
         std::get<0>(value) == kESCompactionRecordValue, std::get<0>(value), std::get<1>(value)};
 }
 
+StatusWith<ESCDocument> ESCCollection::decryptAnchorDocument(
+    const ESCTwiceDerivedValueToken& valueToken, BSONObj& doc) {
+    return ESCCollection::decryptDocument(valueToken, doc);
+}
 
 boost::optional<uint64_t> ESCCollection::emuBinary(const FLEStateCollectionReader& reader,
                                                    ESCTwiceDerivedTagToken tagToken,
diff --git a/src/mongo/crypto/fle_crypto.h b/src/mongo/crypto/fle_crypto.h
index bc8824551f0..c1e555af40d 100644
--- a/src/mongo/crypto/fle_crypto.h
+++ b/src/mongo/crypto/fle_crypto.h
@@ -284,6 +284,29 @@ public:
  *    uint64_t count_type;
  *    uint64_t count;
  * }
+ *
+ * ===== Protocol Version 2 =====
+ * Positional values:
+ *   cpos = position of non-anchor record in the range [1..UINT64_MAX]
+ *   apos = position of anchor record in the range [1..UINT64_MAX]
+ *
+ * Non-anchor record:
+ * {
+ *    _id : HMAC(ESCTwiceDerivedTagToken, cpos)
+ * }
+ *
+ * Non-null anchor record:
+ * {
+ *    _id : HMAC(ESCTwiceDerivedTagToken, (0 || apos))
+ *    value : Encrypt(ESCTwiceDerivedValueToken, (0 || cpos))
+ * }
+ *
+ * Null anchor record:
+ * {
+ *    _id : HMAC(ESCTwiceDerivedTagToken, (0 || 0))
+ *    value : Encrypt(ESCTwiceDerivedValueToken, (apos || cpos))
+ * }
+ *
  */
 
 
@@ -381,6 +404,50 @@ public:
     static boost::optional<uint64_t> emuBinary(const FLEStateCollectionReader& reader,
                                                ESCTwiceDerivedTagToken tagToken,
                                                ESCTwiceDerivedValueToken valueToken);
+
+    // ===== Protocol Version 2 =====
+    /**
+     * Generate the _id value for a non-anchor record
+     */
+    static PrfBlock generateNonAnchorId(const ESCTwiceDerivedTagToken& tagToken, uint64_t cpos);
+
+    /**
+     * Generate the _id value for an anchor record
+     */
+    static PrfBlock generateAnchorId(const ESCTwiceDerivedTagToken& tagToken, uint64_t apos);
+
+    /**
+     * Generate the _id value for a null anchor record
+     */
+    static PrfBlock generateNullAnchorId(const ESCTwiceDerivedTagToken& tagToken);
+
+    /**
+     * Generate a non-anchor ESC document for inserts.
+     */
+    static BSONObj generateNonAnchorDocument(const ESCTwiceDerivedTagToken& tagToken,
+                                             uint64_t cpos);
+
+    /**
+     * Generate an anchor ESC document for compacts.
+     */
+    static BSONObj generateAnchorDocument(const ESCTwiceDerivedTagToken& tagToken,
+                                          const ESCTwiceDerivedValueToken& valueToken,
+                                          uint64_t apos,
+                                          uint64_t cpos);
+
+    /**
+     * Generate a null anchor ESC document for cleanups.
+     */
+    static BSONObj generateNullAnchorDocument(const ESCTwiceDerivedTagToken& tagToken,
+                                              const ESCTwiceDerivedValueToken& valueToken,
+                                              uint64_t apos,
+                                              uint64_t cpos);
+
+    /**
+     * Decrypts an anchor document (either null or non-null).
+     */
+    static StatusWith<ESCDocument> decryptAnchorDocument(
+        const ESCTwiceDerivedValueToken& valueToken, BSONObj& doc);
 };
 
 
diff --git a/src/mongo/crypto/fle_crypto_test.cpp b/src/mongo/crypto/fle_crypto_test.cpp
index 69b08eb831c..792b1f2143a 100644
--- a/src/mongo/crypto/fle_crypto_test.cpp
+++ b/src/mongo/crypto/fle_crypto_test.cpp
@@ -428,6 +428,11 @@ TEST(FLETokens, TestVectorESCCollectionDecryptDocument) {
     ASSERT_EQ(swDoc.getValue().compactionPlaceholder, false);
     ASSERT_EQ(swDoc.getValue().position, 0);
     ASSERT_EQ(swDoc.getValue().count, 123456789);
+
+    auto swAnchorDoc = ESCCollection::decryptAnchorDocument(escTwiceValue, doc);
+    ASSERT_OK(swDoc.getStatus());
+    ASSERT_EQ(swDoc.getValue().position, 0);
+    ASSERT_EQ(swDoc.getValue().count, 123456789);
 }
 
 TEST(FLETokens, TestVectorECCCollectionDecryptDocument) {
@@ -509,6 +514,32 @@ TEST(FLE_ESC, RoundTrip) {
         ASSERT_EQ(swDoc.getValue().position, std::numeric_limits<uint64_t>::max());
         ASSERT_EQ(swDoc.getValue().count, 456789);
     }
+
+    {
+        // Non-anchor documents don't work with decryptAnchorDocument()
+        BSONObj doc = ESCCollection::generateNonAnchorDocument(escTwiceTag, 123);
+        auto swDoc = ESCCollection::decryptAnchorDocument(escTwiceValue, doc);
+        ASSERT_NOT_OK(swDoc.getStatus());
+        ASSERT_EQ(ErrorCodes::Error::NoSuchKey, swDoc.getStatus().code());
+    }
+
+    {
+        BSONObj doc =
+            ESCCollection::generateAnchorDocument(escTwiceTag, escTwiceValue, 123, 456789);
+        auto swDoc = ESCCollection::decryptAnchorDocument(escTwiceValue, doc);
+        ASSERT_OK(swDoc.getStatus());
+        ASSERT_EQ(swDoc.getValue().position, 0);
+        ASSERT_EQ(swDoc.getValue().count, 456789);
+    }
+
+    {
+        BSONObj doc =
+            ESCCollection::generateNullAnchorDocument(escTwiceTag, escTwiceValue, 123, 456789);
+        auto swDoc = ESCCollection::decryptAnchorDocument(escTwiceValue, doc);
+        ASSERT_OK(swDoc.getStatus());
+        ASSERT_EQ(swDoc.getValue().position, 123);
+        ASSERT_EQ(swDoc.getValue().count, 456789);
+    }
 }
 
 TEST(FLE_ECC, RoundTrip) {