SERVER-43005 : getKeyVault.createKey() should return the created data key

(cherry picked from commit a605bfb8b3ecdff282508263b119f1fcfd13f27e)

7 files changed, 17 insertions, 22 deletions

diff --git a/jstests/client_encrypt/fle_auto_decrypt.js b/jstests/client_encrypt/fle_auto_decrypt.js
index 182c8af3ac4..86d3bd9bf91 100644
--- a/jstests/client_encrypt/fle_auto_decrypt.js
+++ b/jstests/client_encrypt/fle_auto_decrypt.js
@@ -34,11 +34,9 @@ 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 keyId = 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
diff --git a/jstests/client_encrypt/fle_aws_faults.js b/jstests/client_encrypt/fle_aws_faults.js
index 1d9a621b42a..e0f21d9f99b 100644
--- a/jstests/client_encrypt/fle_aws_faults.js
+++ b/jstests/client_encrypt/fle_aws_faults.js
@@ -84,9 +84,8 @@ function testBadDecryptResult(fault) {
 
     runKMS(mock_kms, (shell) => {
         const keyVault = shell.getKeyVault();
-        assert.writeOK(
-            keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ["mongoKey"]));
-        const keyId = keyVault.getKeys("mongoKey").toArray()[0]._id;
+        const keyId =
+            keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ["mongoKey"]);
         const str = "mongo";
         assert.throws(() => {
             const encStr = shell.getClientEncryption().encrypt(keyId, str, randomAlgorithm);
@@ -102,8 +101,7 @@ function testBadDecryptKeyResult(fault) {
     runKMS(mock_kms, (shell, cleanCacheShell) => {
         const keyVault = shell.getKeyVault();
 
-        assert.writeOK(
-            keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ["mongoKey"]));
+        keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ["mongoKey"]);
         const keyId = keyVault.getKeys("mongoKey").toArray()[0]._id;
         const str = "mongo";
         const encStr = shell.getClientEncryption().encrypt(keyId, str, randomAlgorithm);
@@ -123,8 +121,7 @@ function testBadDecryptError() {
 
     runKMS(mock_kms, (shell) => {
         const keyVault = shell.getKeyVault();
-        assert.writeOK(
-            keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ["mongoKey"]));
+        keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ["mongoKey"]);
         const keyId = keyVault.getKeys("mongoKey").toArray()[0]._id;
         const str = "mongo";
         let error = assert.throws(() => {
diff --git a/jstests/client_encrypt/fle_encrypt_decrypt_shell.js b/jstests/client_encrypt/fle_encrypt_decrypt_shell.js
index f67bc72dccc..ae323b0fd84 100644
--- a/jstests/client_encrypt/fle_encrypt_decrypt_shell.js
+++ b/jstests/client_encrypt/fle_encrypt_decrypt_shell.js
@@ -87,9 +87,8 @@ for (const kmsType of kmsTypes) {
     for (const encryptionAlgorithm of encryptionAlgorithms) {
         collection.drop();
 
-        assert.writeOK(
-            keyVault.createKey(kmsType, "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']));
-        const keyId = keyVault.getKeyByAltName("mongoKey").toArray()[0]._id;
+        const keyId =
+            keyVault.createKey(kmsType, "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']);
 
         let pass;
         let fail;
diff --git a/jstests/client_encrypt/fle_key_faults.js b/jstests/client_encrypt/fle_key_faults.js
index 41b2505c8ba..8f8133cbaf1 100644
--- a/jstests/client_encrypt/fle_key_faults.js
+++ b/jstests/client_encrypt/fle_key_faults.js
@@ -48,8 +48,7 @@ function testFault(kmsType, func) {
     const shell = Mongo(conn.host, clientSideFLEOptions);
     const keyVault = shell.getKeyVault();
 
-    assert.writeOK(
-        keyVault.createKey(kmsType, "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']));
+    keyVault.createKey(kmsType, "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']);
     const keyId = keyVault.getKeyByAltName("mongoKey").toArray()[0]._id;
 
     func(keyId, shell);
diff --git a/jstests/client_encrypt/fle_keys.js b/jstests/client_encrypt/fle_keys.js
index 646b95141ad..d7039e2dfaf 100644
--- a/jstests/client_encrypt/fle_keys.js
+++ b/jstests/client_encrypt/fle_keys.js
@@ -39,7 +39,7 @@ const conn_str = "mongodb://" + conn.host + "/?ssl=true";
 const shell = Mongo(conn_str, clientSideFLEOptions);
 const keyVault = shell.getKeyVault();
 
-var key = keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']);
+keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']);
 assert.eq(1, keyVault.getKeys().itcount());
 
 var result = keyVault.createKey("aws", "arn:aws:kms:us-east-4:fake:fake:fake", {});
@@ -68,9 +68,9 @@ result = keyVault.deleteKey(keyId);
 assert.eq(0, keyVault.getKey(keyId).itcount());
 assert.eq(0, keyVault.getKeys().itcount());
 
-assert.writeOK(keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake1"));
-assert.writeOK(keyVault.createKey("aws", "arn:aws:kms:us-east-2:fake:fake:fake2"));
-assert.writeOK(keyVault.createKey("aws", "arn:aws:kms:us-east-3:fake:fake:fake3"));
+keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake1");
+keyVault.createKey("aws", "arn:aws:kms:us-east-2:fake:fake:fake2");
+keyVault.createKey("aws", "arn:aws:kms:us-east-3:fake:fake:fake3");
 
 assert.eq(3, keyVault.getKeys().itcount());
 
diff --git a/jstests/client_encrypt/lib/fle_command_line_explicit_encryption.js b/jstests/client_encrypt/lib/fle_command_line_explicit_encryption.js
index da83d69c87b..e4e2d159e80 100644
--- a/jstests/client_encrypt/lib/fle_command_line_explicit_encryption.js
+++ b/jstests/client_encrypt/lib/fle_command_line_explicit_encryption.js
@@ -51,7 +51,7 @@ const failTestCases = [null, undefined, MinKey(), MaxKey(), DBRef("test", "test"
 for (const encryptionAlgorithm of encryptionAlgorithms) {
     collection.drop();
 
-    assert.writeOK(keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']));
+    keyVault.createKey("aws", "arn:aws:kms:us-east-1:fake:fake:fake", ['mongoKey']);
     const keyId = keyVault.getKeyByAltName("mongoKey").toArray()[0]._id;
 
     let pass;
diff --git a/src/mongo/shell/keyvault.js b/src/mongo/shell/keyvault.js
index 4247a14c798..959a4e8c7f4 100644
--- a/src/mongo/shell/keyvault.js
+++ b/src/mongo/shell/keyvault.js
@@ -28,9 +28,10 @@ class KeyVault {
         var masterKey = masterKeyAndMaterial.masterKey;
 
         var current = ISODate();
+        var uuid = UUID();
 
         var doc = {
-            "_id": UUID(),
+            "_id": uuid,
             "keyMaterial": masterKeyAndMaterial.keyMaterial,
             "creationDate": current,
             "updateDate": current,
@@ -54,7 +55,8 @@ class KeyVault {
             doc.keyAltNames = keyAltNames;
         }
 
-        return this.keyColl.insert(doc);
+        this.keyColl.insert(doc);
+        return uuid;
     }
 
     getKey(keyId) {