misc: made project key and data key creation concurrency safe

backend/src/services/kms/kms-service.ts

@@ -134,34 +134,32 @@ export const kmsServiceFactory = ({
 
           org = await orgDAL.findById(orgId);
         } else {
-          org = await orgDAL.findById(orgId);
+          const keyId = await orgDAL.transaction(async (tx) => {
-          if (!org.kmsDefaultKeyId) {
+            org = await orgDAL.findById(orgId, tx);
-            const keyId = await orgDAL.transaction(async (tx) => {
+            if (org.kmsDefaultKeyId) {
-              const key = await generateKmsKey({
+              return org.kmsDefaultKeyId;
-                isReserved: true,
+            }
-                orgId: org.id,
-                tx
-              });
 
-              await orgDAL.updateById(
+            const key = await generateKmsKey({
-                org.id,
+              isReserved: true,
-                {
+              orgId: org.id,
-                  kmsDefaultKeyId: key.id
+              tx
-                },
-                tx
-              );
-
-              await keyStore.setItemWithExpiry(
-                `${KeyStorePrefixes.WaitUntilReadyKmsOrgKeyCreation}${orgId}`,
-                10,
-                "true"
-              );
-
-              return key.id;
             });
 
-            return keyId;
+            await orgDAL.updateById(
-          }
+              org.id,
+              {
+                kmsDefaultKeyId: key.id
+              },
+              tx
+            );
+
+            await keyStore.setItemWithExpiry(`${KeyStorePrefixes.WaitUntilReadyKmsOrgKeyCreation}${orgId}`, 10, "true");
+
+            return key.id;
+          });
+
+          return keyId;
         }
       } finally {
         await lock?.release();
@@ -192,7 +190,6 @@ export const kmsServiceFactory = ({
         kmsId: kmsDoc.orgKms.id
       });
 
-      // fetch encryptedDataKey straight from kmsDoc by joining it in query :D
       const orgKmsDataKey = await orgKmsDecryptor({
         cipherTextBlob: kmsDoc.orgKms.encryptedDataKey
       });
@@ -326,38 +323,42 @@ export const kmsServiceFactory = ({
 
           org = await orgDAL.findById(orgId);
         } else {
-          org = await orgDAL.findById(orgId);
+          const orgDataKey = await orgDAL.transaction(async (tx) => {
-          if (!org.kmsEncryptedDataKey) {
+            org = await orgDAL.findById(orgId, tx);
-            const orgDataKey = await orgDAL.transaction(async (tx) => {
+            if (org.kmsEncryptedDataKey) {
-              const dataKey = randomSecureBytes();
+              return;
-              const kmsEncryptor = await encryptWithKmsKey(
+            }
-                {
-                  kmsId: kmsKeyId
-                },
-                tx
-              );
 
-              const { cipherTextBlob } = await kmsEncryptor({
+            const dataKey = randomSecureBytes();
-                plainText: dataKey
+            const kmsEncryptor = await encryptWithKmsKey(
-              });
+              {
+                kmsId: kmsKeyId
+              },
+              tx
+            );
 
-              await orgDAL.updateById(
+            const { cipherTextBlob } = await kmsEncryptor({
-                org.id,
+              plainText: dataKey
-                {
-                  kmsEncryptedDataKey: cipherTextBlob
-                },
-                tx
-              );
-
-              await keyStore.setItemWithExpiry(
-                `${KeyStorePrefixes.WaitUntilReadyKmsOrgDataKeyCreation}${orgId}`,
-                10,
-                "true"
-              );
-
-              return dataKey;
             });
 
+            await orgDAL.updateById(
+              org.id,
+              {
+                kmsEncryptedDataKey: cipherTextBlob
+              },
+              tx
+            );
+
+            await keyStore.setItemWithExpiry(
+              `${KeyStorePrefixes.WaitUntilReadyKmsOrgDataKeyCreation}${orgId}`,
+              10,
+              "true"
+            );
+
+            return dataKey;
+          });
+
+          if (orgDataKey) {
             return orgDataKey;
           }
         }
@@ -401,6 +402,11 @@ export const kmsServiceFactory = ({
           project = await projectDAL.findById(projectId);
         } else {
           const kmsKeyId = await projectDAL.transaction(async (tx) => {
+            project = await projectDAL.findById(projectId, tx);
+            if (project.kmsSecretManagerKeyId) {
+              return project.kmsSecretManagerKeyId;
+            }
+
             const key = await generateKmsKey({
               isReserved: true,
               orgId: project.orgId,
@@ -464,25 +470,40 @@ export const kmsServiceFactory = ({
 
           project = await projectDAL.findById(projectId);
         } else {
-          const dataKey = randomSecureBytes();
+          const projectDataKey = await projectDAL.transaction(async (tx) => {
-          const kmsEncryptor = await encryptWithKmsKey({
+            project = await projectDAL.findById(projectId, tx);
-            kmsId: kmsKeyId
+            if (project.kmsSecretManagerEncryptedDataKey) {
+              return;
+            }
+
+            const dataKey = randomSecureBytes();
+            const kmsEncryptor = await encryptWithKmsKey({
+              kmsId: kmsKeyId
+            });
+
+            const { cipherTextBlob } = await kmsEncryptor({
+              plainText: dataKey
+            });
+
+            await projectDAL.updateById(
+              projectId,
+              {
+                kmsSecretManagerEncryptedDataKey: cipherTextBlob
+              },
+              tx
+            );
+
+            await keyStore.setItemWithExpiry(
+              `${KeyStorePrefixes.WaitUntilReadyKmsProjectDataKeyCreation}${projectId}`,
+              10,
+              "true"
+            );
+            return dataKey;
           });
 
-          const { cipherTextBlob } = await kmsEncryptor({
+          if (projectDataKey) {
-            plainText: dataKey
+            return projectDataKey;
-          });
+          }
-
-          await projectDAL.updateById(projectId, {
-            kmsSecretManagerEncryptedDataKey: cipherTextBlob
-          });
-
-          await keyStore.setItemWithExpiry(
-            `${KeyStorePrefixes.WaitUntilReadyKmsProjectDataKeyCreation}${projectId}`,
-            10,
-            "true"
-          );
-          return dataKey;
         }
       } finally {
         await lock?.release();