optimized UA login for higher throughput

backend/src/services/identity-ua/identity-ua-service.ts

@@ -84,69 +84,70 @@ export const identityUaServiceFactory = ({
 
     const LOCKOUT_KEY = `lockout:identity:${identityUa.identityId}:${IdentityAuthMethod.UNIVERSAL_AUTH}:${clientId}`;
 
-    let lock: Awaited<ReturnType<typeof keyStore.acquireLock>> | undefined;
+    const lockoutRaw = await keyStore.getItem(LOCKOUT_KEY);
-    if (identityUa.lockoutEnabled) {
+
-      try {
+    let lockout: LockoutObject | undefined;
-        lock = await keyStore.acquireLock([KeyStorePrefixes.IdentityLockoutLock(LOCKOUT_KEY)], 500, {
+    if (lockoutRaw) {
-          retryCount: 3,
+      lockout = JSON.parse(lockoutRaw) as LockoutObject;
-          retryDelay: 300,
+    }
-          retryJitter: 100
+
-        });
+    if (lockout && lockout.lockedOut) {
-      } catch (e) {
+      throw new UnauthorizedError({
-        logger.info(
+        message: "This identity auth method is temporarily locked, please try again later"
-          `identity login failed to acquire lock [identityId=${identityUa.identityId}] [authMethod=${IdentityAuthMethod.UNIVERSAL_AUTH}]`
+      });
-        );
+    }
-        throw new RateLimitError({ message: "Failed to acquire lock: rate limit exceeded" });
+
+    const identityMembershipOrg = await identityOrgMembershipDAL.findOne({ identityId: identityUa.identityId });
+    if (!identityMembershipOrg) {
+      throw new UnauthorizedError({
+        message: "Invalid credentials"
+      });
+    }
+
+    const clientSecretPrefix = clientSecret.slice(0, 4);
+    const clientSecretInfo = await identityUaClientSecretDAL.find({
+      identityUAId: identityUa.id,
+      isClientSecretRevoked: false,
+      clientSecretPrefix
+    });
+
+    let validClientSecretInfo: (typeof clientSecretInfo)[0] | null = null;
+    for await (const info of clientSecretInfo) {
+      const isMatch = await crypto.hashing().compareHash(clientSecret, info.clientSecretHash);
+
+      if (isMatch) {
+        validClientSecretInfo = info;
+        break;
       }
     }
 
-    try {
+    if (!validClientSecretInfo) {
-      const lockoutRaw = await keyStore.getItem(LOCKOUT_KEY);
+      if (identityUa.lockoutEnabled) {
+        let lock: Awaited<ReturnType<typeof keyStore.acquireLock>> | undefined;
+        try {
+          lock = await keyStore.acquireLock([KeyStorePrefixes.IdentityLockoutLock(LOCKOUT_KEY)], 300, {
+            retryCount: 3,
+            retryDelay: 300,
+            retryJitter: 100
+          });
 
-      let lockout: LockoutObject | undefined;
+          // Re-fetch the latest lockout data while holding the lock
-      if (lockoutRaw) {
+          const lockoutRawNew = await keyStore.getItem(LOCKOUT_KEY);
-        lockout = JSON.parse(lockoutRaw) as LockoutObject;
+          if (lockoutRawNew) {
-      }
+            lockout = JSON.parse(lockoutRawNew) as LockoutObject;
-
+          } else {
-      if (lockout && lockout.lockedOut) {
-        throw new UnauthorizedError({
-          message: "This identity auth method is temporarily locked, please try again later"
-        });
-      }
-
-      const identityMembershipOrg = await identityOrgMembershipDAL.findOne({ identityId: identityUa.identityId });
-      if (!identityMembershipOrg) {
-        throw new UnauthorizedError({
-          message: "Invalid credentials"
-        });
-      }
-
-      const clientSecretPrefix = clientSecret.slice(0, 4);
-      const clientSecretInfo = await identityUaClientSecretDAL.find({
-        identityUAId: identityUa.id,
-        isClientSecretRevoked: false,
-        clientSecretPrefix
-      });
-
-      let validClientSecretInfo: (typeof clientSecretInfo)[0] | null = null;
-      for await (const info of clientSecretInfo) {
-        const isMatch = await crypto.hashing().compareHash(clientSecret, info.clientSecretHash);
-
-        if (isMatch) {
-          validClientSecretInfo = info;
-          break;
-        }
-      }
-
-      if (!validClientSecretInfo) {
-        if (identityUa.lockoutEnabled) {
-          if (!lockout) {
             lockout = {
               lockedOut: false,
               failedAttempts: 0
             };
           }
 
+          if (lockout.lockedOut) {
+            throw new UnauthorizedError({
+              message: "This identity auth method is temporarily locked, please try again later"
+            });
+          }
+
           lockout.failedAttempts += 1;
           if (lockout.failedAttempts >= identityUa.lockoutThreshold) {
             lockout.lockedOut = true;
@@ -157,110 +158,121 @@ export const identityUaServiceFactory = ({
             lockout.lockedOut ? identityUa.lockoutDurationSeconds : identityUa.lockoutCounterResetSeconds,
             JSON.stringify(lockout)
           );
-        }
+        } catch (e) {
-
+          if (lock === undefined) {
-        throw new UnauthorizedError({ message: "Invalid credentials" });
+            logger.info(
-      } else if (lockout) {
+              `identity login failed to acquire lock [identityId=${identityUa.identityId}] [authMethod=${IdentityAuthMethod.UNIVERSAL_AUTH}]`
-        await keyStore.deleteItem(LOCKOUT_KEY);
+            );
-      }
+            throw new RateLimitError({ message: "Failed to acquire lock: rate limit exceeded" });
-
+          }
-      const { clientSecretTTL, clientSecretNumUses, clientSecretNumUsesLimit } = validClientSecretInfo;
+          throw e;
-      if (Number(clientSecretTTL) > 0) {
+        } finally {
-        const clientSecretCreated = new Date(validClientSecretInfo.createdAt);
+          if (lock) {
-        const ttlInMilliseconds = Number(clientSecretTTL) * 1000;
+            await lock.release();
-        const currentDate = new Date();
+          }
-        const expirationTime = new Date(clientSecretCreated.getTime() + ttlInMilliseconds);
-
-        if (currentDate > expirationTime) {
-          await identityUaClientSecretDAL.updateById(validClientSecretInfo.id, {
-            isClientSecretRevoked: true
-          });
-
-          throw new UnauthorizedError({
-            message: "Access denied due to expired client secret"
-          });
         }
       }
 
-      if (clientSecretNumUsesLimit > 0 && clientSecretNumUses >= clientSecretNumUsesLimit) {
+      throw new UnauthorizedError({ message: "Invalid credentials" });
-        // number of times client secret can be used for
+    } else if (lockout) {
-        // a login operation reached
+      // If credentials are valid, clear any existing lockout record
+      await keyStore.deleteItem(LOCKOUT_KEY);
+    }
+
+    const { clientSecretTTL, clientSecretNumUses, clientSecretNumUsesLimit } = validClientSecretInfo;
+    if (Number(clientSecretTTL) > 0) {
+      const clientSecretCreated = new Date(validClientSecretInfo.createdAt);
+      const ttlInMilliseconds = Number(clientSecretTTL) * 1000;
+      const currentDate = new Date();
+      const expirationTime = new Date(clientSecretCreated.getTime() + ttlInMilliseconds);
+
+      if (currentDate > expirationTime) {
         await identityUaClientSecretDAL.updateById(validClientSecretInfo.id, {
           isClientSecretRevoked: true
         });
+
         throw new UnauthorizedError({
-          message: "Access denied due to client secret usage limit reached"
+          message: "Access denied due to expired client secret"
         });
       }
+    }
 
-      const accessTokenTTLParams =
+    if (clientSecretNumUsesLimit > 0 && clientSecretNumUses >= clientSecretNumUsesLimit) {
-        Number(identityUa.accessTokenPeriod) === 0
+      // number of times client secret can be used for
-          ? {
+      // a login operation reached
-              accessTokenTTL: identityUa.accessTokenTTL,
+      await identityUaClientSecretDAL.updateById(validClientSecretInfo.id, {
-              accessTokenMaxTTL: identityUa.accessTokenMaxTTL
+        isClientSecretRevoked: true
-            }
-          : {
-              accessTokenTTL: identityUa.accessTokenPeriod,
-              // We set a very large Max TTL for periodic tokens to ensure that clients (even outdated ones) can always renew their token
-              // without them having to update their SDKs, CLIs, etc. This workaround sets it to 30 years to emulate "forever"
-              accessTokenMaxTTL: 1000000000
-            };
-
-      const identityAccessToken = await identityUaDAL.transaction(async (tx) => {
-        const uaClientSecretDoc = await identityUaClientSecretDAL.incrementUsage(validClientSecretInfo!.id, tx);
-        await identityOrgMembershipDAL.updateById(
-          identityMembershipOrg.id,
-          {
-            lastLoginAuthMethod: IdentityAuthMethod.UNIVERSAL_AUTH,
-            lastLoginTime: new Date()
-          },
-          tx
-        );
-        const newToken = await identityAccessTokenDAL.create(
-          {
-            identityId: identityUa.identityId,
-            isAccessTokenRevoked: false,
-            identityUAClientSecretId: uaClientSecretDoc.id,
-            accessTokenNumUses: 0,
-            accessTokenNumUsesLimit: identityUa.accessTokenNumUsesLimit,
-            accessTokenPeriod: identityUa.accessTokenPeriod,
-            authMethod: IdentityAuthMethod.UNIVERSAL_AUTH,
-            ...accessTokenTTLParams
-          },
-          tx
-        );
-
-        return newToken;
       });
+      throw new UnauthorizedError({
+        message: "Access denied due to client secret usage limit reached"
+      });
+    }
 
-      const appCfg = getConfig();
+    const accessTokenTTLParams =
-      const accessToken = crypto.jwt().sign(
+      Number(identityUa.accessTokenPeriod) === 0
+        ? {
+            accessTokenTTL: identityUa.accessTokenTTL,
+            accessTokenMaxTTL: identityUa.accessTokenMaxTTL
+          }
+        : {
+            accessTokenTTL: identityUa.accessTokenPeriod,
+            // We set a very large Max TTL for periodic tokens to ensure that clients (even outdated ones) can always renew their token
+            // without them having to update their SDKs, CLIs, etc. This workaround sets it to 30 years to emulate "forever"
+            accessTokenMaxTTL: 1000000000
+          };
+
+    const identityAccessToken = await identityUaDAL.transaction(async (tx) => {
+      const uaClientSecretDoc = await identityUaClientSecretDAL.incrementUsage(validClientSecretInfo!.id, tx);
+      await identityOrgMembershipDAL.updateById(
+        identityMembershipOrg.id,
+        {
+          lastLoginAuthMethod: IdentityAuthMethod.UNIVERSAL_AUTH,
+          lastLoginTime: new Date()
+        },
+        tx
+      );
+      const newToken = await identityAccessTokenDAL.create(
         {
           identityId: identityUa.identityId,
-          clientSecretId: validClientSecretInfo.id,
+          isAccessTokenRevoked: false,
-          identityAccessTokenId: identityAccessToken.id,
+          identityUAClientSecretId: uaClientSecretDoc.id,
-          authTokenType: AuthTokenType.IDENTITY_ACCESS_TOKEN
+          accessTokenNumUses: 0,
-        } as TIdentityAccessTokenJwtPayload,
+          accessTokenNumUsesLimit: identityUa.accessTokenNumUsesLimit,
-        appCfg.AUTH_SECRET,
+          accessTokenPeriod: identityUa.accessTokenPeriod,
-        // akhilmhdh: for non-expiry tokens you should not even set the value, including undefined. Even for undefined jsonwebtoken throws error
+          authMethod: IdentityAuthMethod.UNIVERSAL_AUTH,
-        Number(identityAccessToken.accessTokenTTL) === 0
+          ...accessTokenTTLParams
-          ? undefined
+        },
-          : {
+        tx
-              expiresIn: Number(identityAccessToken.accessTokenTTL)
-            }
       );
 
-      return {
+      return newToken;
-        accessToken,
+    });
-        identityUa,
+
-        validClientSecretInfo,
+    const appCfg = getConfig();
-        identityAccessToken,
+    const accessToken = crypto.jwt().sign(
-        identityMembershipOrg,
+      {
-        ...accessTokenTTLParams
+        identityId: identityUa.identityId,
-      };
+        clientSecretId: validClientSecretInfo.id,
-    } finally {
+        identityAccessTokenId: identityAccessToken.id,
-      if (lock) await lock.release();
+        authTokenType: AuthTokenType.IDENTITY_ACCESS_TOKEN
-    }
+      } as TIdentityAccessTokenJwtPayload,
+      appCfg.AUTH_SECRET,
+      // akhilmhdh: for non-expiry tokens you should not even set the value, including undefined. Even for undefined jsonwebtoken throws error
+      Number(identityAccessToken.accessTokenTTL) === 0
+        ? undefined
+        : {
+            expiresIn: Number(identityAccessToken.accessTokenTTL)
+          }
+    );
+
+    return {
+      accessToken,
+      identityUa,
+      validClientSecretInfo,
+      identityAccessToken,
+      identityMembershipOrg,
+      ...accessTokenTTLParams
+    };
   };
 
   const attachUniversalAuth = async ({