SERVER-63467 Create a shell helper that can be used to call compact encryption data

9 files changed, 130 insertions, 1 deletions

diff --git a/jstests/fle2/compact_collection.js b/jstests/fle2/compact_collection.js
new file mode 100644
index 00000000000..0969fc216d2
--- /dev/null
+++ b/jstests/fle2/compact_collection.js
@@ -0,0 +1,43 @@
+// Verify compact collection capability in client side
+
+/**
+ * @tags: [
+ *  featureFlagFLE2,
+ * ]
+ */
+load("jstests/fle2/libs/encrypted_client_util.js");
+
+(function() {
+'use strict';
+
+if (!isFLE2Enabled()) {
+    return;
+}
+
+const dbName = 'compact_collection_db';
+const dbTest = db.getSiblingDB(dbName);
+dbTest.dropDatabase();
+
+const client = new EncryptedClient(db.getMongo(), dbName);
+const edb = client.getDB();
+
+const sampleEncryptedFields = {
+    fields: [
+        {path: "firstName", bsonType: "string", queries: {"queryType": "equality"}},
+        {path: "a.b.c", bsonType: "int", queries: {"queryType": "equality"}},
+    ]
+};
+
+assert.commandWorked(
+    client.createEncryptionCollection("encrypted", {encryptedFields: sampleEncryptedFields}));
+assert.commandWorked(edb.createCollection("unencrypted"));
+
+assert.commandFailedWithCode(edb.unencrypted.compact(), ErrorCodes.BadValue);
+
+const res = edb.encrypted.compact();
+assert.commandWorked(res);
+assert(res.hasOwnProperty("stats"));
+assert(res.stats.hasOwnProperty("esc"));
+assert(res.stats.hasOwnProperty("ecc"));
+assert(res.stats.hasOwnProperty("ecoc"));
+}());
diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index 7b4fdfa851d..33813becaf1 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -1232,6 +1232,20 @@ BSONObj FLEClientCrypto::generateInsertOrUpdateFromPlaceholders(const BSONObj& o
     return ret;
 }
 
+BSONObj FLEClientCrypto::generateCompactionTokens(const EncryptedFieldConfig& cfg,
+                                                  FLEKeyVault* keyVault) {
+    BSONObjBuilder tokensBuilder;
+    auto& fields = cfg.getFields();
+    for (const auto& field : fields) {
+        auto indexKey = keyVault->getIndexKeyById(field.getKeyId());
+        auto collToken = FLELevel1TokenGenerator::generateCollectionsLevel1Token(indexKey.key);
+        auto ecocToken = FLECollectionTokenGenerator::generateECOCToken(collToken);
+        auto tokenCdr = ecocToken.toCDR();
+        tokensBuilder.appendBinData(
+            field.getPath(), tokenCdr.length(), BinDataType::BinDataGeneral, tokenCdr.data());
+    }
+    return tokensBuilder.obj();
+}
 
 std::pair<BSONType, std::vector<uint8_t>> FLEClientCrypto::decrypt(BSONElement element,
                                                                    FLEKeyVault* keyVault) {
diff --git a/src/mongo/crypto/fle_crypto.h b/src/mongo/crypto/fle_crypto.h
index 2460e0b1436..de2b9e93c13 100644
--- a/src/mongo/crypto/fle_crypto.h
+++ b/src/mongo/crypto/fle_crypto.h
@@ -733,6 +733,17 @@ public:
     static BSONObj generateInsertOrUpdateFromPlaceholders(const BSONObj& obj,
                                                           FLEKeyVault* keyVault);
 
+
+    /**
+     * For every encrypted field path in the EncryptedFieldConfig, this generates
+     * a compaction token derived from the field's index key, which is retrieved from
+     * the supplied FLEKeyVault using the field's key ID.
+     *
+     * Returns a BSON object mapping the encrypted field path to its compaction token,
+     * which is a general BinData value.
+     */
+    static BSONObj generateCompactionTokens(const EncryptedFieldConfig& cfg, FLEKeyVault* keyVault);
+
     /**
      * Decrypts a document. Only supports FLE2.
      */
diff --git a/src/mongo/scripting/mozjs/mongo.cpp b/src/mongo/scripting/mozjs/mongo.cpp
index b06f404c383..00818afb115 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -62,6 +62,7 @@ namespace mozjs {
 const JSFunctionSpec MongoBase::methods[] = {
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(auth, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(close, MongoExternalInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(compact, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(cursorHandleFromId, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(find, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(generateDataKey, MongoExternalInfo),
@@ -407,6 +408,12 @@ void MongoBase::Functions::decrypt::call(JSContext* cx, JS::CallArgs args) {
     ptr->decrypt(scope, cx, args);
 }
 
+void MongoBase::Functions::compact::call(JSContext* cx, JS::CallArgs args) {
+    auto conn = getConnection(args);
+    auto ptr = getEncryptionCallbacks(conn);
+    ptr->compact(cx, args);
+}
+
 void MongoBase::Functions::logout::call(JSContext* cx, JS::CallArgs args) {
     if (args.length() != 1)
         uasserted(ErrorCodes::BadValue, "logout needs 1 arg");
diff --git a/src/mongo/scripting/mozjs/mongo.h b/src/mongo/scripting/mozjs/mongo.h
index 24dadb9c232..b79ba3f061a 100644
--- a/src/mongo/scripting/mozjs/mongo.h
+++ b/src/mongo/scripting/mozjs/mongo.h
@@ -55,6 +55,7 @@ struct MongoBase : public BaseInfo {
     struct Functions {
         MONGO_DECLARE_JS_FUNCTION(auth);
         MONGO_DECLARE_JS_FUNCTION(close);
+        MONGO_DECLARE_JS_FUNCTION(compact);
         MONGO_DECLARE_JS_FUNCTION(cursorHandleFromId);
         MONGO_DECLARE_JS_FUNCTION(find);
         MONGO_DECLARE_JS_FUNCTION(generateDataKey);
@@ -77,7 +78,7 @@ struct MongoBase : public BaseInfo {
         MONGO_DECLARE_JS_FUNCTION(_startSession);
     };
 
-    static const JSFunctionSpec methods[20];
+    static const JSFunctionSpec methods[21];
 
     static const char* const className;
     static const unsigned classFlags = JSCLASS_HAS_PRIVATE;
@@ -104,6 +105,7 @@ public:
     virtual void getDataKeyCollection(JSContext* cx, JS::CallArgs args) = 0;
     virtual void encrypt(MozJSImplScope* scope, JSContext* cx, JS::CallArgs args) = 0;
     virtual void decrypt(MozJSImplScope* scope, JSContext* cx, JS::CallArgs args) = 0;
+    virtual void compact(JSContext* cx, JS::CallArgs args) = 0;
     virtual void trace(JSTracer* trc) = 0;
 };
 
diff --git a/src/mongo/shell/SConscript b/src/mongo/shell/SConscript
index 9bdfdad9c17..aaa9e39f54a 100644
--- a/src/mongo/shell/SConscript
+++ b/src/mongo/shell/SConscript
@@ -205,6 +205,7 @@ if get_option('ssl') == 'on':
             LIBDEPS_PRIVATE=[
                 '$BUILD_DIR/mongo/client/clientdriver_minimal',
                 '$BUILD_DIR/mongo/crypto/aead_encryption',
+                '$BUILD_DIR/mongo/crypto/encrypted_field_config',
                 '$BUILD_DIR/mongo/crypto/fle_crypto',
                 '$BUILD_DIR/mongo/crypto/fle_fields',
                 '$BUILD_DIR/mongo/crypto/symmetric_crypto',
diff --git a/src/mongo/shell/collection.js b/src/mongo/shell/collection.js
index 13e83d59855..85c6b92c09d 100644
--- a/src/mongo/shell/collection.js
+++ b/src/mongo/shell/collection.js
@@ -13,6 +13,10 @@ if ((typeof DBCollection) == "undefined") {
     };
 }
 
+DBCollection.prototype.compact = function() {
+    return this._db.getMongo().compact(this._fullName);
+};
+
 DBCollection.prototype.verify = function() {
     assert(this._fullName, "no fullName");
     assert(this._shortName, "no shortName");
diff --git a/src/mongo/shell/encrypted_dbclient_base.cpp b/src/mongo/shell/encrypted_dbclient_base.cpp
index c38a8df258b..137d0c482ab 100644
--- a/src/mongo/shell/encrypted_dbclient_base.cpp
+++ b/src/mongo/shell/encrypted_dbclient_base.cpp
@@ -515,6 +515,48 @@ void EncryptedDBClientBase::decrypt(mozjs::MozJSImplScope* scope,
     }
 }
 
+boost::optional<EncryptedFieldConfig> EncryptedDBClientBase::getEncryptedFieldConfig(
+    const NamespaceString& nss) {
+    auto collsList = _conn->getCollectionInfos(nss.db().toString(), BSON("name" << nss.coll()));
+    uassert(ErrorCodes::BadValue,
+            str::stream() << "Namespace not found: " << nss.toString(),
+            !collsList.empty());
+    auto info = collsList.front();
+    auto opts = info.getField("options");
+    if (opts.eoo() || !opts.isABSONObj()) {
+        return boost::none;
+    }
+    auto efc = opts.Obj().getField("encryptedFields");
+    if (efc.eoo() || !efc.isABSONObj()) {
+        return boost::none;
+    }
+    return EncryptedFieldConfig::parse(IDLParserErrorContext("encryptedFields"), efc.Obj());
+}
+
+void EncryptedDBClientBase::compact(JSContext* cx, JS::CallArgs args) {
+    if (args.length() != 1) {
+        uasserted(ErrorCodes::BadValue, "compact requires 1 arg");
+    }
+    if (!args.get(0).isString()) {
+        uasserted(ErrorCodes::BadValue, "1st param to compact has to be a string");
+    }
+    std::string fullName = mozjs::ValueWriter(cx, args.get(0)).toString();
+    NamespaceString nss(fullName);
+    uassert(
+        ErrorCodes::BadValue, str::stream() << "Invalid namespace: " << fullName, nss.isValid());
+
+    auto efc = getEncryptedFieldConfig(nss);
+    BSONObjBuilder builder;
+    builder.append("compactStructuredEncryptionData", nss.coll());
+    builder.append("compactionTokens",
+                   efc ? FLEClientCrypto::generateCompactionTokens(*efc, this) : BSONObj());
+
+    BSONObj reply;
+    runCommand(nss.db().toString(), builder.obj(), reply, 0);
+    reply = reply.getOwned();
+    mozjs::ValueReader(cx, args.rval()).fromBSON(reply, nullptr, false);
+}
+
 void EncryptedDBClientBase::trace(JSTracer* trc) {
     JS::TraceEdge(trc, &_collection, "collection object");
 }
diff --git a/src/mongo/shell/encrypted_dbclient_base.h b/src/mongo/shell/encrypted_dbclient_base.h
index 838d110a477..8a6bab60f59 100644
--- a/src/mongo/shell/encrypted_dbclient_base.h
+++ b/src/mongo/shell/encrypted_dbclient_base.h
@@ -118,6 +118,9 @@ public:
     using EncryptionCallbacks::decrypt;
     void decrypt(mozjs::MozJSImplScope* scope, JSContext* cx, JS::CallArgs args) final;
 
+    using EncryptionCallbacks::compact;
+    void compact(JSContext* cx, JS::CallArgs args) final;
+
     using EncryptionCallbacks::trace;
     void trace(JSTracer* trc) final;
 
@@ -201,6 +204,8 @@ private:
     std::shared_ptr<SymmetricKey> getDataKeyFromDisk(const UUID& uuid);
     SecureVector<uint8_t> getKeyMaterialFromDisk(const UUID& uuid);
 
+    boost::optional<EncryptedFieldConfig> getEncryptedFieldConfig(const NamespaceString& nss);
+
 protected:
     std::unique_ptr<DBClientBase> _conn;
     ClientSideFLEOptions _encryptionOptions;