SERVER-42111 Enable auto-decrypt in community

3 files changed, 199 insertions, 2 deletions

diff --git a/jstests/client_encrypt/fle_auto_decrypt.js b/jstests/client_encrypt/fle_auto_decrypt.js
new file mode 100644
index 00000000000..73da503a84b
--- /dev/null
+++ b/jstests/client_encrypt/fle_auto_decrypt.js
@@ -0,0 +1,56 @@
+// Test to ensure that the client community shell auto decrypts an encrypted field
+// stored in the database if it has the correct credentials.
+
+load("jstests/client_encrypt/lib/mock_kms.js");
+load('jstests/ssl/libs/ssl_helpers.js');
+
+(function() {
+    "use strict";
+
+    const x509_options = {sslMode: "requireSSL", sslPEMKeyFile: SERVER_CERT, sslCAFile: CA_CERT};
+
+    const conn = MongoRunner.runMongod(x509_options);
+
+    let localKMS = {
+        key: BinData(
+            0,
+            "/i8ytmWQuCe1zt3bIuVa4taPGKhqasVp0/0yI4Iy0ixQPNmeDF1J5qPUbBYoueVUJHMqj350eRTwztAWXuBdSQ=="),
+    };
+
+    const clientSideFLEOptions = {
+        kmsProviders: {
+            local: localKMS,
+        },
+        keyVaultNamespace: "test.coll",
+        schemaMap: {}
+    };
+
+    const deterministicAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic";
+
+    const shell = Mongo(conn.host, clientSideFLEOptions);
+    const keyVault = shell.getKeyVault();
+
+    assert.writeOK(
+        keyVault.createKey("local", "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']));
+
+    const clientEncrypt = shell.getClientEncryption();
+    const keyId = keyVault.getKeyByAltName("mongoKey").toArray()[0]._id;
+
+    const encryptedStr = clientEncrypt.encrypt(keyId, "mongodb", deterministicAlgorithm);
+
+    // Insert encrypted string into database
+    const collection = conn.getDB("test").getCollection("collection");
+
+    for (var i = 0; i < 150; i++) {
+        assert.writeOK(collection.insert({string: encryptedStr, id: 1}));
+    }
+
+    // Ensure string is auto decrypted
+    const encryptedCollection = shell.getDB("test").getCollection("collection");
+    const result = encryptedCollection.find({id: 1}).toArray();
+    result.forEach(function(entry) {
+        assert.eq(entry.string, "mongodb");
+    });
+
+    MongoRunner.stopMongod(conn);
+}());
\ No newline at end of file
diff --git a/src/mongo/shell/encrypted_dbclient_base.cpp b/src/mongo/shell/encrypted_dbclient_base.cpp
index 1ee17b622b1..c8858f8a9f3 100644
--- a/src/mongo/shell/encrypted_dbclient_base.cpp
+++ b/src/mongo/shell/encrypted_dbclient_base.cpp
@@ -117,11 +117,127 @@ bool EncryptedDBClientBase::lazySupported() const {
     return _conn->lazySupported();
 }
 
+BSONObj EncryptedDBClientBase::encryptDecryptCommand(const BSONObj& object,
+                                                     bool encrypt,
+                                                     const StringData databaseName) {
+    std::stack<std::pair<BSONObjIterator, BSONObjBuilder>> frameStack;
+
+    // The encryptDecryptCommand frameStack requires a guard because  if encryptMarking or
+    // decrypt payload throw an exception, the stack's destructor will fire. Because a stack's
+    // variables are not guaranteed to be destroyed in any order, we need to add a guard
+    // to ensure the stack is destroyed in order.
+    const auto frameStackGuard = makeGuard([&] {
+        while (!frameStack.empty()) {
+            frameStack.pop();
+        }
+    });
+
+    frameStack.emplace(BSONObjIterator(object), BSONObjBuilder());
+
+    while (frameStack.size() > 1 || frameStack.top().first.more()) {
+        uassert(31096,
+                "Object too deep to be encrypted. Exceeded stack depth.",
+                frameStack.size() < BSONDepth::kDefaultMaxAllowableDepth);
+        auto & [ iterator, builder ] = frameStack.top();
+        if (iterator.more()) {
+            BSONElement elem = iterator.next();
+            if (elem.type() == BSONType::Object) {
+                frameStack.emplace(BSONObjIterator(elem.Obj()),
+                                   BSONObjBuilder(builder.subobjStart(elem.fieldNameStringData())));
+            } else if (elem.type() == BSONType::Array) {
+                frameStack.emplace(
+                    BSONObjIterator(elem.Obj()),
+                    BSONObjBuilder(builder.subarrayStart(elem.fieldNameStringData())));
+            } else if (elem.isBinData(BinDataType::Encrypt)) {
+                int len;
+                const char* data(elem.binData(len));
+                uassert(31178, "Invalid intentToEncrypt object from Query Analyzer", len >= 1);
+                if ((*data == kRandomEncryptionBit || *data == kDeterministicEncryptionBit) &&
+                    !encrypt) {
+                    ConstDataRange dataCursor(data, len);
+                    decryptPayload(dataCursor, &builder, elem.fieldNameStringData());
+                } else if (*data == kIntentToEncryptBit && encrypt) {
+                    BSONObj obj = BSONObj(data + 1);
+                    encryptMarking(obj, &builder, elem.fieldNameStringData());
+                } else {
+                    builder.append(elem);
+                }
+            } else {
+                builder.append(elem);
+            }
+        } else {
+            frameStack.pop();
+        }
+    }
+    invariant(frameStack.size() == 1);
+    frameStack.top().second.append("$db", databaseName);
+    return frameStack.top().second.obj();
+}
+
+void EncryptedDBClientBase::encryptMarking(const BSONObj& elem,
+                                           BSONObjBuilder* builder,
+                                           StringData elemName) {
+    MONGO_UNREACHABLE;
+}
+
+void EncryptedDBClientBase::decryptPayload(ConstDataRange data,
+                                           BSONObjBuilder* builder,
+                                           StringData elemName) {
+    uassert(ErrorCodes::BadValue, "Invalid decryption blob", data.length() > kAssociatedDataLength);
+    ConstDataRange uuidCdr = ConstDataRange(data.data() + 1, 16);
+    UUID uuid = UUID::fromCDR(uuidCdr);
+
+    auto key = getDataKey(uuid);
+    std::vector<uint8_t> out(data.length() - kAssociatedDataLength);
+    size_t outLen = out.size();
+
+    uassertStatusOK(
+        crypto::aeadDecrypt(*key,
+                            reinterpret_cast<const uint8_t*>(data.data() + kAssociatedDataLength),
+                            data.length() - kAssociatedDataLength,
+                            reinterpret_cast<const uint8_t*>(data.data()),
+                            kAssociatedDataLength,
+                            out.data(),
+                            &outLen));
+
+    // extract type byte
+    const uint8_t bsonType = static_cast<const uint8_t>(*(data.data() + 17));
+    BSONObj decryptedObj = validateBSONElement(ConstDataRange(out.data(), outLen), bsonType);
+    if (bsonType == BSONType::Object) {
+        builder->append(elemName, decryptedObj);
+    } else {
+        builder->appendAs(decryptedObj.firstElement(), elemName);
+    }
+}
+
+std::pair<rpc::UniqueReply, DBClientBase*> EncryptedDBClientBase::processResponse(
+    rpc::UniqueReply result, const StringData databaseName) {
+    auto rawReply = result->getCommandReply();
+    BSONObj decryptedDoc = encryptDecryptCommand(rawReply, false, databaseName);
+
+    rpc::OpMsgReplyBuilder replyBuilder;
+    replyBuilder.setCommandReply(StatusWith<BSONObj>(decryptedDoc));
+    auto msg = replyBuilder.done();
+
+    auto host = _conn->getServerAddress();
+    auto reply = _conn->parseCommandReplyMessage(host, msg);
+
+    return {std::move(reply), this};
+}
+
 std::pair<rpc::UniqueReply, DBClientBase*> EncryptedDBClientBase::runCommandWithTarget(
     OpMsgRequest request) {
-    return _conn->runCommandWithTarget(std::move(request));
-}
+    std::string commandName = request.getCommandName().toString();
+    std::string databaseName = request.getDatabase().toString();
 
+    if (std::find(kEncryptedCommands.begin(), kEncryptedCommands.end(), StringData(commandName)) ==
+        std::end(kEncryptedCommands)) {
+        return _conn->runCommandWithTarget(std::move(request));
+    }
+
+    auto result = _conn->runCommandWithTarget(std::move(request)).first;
+    return processResponse(std::move(result), databaseName);
+}
 
 /**
  *
diff --git a/src/mongo/shell/encrypted_dbclient_base.h b/src/mongo/shell/encrypted_dbclient_base.h
index f72b8f6cddc..1cd05dca2cf 100644
--- a/src/mongo/shell/encrypted_dbclient_base.h
+++ b/src/mongo/shell/encrypted_dbclient_base.h
@@ -64,6 +64,20 @@ constexpr uint8_t kIntentToEncryptBit = 0x00;
 constexpr uint8_t kDeterministicEncryptionBit = 0x01;
 constexpr uint8_t kRandomEncryptionBit = 0x02;
 
+static constexpr auto kExplain = "explain"_sd;
+
+constexpr std::array<StringData, 11> kEncryptedCommands = {"aggregate"_sd,
+                                                           "count"_sd,
+                                                           "delete"_sd,
+                                                           "distinct"_sd,
+                                                           kExplain,
+                                                           "find"_sd,
+                                                           "findandmodify"_sd,
+                                                           "findAndModify"_sd,
+                                                           "getMore"_sd,
+                                                           "insert"_sd,
+                                                           "update"_sd};
+
 class EncryptedDBClientBase : public DBClientBase, public mozjs::EncryptionCallbacks {
 public:
     EncryptedDBClientBase(std::unique_ptr<DBClientBase> conn,
@@ -126,6 +140,13 @@ public:
     bool isMongos() const final;
 
 protected:
+    std::pair<rpc::UniqueReply, DBClientBase*> processResponse(rpc::UniqueReply result,
+                                                               const StringData databaseName);
+
+    BSONObj encryptDecryptCommand(const BSONObj& object,
+                                  bool encrypt,
+                                  const StringData databaseName);
+
     JS::Value getCollection() const;
 
     BSONObj validateBSONElement(ConstDataRange out, uint8_t bsonType);
@@ -141,6 +162,10 @@ protected:
                                         int32_t algorithm);
 
 private:
+    virtual void encryptMarking(const BSONObj& elem, BSONObjBuilder* builder, StringData elemName);
+
+    void decryptPayload(ConstDataRange data, BSONObjBuilder* builder, StringData elemName);
+
     std::vector<uint8_t> getBinDataArg(mozjs::MozJSImplScope* scope,
                                        JSContext* cx,
                                        JS::CallArgs args,