misc: made org key and data key concurrency safe

backend/src/keystore/keystore.ts

@@ -10,7 +10,11 @@ export enum KeyStorePrefixes {
   KmsProjectDataKeyCreation = "kms-project-data-key-creation-lock",
   KmsProjectKeyCreation = "kms-project-key-creation-lock",
   WaitUntilReadyKmsProjectDataKeyCreation = "wait-until-ready-kms-project-data-key-creation-",
-  WaitUntilReadyKmsProjectKeyCreation = "wait-until-ready-kms-project-key-creation-"
+  WaitUntilReadyKmsProjectKeyCreation = "wait-until-ready-kms-project-key-creation-",
+  KmsOrgKeyCreation = "kms-org-key-creation-lock",
+  KmsOrgDataKeyCreation = "kms-org-data-key-creation-lock",
+  WaitUntilReadyKmsOrgKeyCreation = "wait-until-ready-kms-org-key-creation-",
+  WaitUntilReadyKmsOrgDataKeyCreation = "wait-until-ready-kms-org-data-key-creation-"
 }
 
 type TWaitTillReady = {

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

@@ -113,14 +113,30 @@ export const kmsServiceFactory = ({
   };
 
   const getOrgKmsKeyId = async (orgId: string) => {
-    const keyId = await orgDAL.transaction(async (tx) => {
+    let org = await orgDAL.findById(orgId);
-      const org = await orgDAL.findById(orgId, tx);
+
     if (!org) {
       throw new NotFoundError({ message: "Org not found" });
     }
 
     if (!org.kmsDefaultKeyId) {
-        // create default kms key for certificate service
+      const lock = await keyStore
+        .acquireLock([KeyStorePrefixes.KmsOrgKeyCreation, orgId], 3000, { retryCount: 3 })
+        .catch(() => null);
+
+      try {
+        if (!lock) {
+          await keyStore.waitTillReady({
+            key: `${KeyStorePrefixes.WaitUntilReadyKmsOrgKeyCreation}${orgId}`,
+            keyCheckCb: (val) => val === "true",
+            waitingCb: () => logger.info("KMS. Waiting for org key to be created")
+          });
+
+          org = await orgDAL.findById(orgId);
+        } else {
+          org = await orgDAL.findById(orgId);
+          if (!org.kmsDefaultKeyId) {
+            const keyId = await orgDAL.transaction(async (tx) => {
               const key = await generateKmsKey({
                 isReserved: true,
                 orgId: org.id,
@@ -135,13 +151,28 @@ export const kmsServiceFactory = ({
                 tx
               );
 
-        return key.id;
+              await keyStore.setItemWithExpiry(
-      }
+                `${KeyStorePrefixes.WaitUntilReadyKmsOrgKeyCreation}${orgId}`,
+                10,
+                "true"
+              );
 
-      return org.kmsDefaultKeyId;
+              return key.id;
             });
 
             return keyId;
+          }
+        }
+      } finally {
+        await lock?.release();
+      }
+    }
+
+    if (!org.kmsDefaultKeyId) {
+      throw new Error("Invalid organization KMS");
+    }
+
+    return org.kmsDefaultKeyId;
   };
 
   const decryptWithKmsKey = async ({ kmsId }: Omit<TDecryptWithKmsDTO, "cipherTextBlob">) => {
@@ -274,15 +305,30 @@ export const kmsServiceFactory = ({
 
   const getOrgKmsDataKey = async (orgId: string) => {
     const kmsKeyId = await getOrgKmsKeyId(orgId);
-    const orgKmsDataKey = await orgDAL.transaction(async (tx) => {
+    let org = await orgDAL.findById(orgId);
-      const org = await orgDAL.findById(orgId, tx);
 
     if (!org) {
       throw new NotFoundError({ message: "Org not found" });
     }
 
-      let encryptedDataKey = org.kmsEncryptedDataKey;
+    if (!org.kmsEncryptedDataKey) {
-      if (!encryptedDataKey) {
+      const lock = await keyStore
+        .acquireLock([KeyStorePrefixes.KmsOrgDataKeyCreation, orgId], 3000, { retryCount: 3 })
+        .catch(() => null);
+
+      try {
+        if (!lock) {
+          await keyStore.waitTillReady({
+            key: `${KeyStorePrefixes.WaitUntilReadyKmsOrgDataKeyCreation}${orgId}`,
+            keyCheckCb: (val) => val === "true",
+            waitingCb: () => logger.info("KMS. Waiting for org data key to be created")
+          });
+
+          org = await orgDAL.findById(orgId);
+        } else {
+          org = await orgDAL.findById(orgId);
+          if (!org.kmsEncryptedDataKey) {
+            const orgDataKey = await orgDAL.transaction(async (tx) => {
               const dataKey = randomSecureBytes();
               const kmsEncryptor = await encryptWithKmsKey(
                 {
@@ -295,16 +341,33 @@ export const kmsServiceFactory = ({
                 plainText: dataKey
               });
 
-        encryptedDataKey = cipherTextBlob;
               await orgDAL.updateById(
                 org.id,
                 {
-            kmsEncryptedDataKey: encryptedDataKey
+                  kmsEncryptedDataKey: cipherTextBlob
                 },
                 tx
               );
 
+              await keyStore.setItemWithExpiry(
+                `${KeyStorePrefixes.WaitUntilReadyKmsOrgDataKeyCreation}${orgId}`,
+                10,
+                "true"
+              );
+
               return dataKey;
+            });
+
+            return orgDataKey;
+          }
+        }
+      } finally {
+        await lock?.release();
+      }
+    }
+
+    if (!org.kmsEncryptedDataKey) {
+      throw new Error("Invalid organization KMS");
     }
 
     const kmsDecryptor = await decryptWithKmsKey({
@@ -312,11 +375,8 @@ export const kmsServiceFactory = ({
     });
 
     return kmsDecryptor({
-        cipherTextBlob: encryptedDataKey
+      cipherTextBlob: org.kmsEncryptedDataKey
     });
-    });
-
-    return orgKmsDataKey;
   };
 
   const getProjectSecretManagerKmsKeyId = async (projectId: string) => {