2 files changed, 26 insertions, 1 deletions
diff --git a/src/mongo/shell/encrypted_dbclient_base.cpp b/src/mongo/shell/encrypted_dbclient_base.cpp
index 844c44fea36..a2e2ed037b0 100644
--- a/src/mongo/shell/encrypted_dbclient_base.cpp
+++ b/src/mongo/shell/encrypted_dbclient_base.cpp
@@ -693,7 +693,7 @@ DBClientBase* EncryptedDBClientBase::getRawConnection() {
     return _conn.get();
 }
 
-SecureVector<uint8_t> EncryptedDBClientBase::getKeyMaterialFromDisk(const UUID& uuid) {
+BSONObj EncryptedDBClientBase::getEncryptedKey(const UUID& uuid) {
     NamespaceString fullNameNS = getCollectionNS();
     FindCommandRequest findCmd{fullNameNS};
     findCmd.setFilter(BSON("_id" << uuid));
@@ -720,6 +720,15 @@ SecureVector<uint8_t> EncryptedDBClientBase::getKeyMaterialFromDisk(const UUID&
     auto dataKey = keyStoreRecord.getKeyMaterial();
     uassert(ErrorCodes::BadValue, "Invalid data key.", dataKey.length() != 0);
 
+    return keyStoreRecord.toBSON();
+}
+
+SecureVector<uint8_t> EncryptedDBClientBase::getKeyMaterialFromDisk(const UUID& uuid) {
+    auto rawKey = getEncryptedKey(uuid);
+    auto keyStoreRecord = KeyStoreRecord::parse(IDLParserContext("root"), rawKey);
+
+    auto dataKey = keyStoreRecord.getKeyMaterial();
+
     std::unique_ptr<KMSService> kmsService = KMSServiceController::createFromDisk(
         _encryptionOptions.getKmsProviders().toBSON(), keyStoreRecord.getMasterKey());
     SecureVector<uint8_t> decryptedKey =
@@ -742,6 +751,18 @@ KeyMaterial EncryptedDBClientBase::getKey(const UUID& uuid) {
     return km;
 }
 
+SymmetricKey& EncryptedDBClientBase::getKMSLocalKey() {
+    if (!_localKey.has_value()) {
+        std::unique_ptr<KMSService> kmsService =
+            KMSServiceController::createFromDisk(_encryptionOptions.getKmsProviders().toBSON(),
+                                                 BSON("provider"
+                                                      << "local"));
+        _localKey = std::move(kmsService->getMasterKey());
+    }
+
+    return _localKey.get();
+}
+
 #ifdef MONGO_CONFIG_SSL
 const SSLConfiguration* EncryptedDBClientBase::getSSLConfiguration() {
     return _conn->getSSLConfiguration();
diff --git a/src/mongo/shell/encrypted_dbclient_base.h b/src/mongo/shell/encrypted_dbclient_base.h
index 4f00b5b0e4b..9ef3420a55a 100644
--- a/src/mongo/shell/encrypted_dbclient_base.h
+++ b/src/mongo/shell/encrypted_dbclient_base.h
@@ -152,6 +152,9 @@ public:
 #endif
 
     KeyMaterial getKey(const UUID& uuid) final;
+    BSONObj getEncryptedKey(const UUID& uuid) final;
+
+    SymmetricKey& getKMSLocalKey() final;
 
 protected:
     BSONObj _decryptResponsePayload(BSONObj& reply, StringData databaseName, bool isFLE2);
@@ -250,6 +253,7 @@ private:
         kEncryptedDBCacheSize};
     JS::Heap<JS::Value> _collection;
     JSContext* _cx;
+    boost::optional<SymmetricKey> _localKey;
 };
 
 using ImplicitEncryptedDBClientCallback =