Merge pull request #2317 from akhilmhdh/feat/debounce-last-used

feat: added identity and service token postgres update debounced

backend/src/keystore/keystore.ts

@@ -19,11 +19,15 @@ export const KeyStorePrefixes = {
   SyncSecretIntegrationLock: (projectId: string, environmentSlug: string, secretPath: string) =>
     `sync-integration-mutex-${projectId}-${environmentSlug}-${secretPath}` as const,
   SyncSecretIntegrationLastRunTimestamp: (projectId: string, environmentSlug: string, secretPath: string) =>
-    `sync-integration-last-run-${projectId}-${environmentSlug}-${secretPath}` as const
+    `sync-integration-last-run-${projectId}-${environmentSlug}-${secretPath}` as const,
+  IdentityAccessTokenStatusUpdate: (identityAccessTokenId: string) =>
+    `identity-access-token-status:${identityAccessTokenId}`,
+  ServiceTokenStatusUpdate: (serviceTokenId: string) => `service-token-status:${serviceTokenId}`
 };
 
 export const KeyStoreTtls = {
-  SetSyncSecretIntegrationLastRunTimestampInSeconds: 10
+  SetSyncSecretIntegrationLastRunTimestampInSeconds: 10,
+  AccessTokenStatusUpdateInSeconds: 120
 };
 
 type TWaitTillReady = {

backend/src/lib/dates/index.ts

@@ -1,3 +1,8 @@
 export const daysToMillisecond = (days: number) => days * 24 * 60 * 60 * 1000;
 
 export const secondsToMillis = (seconds: number) => seconds * 1000;
+
+export const applyJitter = (delayMs: number, jitterMs: number) => {
+  const jitter = Math.floor(Math.random() * (2 * jitterMs)) - jitterMs;
+  return delayMs + jitter;
+};

backend/src/queue/queue-service.ts

@@ -27,7 +27,8 @@ export enum QueueName {
   CaCrlRotation = "ca-crl-rotation",
   SecretReplication = "secret-replication",
   SecretSync = "secret-sync", // parent queue to push integration sync, webhook, and secret replication
-  ProjectV3Migration = "project-v3-migration"
+  ProjectV3Migration = "project-v3-migration",
+  AccessTokenStatusUpdate = "access-token-status-update"
 }
 
 export enum QueueJobs {
@@ -48,7 +49,9 @@ export enum QueueJobs {
   CaCrlRotation = "ca-crl-rotation-job",
   SecretReplication = "secret-replication",
   SecretSync = "secret-sync", // parent queue to push integration sync, webhook, and secret replication
-  ProjectV3Migration = "project-v3-migration"
+  ProjectV3Migration = "project-v3-migration",
+  IdentityAccessTokenStatusUpdate = "identity-access-token-status-update",
+  ServiceTokenStatusUpdate = "service-token-status-update"
 }
 
 export type TQueueJobTypes = {
@@ -148,6 +151,15 @@ export type TQueueJobTypes = {
     name: QueueJobs.ProjectV3Migration;
     payload: { projectId: string };
   };
+  [QueueName.AccessTokenStatusUpdate]:
+    | {
+        name: QueueJobs.IdentityAccessTokenStatusUpdate;
+        payload: { identityAccessTokenId: string; numberOfUses: number };
+      }
+    | {
+        name: QueueJobs.ServiceTokenStatusUpdate;
+        payload: { serviceTokenId: string };
+      };
 };
 
 export type TQueueServiceFactory = ReturnType<typeof queueServiceFactory>;

backend/src/server/routes/index.ts

@@ -73,6 +73,7 @@ import { TKeyStoreFactory } from "@app/keystore/keystore";
 import { getConfig } from "@app/lib/config/env";
 import { TQueueServiceFactory } from "@app/queue";
 import { readLimit } from "@app/server/config/rateLimiter";
+import { accessTokenQueueServiceFactory } from "@app/services/access-token-queue/access-token-queue";
 import { apiKeyDALFactory } from "@app/services/api-key/api-key-dal";
 import { apiKeyServiceFactory } from "@app/services/api-key/api-key-service";
 import { authDALFactory } from "@app/services/auth/auth-dal";
@@ -953,12 +954,20 @@ export const registerRoutes = async (
     kmsService
   });
 
+  const accessTokenQueue = accessTokenQueueServiceFactory({
+    keyStore,
+    identityAccessTokenDAL,
+    queueService,
+    serviceTokenDAL
+  });
+
   const serviceTokenService = serviceTokenServiceFactory({
     projectEnvDAL,
     serviceTokenDAL,
     userDAL,
     permissionService,
-    projectDAL
+    projectDAL,
+    accessTokenQueue
   });
 
   const identityService = identityServiceFactory({
@@ -968,10 +977,13 @@ export const registerRoutes = async (
     identityProjectDAL,
     licenseService
   });
+
   const identityAccessTokenService = identityAccessTokenServiceFactory({
     identityAccessTokenDAL,
-    identityOrgMembershipDAL
+    identityOrgMembershipDAL,
+    accessTokenQueue
   });
+
   const identityProjectService = identityProjectServiceFactory({
     permissionService,
     projectDAL,

backend/src/services/access-token-queue/access-token-queue.ts

@@ -0,0 +1,125 @@
+import { z } from "zod";
+
+import { KeyStorePrefixes, KeyStoreTtls, TKeyStoreFactory } from "@app/keystore/keystore";
+import { applyJitter, secondsToMillis } from "@app/lib/dates";
+import { logger } from "@app/lib/logger";
+import { QueueJobs, QueueName, TQueueServiceFactory } from "@app/queue";
+
+import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
+import { TServiceTokenDALFactory } from "../service-token/service-token-dal";
+
+type TAccessTokenQueueServiceFactoryDep = {
+  queueService: TQueueServiceFactory;
+  keyStore: Pick<TKeyStoreFactory, "getItem" | "setItemWithExpiry">;
+  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "updateById">;
+  serviceTokenDAL: Pick<TServiceTokenDALFactory, "updateById">;
+};
+
+export type TAccessTokenQueueServiceFactory = ReturnType<typeof accessTokenQueueServiceFactory>;
+
+export const AccessTokenStatusSchema = z.object({
+  lastUpdatedAt: z.string().datetime(),
+  numberOfUses: z.number()
+});
+
+export const accessTokenQueueServiceFactory = ({
+  queueService,
+  keyStore,
+  identityAccessTokenDAL,
+  serviceTokenDAL
+}: TAccessTokenQueueServiceFactoryDep) => {
+  const getIdentityTokenDetailsInCache = async (identityAccessTokenId: string) => {
+    const tokenDetailsInCache = await keyStore.getItem(
+      KeyStorePrefixes.IdentityAccessTokenStatusUpdate(identityAccessTokenId)
+    );
+    if (tokenDetailsInCache) {
+      return AccessTokenStatusSchema.parseAsync(JSON.parse(tokenDetailsInCache));
+    }
+  };
+
+  const updateServiceTokenStatus = async (serviceTokenId: string) => {
+    await keyStore.setItemWithExpiry(
+      KeyStorePrefixes.ServiceTokenStatusUpdate(serviceTokenId),
+      KeyStoreTtls.AccessTokenStatusUpdateInSeconds,
+      JSON.stringify({ lastUpdatedAt: new Date() })
+    );
+    await queueService.queue(
+      QueueName.AccessTokenStatusUpdate,
+      QueueJobs.ServiceTokenStatusUpdate,
+      {
+        serviceTokenId
+      },
+      {
+        delay: applyJitter(secondsToMillis(KeyStoreTtls.AccessTokenStatusUpdateInSeconds / 2), secondsToMillis(10)),
+        // https://docs.bullmq.io/guide/jobs/job-ids
+        jobId: KeyStorePrefixes.ServiceTokenStatusUpdate(serviceTokenId).replaceAll(":", "_"),
+        removeOnFail: true,
+        removeOnComplete: true
+      }
+    );
+  };
+
+  const updateIdentityAccessTokenStatus = async (identityAccessTokenId: string, numberOfUses: number) => {
+    await keyStore.setItemWithExpiry(
+      KeyStorePrefixes.IdentityAccessTokenStatusUpdate(identityAccessTokenId),
+      KeyStoreTtls.AccessTokenStatusUpdateInSeconds,
+      JSON.stringify({ lastUpdatedAt: new Date(), numberOfUses })
+    );
+    await queueService.queue(
+      QueueName.AccessTokenStatusUpdate,
+      QueueJobs.IdentityAccessTokenStatusUpdate,
+      {
+        identityAccessTokenId,
+        numberOfUses
+      },
+      {
+        delay: applyJitter(secondsToMillis(KeyStoreTtls.AccessTokenStatusUpdateInSeconds / 2), secondsToMillis(10)),
+        jobId: KeyStorePrefixes.IdentityAccessTokenStatusUpdate(identityAccessTokenId).replaceAll(":", "_"),
+        removeOnFail: true,
+        removeOnComplete: true
+      }
+    );
+  };
+
+  queueService.start(QueueName.AccessTokenStatusUpdate, async (job) => {
+    // for identity token update
+    if (job.name === QueueJobs.IdentityAccessTokenStatusUpdate && "identityAccessTokenId" in job.data) {
+      const { identityAccessTokenId } = job.data;
+      const tokenDetails = { lastUpdatedAt: new Date(job.timestamp), numberOfUses: job.data.numberOfUses };
+      const tokenDetailsInCache = await getIdentityTokenDetailsInCache(identityAccessTokenId);
+      if (tokenDetailsInCache) {
+        tokenDetails.numberOfUses = tokenDetailsInCache.numberOfUses;
+        tokenDetails.lastUpdatedAt = new Date(tokenDetailsInCache.lastUpdatedAt);
+      }
+
+      await identityAccessTokenDAL.updateById(identityAccessTokenId, {
+        accessTokenLastUsedAt: tokenDetails.lastUpdatedAt,
+        accessTokenNumUses: tokenDetails.numberOfUses
+      });
+      return;
+    }
+
+    // for service token
+    if (job.name === QueueJobs.ServiceTokenStatusUpdate && "serviceTokenId" in job.data) {
+      const { serviceTokenId } = job.data;
+      const tokenDetailsInCache = await keyStore.getItem(KeyStorePrefixes.ServiceTokenStatusUpdate(serviceTokenId));
+      let lastUsed = new Date(job.timestamp);
+      if (tokenDetailsInCache) {
+        const tokenDetails = await AccessTokenStatusSchema.pick({ lastUpdatedAt: true }).parseAsync(
+          JSON.parse(tokenDetailsInCache)
+        );
+        lastUsed = new Date(tokenDetails.lastUpdatedAt);
+      }
+
+      await serviceTokenDAL.updateById(serviceTokenId, {
+        lastUsed
+      });
+    }
+  });
+
+  queueService.listen(QueueName.AccessTokenStatusUpdate, "failed", (_, err) => {
+    logger.error(err, `${QueueName.AccessTokenStatusUpdate}: Failed to updated access token status`);
+  });
+
+  return { updateIdentityAccessTokenStatus, updateServiceTokenStatus, getIdentityTokenDetailsInCache };
+};

backend/src/services/identity-access-token/identity-access-token-service.ts

@@ -5,6 +5,7 @@ import { getConfig } from "@app/lib/config/env";
 import { BadRequestError, UnauthorizedError } from "@app/lib/errors";
 import { checkIPAgainstBlocklist, TIp } from "@app/lib/ip";
 
+import { TAccessTokenQueueServiceFactory } from "../access-token-queue/access-token-queue";
 import { AuthTokenType } from "../auth/auth-type";
 import { TIdentityOrgDALFactory } from "../identity/identity-org-dal";
 import { TIdentityAccessTokenDALFactory } from "./identity-access-token-dal";
@@ -13,19 +14,24 @@ import { TIdentityAccessTokenJwtPayload, TRenewAccessTokenDTO } from "./identity
 type TIdentityAccessTokenServiceFactoryDep = {
   identityAccessTokenDAL: TIdentityAccessTokenDALFactory;
   identityOrgMembershipDAL: TIdentityOrgDALFactory;
+  accessTokenQueue: Pick<
+    TAccessTokenQueueServiceFactory,
+    "updateIdentityAccessTokenStatus" | "getIdentityTokenDetailsInCache"
+  >;
 };
 
 export type TIdentityAccessTokenServiceFactory = ReturnType<typeof identityAccessTokenServiceFactory>;
 
 export const identityAccessTokenServiceFactory = ({
   identityAccessTokenDAL,
-  identityOrgMembershipDAL
+  identityOrgMembershipDAL,
+  accessTokenQueue
 }: TIdentityAccessTokenServiceFactoryDep) => {
   const validateAccessTokenExp = async (identityAccessToken: TIdentityAccessTokens) => {
     const {
       id: tokenId,
-      accessTokenTTL,
       accessTokenNumUses,
+      accessTokenTTL,
       accessTokenNumUsesLimit,
       accessTokenLastRenewedAt,
       createdAt: accessTokenCreatedAt
@@ -83,7 +89,12 @@ export const identityAccessTokenServiceFactory = ({
     });
     if (!identityAccessToken) throw new UnauthorizedError();
 
-    await validateAccessTokenExp(identityAccessToken);
+    let { accessTokenNumUses } = identityAccessToken;
+    const tokenStatusInCache = await accessTokenQueue.getIdentityTokenDetailsInCache(identityAccessToken.id);
+    if (tokenStatusInCache) {
+      accessTokenNumUses = tokenStatusInCache.numberOfUses;
+    }
+    await validateAccessTokenExp({ ...identityAccessToken, accessTokenNumUses });
 
     const { accessTokenMaxTTL, createdAt: accessTokenCreatedAt, accessTokenTTL } = identityAccessToken;
 
@@ -164,14 +175,14 @@ export const identityAccessTokenServiceFactory = ({
       throw new UnauthorizedError({ message: "Identity does not belong to any organization" });
     }
 
-    await validateAccessTokenExp(identityAccessToken);
+    let { accessTokenNumUses } = identityAccessToken;
+    const tokenStatusInCache = await accessTokenQueue.getIdentityTokenDetailsInCache(identityAccessToken.id);
+    if (tokenStatusInCache) {
+      accessTokenNumUses = tokenStatusInCache.numberOfUses;
+    }
+    await validateAccessTokenExp({ ...identityAccessToken, accessTokenNumUses });
 
-    await identityAccessTokenDAL.updateById(identityAccessToken.id, {
+    await accessTokenQueue.updateIdentityAccessTokenStatus(identityAccessToken.id, Number(accessTokenNumUses) + 1);
-      accessTokenLastUsedAt: new Date(),
-      $incr: {
-        accessTokenNumUses: 1
-      }
-    });
     return { ...identityAccessToken, orgId: identityOrgMembership.orgId };
   };
 

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

@@ -8,6 +8,7 @@ import { ProjectPermissionActions, ProjectPermissionSub } from "@app/ee/services
 import { getConfig } from "@app/lib/config/env";
 import { BadRequestError, UnauthorizedError } from "@app/lib/errors";
 
+import { TAccessTokenQueueServiceFactory } from "../access-token-queue/access-token-queue";
 import { ActorType } from "../auth/auth-type";
 import { TProjectDALFactory } from "../project/project-dal";
 import { TProjectEnvDALFactory } from "../project-env/project-env-dal";
@@ -26,6 +27,7 @@ type TServiceTokenServiceFactoryDep = {
   permissionService: Pick<TPermissionServiceFactory, "getProjectPermission">;
   projectEnvDAL: Pick<TProjectEnvDALFactory, "findBySlugs">;
   projectDAL: Pick<TProjectDALFactory, "findById">;
+  accessTokenQueue: Pick<TAccessTokenQueueServiceFactory, "updateServiceTokenStatus">;
 };
 
 export type TServiceTokenServiceFactory = ReturnType<typeof serviceTokenServiceFactory>;
@@ -35,7 +37,8 @@ export const serviceTokenServiceFactory = ({
   userDAL,
   permissionService,
   projectEnvDAL,
-  projectDAL
+  projectDAL,
+  accessTokenQueue
 }: TServiceTokenServiceFactoryDep) => {
   const createServiceToken = async ({
     iv,
@@ -166,11 +169,9 @@ export const serviceTokenServiceFactory = ({
 
     const isMatch = await bcrypt.compare(TOKEN_SECRET, serviceToken.secretHash);
     if (!isMatch) throw new UnauthorizedError();
-    const updatedToken = await serviceTokenDAL.updateById(serviceToken.id, {
+    await accessTokenQueue.updateServiceTokenStatus(serviceToken.id);
-      lastUsed: new Date()
-    });
 
-    return { ...serviceToken, lastUsed: updatedToken.lastUsed, orgId: project.orgId };
+    return { ...serviceToken, lastUsed: new Date(), orgId: project.orgId };
   };
 
   return {