feat: added migration and identity check

2026-09-22 13:39:35 +00:00 · 2025-10-20 16:20:39 +05:30
parent 2ec851fd34
commit 554c87dc9b
27 changed files with 181 additions and 198 deletions
--- a/backend/src/db/migrations/20251018061215_sub-org.ts
+++ b/backend/src/db/migrations/20251018061215_sub-org.ts
@@ -1,6 +1,6 @@
 import { Knex } from "knex";

-import { TableName } from "../schemas";
+import { AccessScope, TableName } from "../schemas";

 export async function up(knex: Knex): Promise<void> {
  const hasParentOrgId = await knex.schema.hasColumn(TableName.Organization, "parentOrgId");
@@ -18,9 +18,25 @@ export async function up(knex: Knex): Promise<void> {
  const hasIdentityOrgCol = await knex.schema.hasColumn(TableName.Identity, "orgId");
  if (!hasIdentityOrgCol) {
    await knex.schema.alterTable(TableName.Identity, (t) => {
-      t.uuid("orgId").notNullable();
+      t.uuid("orgId");
      t.foreign("orgId").references("id").inTable(TableName.Organization).onDelete("CASCADE");
    });
+
+    await knex.raw(
+      `
+  UPDATE ?? AS identity
+  SET "orgId" = membership."scopeOrgId"
+  FROM ?? AS membership
+  WHERE 
+    membership."actorIdentityId" = identity."id"
+    AND membership."scope" = ?
+`,
+      [TableName.Identity, TableName.Membership, AccessScope.Organization]
+    );
+
+    await knex.schema.alterTable(TableName.Identity, (t) => {
+      t.uuid("orgId").notNullable();
+    });
  }
 }

--- a/backend/src/server/routes/index.ts
+++ b/backend/src/server/routes/index.ts
@@ -1618,6 +1618,7 @@ export const registerRoutes = async (
  });

  const identityTokenAuthService = identityTokenAuthServiceFactory({
+    identityDAL,
    identityTokenAuthDAL,
    identityAccessTokenDAL,
    permissionService,
@@ -1627,6 +1628,7 @@ export const registerRoutes = async (
  });

  const identityUaService = identityUaServiceFactory({
+    identityDAL,
    permissionService,
    identityAccessTokenDAL,
    identityUaClientSecretDAL,
@@ -1638,6 +1640,7 @@ export const registerRoutes = async (
  });

  const identityKubernetesAuthService = identityKubernetesAuthServiceFactory({
+    identityDAL,
    identityKubernetesAuthDAL,
    identityAccessTokenDAL,
    permissionService,
@@ -1651,6 +1654,7 @@ export const registerRoutes = async (
    membershipIdentityDAL
  });
  const identityGcpAuthService = identityGcpAuthServiceFactory({
+    identityDAL,
    identityGcpAuthDAL,
    orgDAL,
    identityAccessTokenDAL,
@@ -1660,6 +1664,7 @@ export const registerRoutes = async (
  });

  const identityAliCloudAuthService = identityAliCloudAuthServiceFactory({
+    identityDAL,
    identityAccessTokenDAL,
    orgDAL,
    identityAliCloudAuthDAL,
@@ -1669,6 +1674,7 @@ export const registerRoutes = async (
  });

  const identityTlsCertAuthService = identityTlsCertAuthServiceFactory({
+    identityDAL,
    identityAccessTokenDAL,
    identityTlsCertAuthDAL,
    licenseService,
@@ -1678,6 +1684,7 @@ export const registerRoutes = async (
  });

  const identityAwsAuthService = identityAwsAuthServiceFactory({
+    identityDAL,
    identityAccessTokenDAL,
    orgDAL,
    identityAwsAuthDAL,
@@ -1687,6 +1694,7 @@ export const registerRoutes = async (
  });

  const identityAzureAuthService = identityAzureAuthServiceFactory({
+    identityDAL,
    identityAzureAuthDAL,
    orgDAL,
    identityAccessTokenDAL,
@@ -1696,6 +1704,7 @@ export const registerRoutes = async (
  });

  const identityOciAuthService = identityOciAuthServiceFactory({
+    identityDAL,
    identityAccessTokenDAL,
    orgDAL,
    identityOciAuthDAL,
@@ -1719,6 +1728,7 @@ export const registerRoutes = async (
  });

  const identityOidcAuthService = identityOidcAuthServiceFactory({
+    identityDAL,
    identityOidcAuthDAL,
    orgDAL,
    identityAccessTokenDAL,
@@ -1729,6 +1739,7 @@ export const registerRoutes = async (
  });

  const identityJwtAuthService = identityJwtAuthServiceFactory({
+    identityDAL,
    identityJwtAuthDAL,
    orgDAL,
    permissionService,
--- a/backend/src/server/routes/v1/identity-alicloud-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-alicloud-auth-router.ts
@@ -73,12 +73,12 @@ export const registerIdentityAliCloudAuthRouter = async (server: FastifyZodProvi
      }
    },
    handler: async (req) => {
-      const { identityAliCloudAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityAliCloudAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityAliCloudAuth.login(req.body);

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_ALICLOUD_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-aws-iam-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-aws-iam-auth-router.ts
@@ -40,12 +40,12 @@ export const registerIdentityAwsAuthRouter = async (server: FastifyZodProvider)
      }
    },
    handler: async (req) => {
-      const { identityAwsAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityAwsAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityAwsAuth.login(req.body);

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_AWS_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-azure-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-azure-auth-router.ts
@@ -35,12 +35,12 @@ export const registerIdentityAzureAuthRouter = async (server: FastifyZodProvider
      }
    },
    handler: async (req) => {
-      const { identityAzureAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityAzureAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityAzureAuth.login(req.body);

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_AZURE_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-gcp-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-gcp-auth-router.ts
@@ -35,12 +35,12 @@ export const registerIdentityGcpAuthRouter = async (server: FastifyZodProvider)
      }
    },
    handler: async (req) => {
-      const { identityGcpAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityGcpAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityGcpAuth.login(req.body);

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_GCP_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-jwt-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-jwt-auth-router.ts
@@ -111,7 +111,7 @@ export const registerIdentityJwtAuthRouter = async (server: FastifyZodProvider)
      }
    },
    handler: async (req) => {
-      const { identityJwtAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityJwtAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityJwtAuth.login({
          identityId: req.body.identityId,
          jwt: req.body.jwt
@@ -119,7 +119,7 @@ export const registerIdentityJwtAuthRouter = async (server: FastifyZodProvider)

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_JWT_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-kubernetes-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-kubernetes-auth-router.ts
@@ -56,7 +56,7 @@ export const registerIdentityKubernetesRouter = async (server: FastifyZodProvide
      }
    },
    handler: async (req) => {
-      const { identityKubernetesAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityKubernetesAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityKubernetesAuth.login({
          identityId: req.body.identityId,
          jwt: req.body.jwt
@@ -64,7 +64,7 @@ export const registerIdentityKubernetesRouter = async (server: FastifyZodProvide

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_KUBERNETES_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-ldap-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-ldap-auth-router.ts
@@ -162,13 +162,13 @@ export const registerIdentityLdapAuthRouter = async (server: FastifyZodProvider)

      const { identityId, user } = req.passportMachineIdentity;

-      const { accessToken, identityLdapAuth, identityMembershipOrg } = await server.services.identityLdapAuth.login({
+      const { accessToken, identityLdapAuth, identity } = await server.services.identityLdapAuth.login({
        identityId
      });

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_LDAP_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-oci-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-oci-auth-router.ts
@@ -52,12 +52,12 @@ export const registerIdentityOciAuthRouter = async (server: FastifyZodProvider)
      }
    },
    handler: async (req) => {
-      const { identityOciAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityOciAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityOciAuth.login(req.body);

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_OCI_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-oidc-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-oidc-auth-router.ts
@@ -59,7 +59,7 @@ export const registerIdentityOidcAuthRouter = async (server: FastifyZodProvider)
      }
    },
    handler: async (req) => {
-      const { identityOidcAuth, accessToken, identityAccessToken, identityMembershipOrg, oidcTokenData } =
+      const { identityOidcAuth, accessToken, identityAccessToken, identity, oidcTokenData } =
        await server.services.identityOidcAuth.login({
          identityId: req.body.identityId,
          jwt: req.body.jwt
@@ -67,7 +67,7 @@ export const registerIdentityOidcAuthRouter = async (server: FastifyZodProvider)

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_OIDC_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-tls-cert-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-tls-cert-auth-router.ts
@@ -64,7 +64,7 @@ export const registerIdentityTlsCertAuthRouter = async (server: FastifyZodProvid
        throw new BadRequestError({ message: "Missing TLS certificate in header" });
      }

-      const { identityTlsCertAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityTlsCertAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityTlsCertAuth.login({
          identityId: req.body.identityId,
          clientCertificate: clientCertificate as string
@@ -72,7 +72,7 @@ export const registerIdentityTlsCertAuthRouter = async (server: FastifyZodProvid

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_TLS_CERT_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-token-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-token-auth-router.ts
@@ -319,7 +319,7 @@ export const registerIdentityTokenAuthRouter = async (server: FastifyZodProvider
      }
    },
    handler: async (req) => {
-      const { identityTokenAuth, accessToken, identityAccessToken, identityMembershipOrg } =
+      const { identityTokenAuth, accessToken, identityAccessToken, identity } =
        await server.services.identityTokenAuth.createTokenAuthToken({
          actor: req.permission.type,
          actorId: req.permission.id,
@@ -332,7 +332,7 @@ export const registerIdentityTokenAuthRouter = async (server: FastifyZodProvider

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.CREATE_TOKEN_IDENTITY_TOKEN_AUTH,
          metadata: {
--- a/backend/src/server/routes/v1/identity-universal-auth-router.ts
+++ b/backend/src/server/routes/v1/identity-universal-auth-router.ts
@@ -52,14 +52,14 @@ export const registerIdentityUaRouter = async (server: FastifyZodProvider) => {
        accessToken,
        identityAccessToken,
        validClientSecretInfo,
-        identityMembershipOrg,
+        identity,
        accessTokenTTL,
        accessTokenMaxTTL
      } = await server.services.identityUa.login(req.body.clientId, req.body.clientSecret, req.realIp);

      await server.services.auditLog.createAuditLog({
        ...req.auditLogInfo,
-        orgId: identityMembershipOrg.scopeOrgId,
+        orgId: identity.orgId,
        event: {
          type: EventType.LOGIN_IDENTITY_UNIVERSAL_AUTH,
          metadata: {
--- a/backend/src/services/identity-alicloud-auth/identity-alicloud-auth-service.ts
+++ b/backend/src/services/identity-alicloud-auth/identity-alicloud-auth-service.ts
@@ -26,6 +26,7 @@ import { logger } from "@app/lib/logger";
 import { ActorType, AuthTokenType } from "../auth/auth-type";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TMembershipIdentityDALFactory } from "../membership-identity/membership-identity-dal";
 import { TOrgDALFactory } from "../org/org-dal";
 import { validateIdentityUpdateForSuperAdminPrivileges } from "../super-admin/super-admin-fns";
@@ -40,12 +41,13 @@ import {
 } from "./identity-alicloud-auth-types";

 type TIdentityAliCloudAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  identityAliCloudAuthDAL: Pick<
    TIdentityAliCloudAuthDALFactory,
    "findOne" | "transaction" | "create" | "updateById" | "delete"
  >;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  orgDAL: Pick<TOrgDALFactory, "findById">;
@@ -54,6 +56,7 @@ type TIdentityAliCloudAuthServiceFactoryDep = {
 export type TIdentityAliCloudAuthServiceFactory = ReturnType<typeof identityAliCloudAuthServiceFactory>;

 export const identityAliCloudAuthServiceFactory = ({
+  identityDAL,
  identityAccessTokenDAL,
  identityAliCloudAuthDAL,
  membershipIdentityDAL,
@@ -69,12 +72,8 @@ export const identityAliCloudAuthServiceFactory = ({
      });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityAliCloudAuth.identityId,
-      scope: AccessScope.Organization
-    });
-
-    if (!identityMembershipOrg) throw new UnauthorizedError({ message: "Identity not attached to a organization" });
+    const identity = await identityDAL.findById(identityAliCloudAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const requestUrl = new URL("https://sts.aliyuncs.com");

@@ -99,8 +98,8 @@ export const identityAliCloudAuthServiceFactory = ({

    // Generate the token
    const identityAccessToken = await identityAliCloudAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
        {
          lastLoginAuthMethod: IdentityAuthMethod.ALICLOUD_AUTH,
          lastLoginTime: new Date()
@@ -141,7 +140,7 @@ export const identityAliCloudAuthServiceFactory = ({
      identityAliCloudAuth,
      accessToken,
      identityAccessToken,
-      identityMembershipOrg
+      identity
    };
  };

--- a/backend/src/services/identity-aws-auth/identity-aws-auth-service.ts
+++ b/backend/src/services/identity-aws-auth/identity-aws-auth-service.ts
@@ -25,6 +25,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { ActorType, AuthTokenType } from "../auth/auth-type";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TMembershipIdentityDALFactory } from "../membership-identity/membership-identity-dal";
 import { TOrgDALFactory } from "../org/org-dal";
 import { validateIdentityUpdateForSuperAdminPrivileges } from "../super-admin/super-admin-fns";
@@ -41,9 +42,10 @@ import {
 } from "./identity-aws-auth-types";

 type TIdentityAwsAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  identityAwsAuthDAL: Pick<TIdentityAwsAuthDALFactory, "findOne" | "transaction" | "create" | "updateById" | "delete">;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  orgDAL: Pick<TOrgDALFactory, "findById">;
@@ -86,6 +88,7 @@ function isValidAwsRegion(region: string | null): boolean {
 }

 export const identityAwsAuthServiceFactory = ({
+  identityDAL,
  identityAccessTokenDAL,
  identityAwsAuthDAL,
  membershipIdentityDAL,
@@ -99,11 +102,8 @@ export const identityAwsAuthServiceFactory = ({
      throw new NotFoundError({ message: "AWS auth method not found for identity, did you configure AWS auth?" });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityAwsAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) throw new UnauthorizedError({ message: "Identity not attached to a organization" });
+    const identity = await identityDAL.findById(identityAwsAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const headers: TAwsGetCallerIdentityHeaders = JSON.parse(Buffer.from(iamRequestHeaders, "base64").toString());
    const body: string = Buffer.from(iamRequestBody, "base64").toString();
@@ -165,8 +165,8 @@ export const identityAwsAuthServiceFactory = ({
    }

    const identityAccessToken = await identityAwsAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
        {
          lastLoginAuthMethod: IdentityAuthMethod.AWS_AUTH,
          lastLoginTime: new Date()
@@ -218,7 +218,7 @@ export const identityAwsAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityAwsAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityAwsAuth, identityAccessToken, identity };
  };

  const attachAwsAuth = async ({
--- a/backend/src/services/identity-azure-auth/identity-azure-auth-service.ts
+++ b/backend/src/services/identity-azure-auth/identity-azure-auth-service.ts
@@ -22,6 +22,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { ActorType, AuthTokenType } from "../auth/auth-type";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TMembershipIdentityDALFactory } from "../membership-identity/membership-identity-dal";
 import { TOrgDALFactory } from "../org/org-dal";
 import { validateIdentityUpdateForSuperAdminPrivileges } from "../super-admin/super-admin-fns";
@@ -36,11 +37,12 @@ import {
 } from "./identity-azure-auth-types";

 type TIdentityAzureAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityAzureAuthDAL: Pick<
    TIdentityAzureAuthDALFactory,
    "findOne" | "transaction" | "create" | "updateById" | "delete"
  >;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
@@ -50,6 +52,7 @@ type TIdentityAzureAuthServiceFactoryDep = {
 export type TIdentityAzureAuthServiceFactory = ReturnType<typeof identityAzureAuthServiceFactory>;

 export const identityAzureAuthServiceFactory = ({
+  identityDAL,
  identityAzureAuthDAL,
  membershipIdentityDAL,
  identityAccessTokenDAL,
@@ -63,11 +66,8 @@ export const identityAzureAuthServiceFactory = ({
      throw new NotFoundError({ message: "Azure auth method not found for identity, did you configure Azure Auth?" });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityAzureAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) throw new UnauthorizedError({ message: "Identity not attached to a organization" });
+    const identity = await identityDAL.findById(identityAzureAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const azureIdentity = await validateAzureIdentity({
      tenantId: identityAzureAuth.tenantId,
@@ -92,8 +92,8 @@ export const identityAzureAuthServiceFactory = ({
    }

    const identityAccessToken = await identityAzureAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
        {
          lastLoginAuthMethod: IdentityAuthMethod.AZURE_AUTH,
          lastLoginTime: new Date()
@@ -131,7 +131,7 @@ export const identityAzureAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityAzureAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityAzureAuth, identityAccessToken, identity };
  };

  const attachAzureAuth = async ({
--- a/backend/src/services/identity-gcp-auth/identity-gcp-auth-service.ts
+++ b/backend/src/services/identity-gcp-auth/identity-gcp-auth-service.ts
@@ -22,6 +22,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { ActorType, AuthTokenType } from "../auth/auth-type";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TMembershipIdentityDALFactory } from "../membership-identity/membership-identity-dal";
 import { TOrgDALFactory } from "../org/org-dal";
 import { validateIdentityUpdateForSuperAdminPrivileges } from "../super-admin/super-admin-fns";
@@ -37,8 +38,9 @@ import {
 } from "./identity-gcp-auth-types";

 type TIdentityGcpAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityGcpAuthDAL: Pick<TIdentityGcpAuthDALFactory, "findOne" | "transaction" | "create" | "updateById" | "delete">;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
@@ -48,6 +50,7 @@ type TIdentityGcpAuthServiceFactoryDep = {
 export type TIdentityGcpAuthServiceFactory = ReturnType<typeof identityGcpAuthServiceFactory>;

 export const identityGcpAuthServiceFactory = ({
+  identityDAL,
  identityGcpAuthDAL,
  membershipIdentityDAL,
  identityAccessTokenDAL,
@@ -61,13 +64,8 @@ export const identityGcpAuthServiceFactory = ({
      throw new NotFoundError({ message: "GCP auth method not found for identity, did you configure GCP auth?" });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityGcpAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) {
-      throw new UnauthorizedError({ message: "Identity does not belong to any organization" });
-    }
+    const identity = await identityDAL.findById(identityGcpAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    let gcpIdentityDetails: TGcpIdentityDetails;
    switch (identityGcpAuth.type) {
@@ -131,8 +129,8 @@ export const identityGcpAuthServiceFactory = ({
    }

    const identityAccessToken = await identityGcpAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
        {
          lastLoginAuthMethod: IdentityAuthMethod.GCP_AUTH,
          lastLoginTime: new Date()
@@ -170,7 +168,7 @@ export const identityGcpAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityGcpAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityGcpAuth, identityAccessToken, identity };
  };

  const attachGcpAuth = async ({
--- a/backend/src/services/identity-jwt-auth/identity-jwt-auth-service.ts
+++ b/backend/src/services/identity-jwt-auth/identity-jwt-auth-service.ts
@@ -24,6 +24,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { getValueByDot } from "@app/lib/template/dot-access";

 import { ActorType, AuthTokenType } from "../auth/auth-type";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
 import { TKmsServiceFactory } from "../kms/kms-service";
@@ -43,8 +44,9 @@ import {
 } from "./identity-jwt-auth-types";

 type TIdentityJwtAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityJwtAuthDAL: TIdentityJwtAuthDALFactory;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
@@ -55,6 +57,7 @@ type TIdentityJwtAuthServiceFactoryDep = {
 export type TIdentityJwtAuthServiceFactory = ReturnType<typeof identityJwtAuthServiceFactory>;

 export const identityJwtAuthServiceFactory = ({
+  identityDAL,
  identityJwtAuthDAL,
  membershipIdentityDAL,
  permissionService,
@@ -69,19 +72,12 @@ export const identityJwtAuthServiceFactory = ({
      throw new NotFoundError({ message: "JWT auth method not found for identity, did you configure JWT auth?" });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityJwtAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) {
-      throw new NotFoundError({
-        message: `Identity organization membership for identity with ID '${identityJwtAuth.identityId}' not found`
-      });
-    }
+    const identity = await identityDAL.findById(identityJwtAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const { decryptor: orgDataKeyDecryptor } = await kmsService.createCipherPairWithDataKey({
      type: KmsDataKey.Organization,
-      orgId: identityMembershipOrg.scopeOrgId
+      orgId: identity.orgId
    });

    const decodedToken = crypto.jwt().decode(jwtValue, { complete: true });
@@ -211,12 +207,9 @@ export const identityJwtAuthServiceFactory = ({
    }

    const identityAccessToken = await identityJwtAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.JWT_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.JWT_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -251,7 +244,7 @@ export const identityJwtAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityJwtAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityJwtAuth, identityAccessToken, identity };
  };

  const attachJwtAuth = async ({
--- a/backend/src/services/identity-kubernetes-auth/identity-kubernetes-auth-service.ts
+++ b/backend/src/services/identity-kubernetes-auth/identity-kubernetes-auth-service.ts
@@ -39,6 +39,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { logger } from "@app/lib/logger";

 import { ActorType, AuthTokenType } from "../auth/auth-type";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
 import { TKmsServiceFactory } from "../kms/kms-service";
@@ -59,12 +60,13 @@ import {
 } from "./identity-kubernetes-auth-types";

 type TIdentityKubernetesAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityKubernetesAuthDAL: Pick<
    TIdentityKubernetesAuthDALFactory,
    "create" | "findOne" | "transaction" | "updateById" | "delete"
  >;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
  kmsService: Pick<TKmsServiceFactory, "createCipherPairWithDataKey">;
@@ -80,6 +82,7 @@ export type TIdentityKubernetesAuthServiceFactory = ReturnType<typeof identityKu
 const GATEWAY_AUTH_DEFAULT_HOST = "https://kubernetes.default.svc.cluster.local";

 export const identityKubernetesAuthServiceFactory = ({
+  identityDAL,
  identityKubernetesAuthDAL,
  membershipIdentityDAL,
  identityAccessTokenDAL,
@@ -186,19 +189,12 @@ export const identityKubernetesAuthServiceFactory = ({
      });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityKubernetesAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) {
-      throw new NotFoundError({
-        message: `Identity organization membership for identity with ID '${identityKubernetesAuth.identityId}' not found`
-      });
-    }
+    const identity = await identityDAL.findById(identityKubernetesAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const { decryptor } = await kmsService.createCipherPairWithDataKey({
      type: KmsDataKey.Organization,
-      orgId: identityMembershipOrg.scopeOrgId
+      orgId: identity.orgId
    });

    let caCert = "";
@@ -441,12 +437,9 @@ export const identityKubernetesAuthServiceFactory = ({
    }

    const identityAccessToken = await identityKubernetesAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.KUBERNETES_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.KUBERNETES_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -486,7 +479,7 @@ export const identityKubernetesAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityKubernetesAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityKubernetesAuth, identityAccessToken, identity };
  };

  const attachKubernetesAuth = async ({
--- a/backend/src/services/identity-ldap-auth/identity-ldap-auth-service.ts
+++ b/backend/src/services/identity-ldap-auth/identity-ldap-auth-service.ts
@@ -57,11 +57,11 @@ type TIdentityLdapAuthServiceFactoryDep = {
    TIdentityLdapAuthDALFactory,
    "findOne" | "transaction" | "create" | "updateById" | "delete"
  >;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  kmsService: TKmsServiceFactory;
-  identityDAL: TIdentityDALFactory;
+  identityDAL: Pick<TIdentityDALFactory, "findById" | "findOne">;
  identityAuthTemplateDAL: TIdentityAuthTemplateDALFactory;
  keyStore: Pick<
    TKeyStoreFactory,
@@ -151,17 +151,6 @@ export const identityLdapAuthServiceFactory = ({
  };

  const login = async ({ identityId }: TLoginLdapAuthDTO) => {
-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityId,
-      scope: AccessScope.Organization
-    });
-
-    if (!identityMembershipOrg) {
-      throw new UnauthorizedError({
-        message: "Invalid credentials"
-      });
-    }
-
    const identityLdapAuth = await identityLdapAuthDAL.findOne({ identityId });

    if (!identityLdapAuth) {
@@ -170,7 +159,10 @@ export const identityLdapAuthServiceFactory = ({
      });
    }

-    const plan = await licenseService.getPlan(identityMembershipOrg.scopeOrgId);
+    const identity = await identityDAL.findById(identityLdapAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });
+
+    const plan = await licenseService.getPlan(identity.orgId);
    if (!plan.ldap) {
      throw new BadRequestError({
        message:
@@ -179,12 +171,9 @@ export const identityLdapAuthServiceFactory = ({
    }

    const identityAccessToken = await identityLdapAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.LDAP_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.LDAP_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -218,7 +207,7 @@ export const identityLdapAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityLdapAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityLdapAuth, identityAccessToken, identity };
  };

  const attachLdapAuth = async ({
--- a/backend/src/services/identity-oci-auth/identity-oci-auth-service.ts
+++ b/backend/src/services/identity-oci-auth/identity-oci-auth-service.ts
@@ -25,6 +25,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { logger } from "@app/lib/logger";

 import { ActorType, AuthTokenType } from "../auth/auth-type";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
 import { TMembershipIdentityDALFactory } from "../membership-identity/membership-identity-dal";
@@ -41,9 +42,10 @@ import {
 } from "./identity-oci-auth-types";

 type TIdentityOciAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  identityOciAuthDAL: Pick<TIdentityOciAuthDALFactory, "findOne" | "transaction" | "create" | "updateById" | "delete">;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  orgDAL: Pick<TOrgDALFactory, "findById">;
@@ -52,6 +54,7 @@ type TIdentityOciAuthServiceFactoryDep = {
 export type TIdentityOciAuthServiceFactory = ReturnType<typeof identityOciAuthServiceFactory>;

 export const identityOciAuthServiceFactory = ({
+  identityDAL,
  identityAccessTokenDAL,
  identityOciAuthDAL,
  membershipIdentityDAL,
@@ -65,11 +68,8 @@ export const identityOciAuthServiceFactory = ({
      throw new NotFoundError({ message: "OCI auth method not found for identity, did you configure OCI auth?" });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityOciAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) throw new UnauthorizedError({ message: "Identity not attached to a organization" });
+    const identity = await identityDAL.findById(identityOciAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    // Validate OCI host format. Ensures that the host is in "identity.<region>.oraclecloud.com" format.
    if (!headers.host || !new RE2("^identity\\.([a-z]{2}-[a-z]+-[1-9])\\.oraclecloud\\.com$").test(headers.host)) {
@@ -104,12 +104,9 @@ export const identityOciAuthServiceFactory = ({

    // Generate the token
    const identityAccessToken = await identityOciAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.OCI_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.OCI_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -146,7 +143,7 @@ export const identityOciAuthServiceFactory = ({
      identityOciAuth,
      accessToken,
      identityAccessToken,
-      identityMembershipOrg
+      identity
    };
  };

--- a/backend/src/services/identity-oidc-auth/identity-oidc-auth-service.ts
+++ b/backend/src/services/identity-oidc-auth/identity-oidc-auth-service.ts
@@ -25,6 +25,7 @@ import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";
 import { getValueByDot } from "@app/lib/template/dot-access";

 import { ActorType, AuthTokenType } from "../auth/auth-type";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
 import { TKmsServiceFactory } from "../kms/kms-service";
@@ -43,8 +44,9 @@ import {
 } from "./identity-oidc-auth-types";

 type TIdentityOidcAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityOidcAuthDAL: TIdentityOidcAuthDALFactory;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
@@ -55,6 +57,7 @@ type TIdentityOidcAuthServiceFactoryDep = {
 export type TIdentityOidcAuthServiceFactory = ReturnType<typeof identityOidcAuthServiceFactory>;

 export const identityOidcAuthServiceFactory = ({
+  identityDAL,
  identityOidcAuthDAL,
  membershipIdentityDAL,
  permissionService,
@@ -69,19 +72,12 @@ export const identityOidcAuthServiceFactory = ({
      throw new NotFoundError({ message: "OIDC auth method not found for identity, did you configure OIDC auth?" });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityOidcAuth.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) {
-      throw new NotFoundError({
-        message: `Identity organization membership for identity with ID '${identityOidcAuth.identityId}' not found`
-      });
-    }
+    const identity = await identityDAL.findById(identityOidcAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const { decryptor } = await kmsService.createCipherPairWithDataKey({
      type: KmsDataKey.Organization,
-      orgId: identityMembershipOrg.scopeOrgId
+      orgId: identity.orgId
    });

    let caCert = "";
@@ -182,12 +178,9 @@ export const identityOidcAuthServiceFactory = ({
    }

    const identityAccessToken = await identityOidcAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.OIDC_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.OIDC_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -226,7 +219,7 @@ export const identityOidcAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityOidcAuth, identityAccessToken, identityMembershipOrg, oidcTokenData: tokenData };
+    return { accessToken, identityOidcAuth, identityAccessToken, identity, oidcTokenData: tokenData };
  };

  const attachOidcAuth = async ({
--- a/backend/src/services/identity-tls-cert-auth/identity-tls-cert-auth-service.ts
+++ b/backend/src/services/identity-tls-cert-auth/identity-tls-cert-auth-service.ts
@@ -21,6 +21,7 @@ import {
 import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";

 import { ActorType, AuthTokenType } from "../auth/auth-type";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
 import { TKmsServiceFactory } from "../kms/kms-service";
@@ -31,12 +32,13 @@ import { TIdentityTlsCertAuthDALFactory } from "./identity-tls-cert-auth-dal";
 import { TIdentityTlsCertAuthServiceFactory } from "./identity-tls-cert-auth-types";

 type TIdentityTlsCertAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityAccessTokenDAL: Pick<TIdentityAccessTokenDALFactory, "create" | "delete">;
  identityTlsCertAuthDAL: Pick<
    TIdentityTlsCertAuthDALFactory,
    "findOne" | "transaction" | "create" | "updateById" | "delete"
  >;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  licenseService: Pick<TLicenseServiceFactory, "getPlan">;
  permissionService: Pick<TPermissionServiceFactory, "getOrgPermission">;
  kmsService: Pick<TKmsServiceFactory, "createCipherPairWithDataKey">;
@@ -52,6 +54,7 @@ const parseSubjectDetails = (data: string) => {
 };

 export const identityTlsCertAuthServiceFactory = ({
+  identityDAL,
  identityAccessTokenDAL,
  identityTlsCertAuthDAL,
  membershipIdentityDAL,
@@ -67,20 +70,12 @@ export const identityTlsCertAuthServiceFactory = ({
      });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityTlsCertAuth.identityId,
-      scope: AccessScope.Organization
-    });
-
-    if (!identityMembershipOrg) {
-      throw new NotFoundError({
-        message: `Identity organization membership for identity with ID '${identityTlsCertAuth.identityId}' not found`
-      });
-    }
+    const identity = await identityDAL.findById(identityTlsCertAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });

    const { decryptor } = await kmsService.createCipherPairWithDataKey({
      type: KmsDataKey.Organization,
-      orgId: identityMembershipOrg.scopeOrgId
+      orgId: identity.orgId
    });

    const caCertificate = decryptor({
@@ -125,12 +120,9 @@ export const identityTlsCertAuthServiceFactory = ({

    // Generate the token
    const identityAccessToken = await identityTlsCertAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.TLS_CERT_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.TLS_CERT_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -167,7 +159,7 @@ export const identityTlsCertAuthServiceFactory = ({
      identityTlsCertAuth,
      accessToken,
      identityAccessToken,
-      identityMembershipOrg
+      identity
    };
  };

--- a/backend/src/services/identity-tls-cert-auth/identity-tls-cert-auth-types.ts
+++ b/backend/src/services/identity-tls-cert-auth/identity-tls-cert-auth-types.ts
@@ -1,4 +1,4 @@
-import { TIdentityAccessTokens, TIdentityTlsCertAuths, TMemberships } from "@app/db/schemas";
+import { TIdentities, TIdentityAccessTokens, TIdentityTlsCertAuths } from "@app/db/schemas";
 import { TProjectPermission } from "@app/lib/types";

 export type TLoginTlsCertAuthDTO = {
@@ -40,7 +40,7 @@ export type TIdentityTlsCertAuthServiceFactory = {
    identityTlsCertAuth: TIdentityTlsCertAuths;
    accessToken: string;
    identityAccessToken: TIdentityAccessTokens;
-    identityMembershipOrg: TMemberships;
+    identity: TIdentities;
  }>;
  attachTlsCertAuth: (dto: TAttachTlsCertAuthDTO) => Promise<TIdentityTlsCertAuths>;
  updateTlsCertAuth: (dto: TUpdateTlsCertAuthDTO) => Promise<TIdentityTlsCertAuths>;
--- a/backend/src/services/identity-token-auth/identity-token-auth-service.ts
+++ b/backend/src/services/identity-token-auth/identity-token-auth-service.ts
@@ -10,10 +10,17 @@ import {
 import { TPermissionServiceFactory } from "@app/ee/services/permission/permission-service-types";
 import { getConfig } from "@app/lib/config/env";
 import { crypto } from "@app/lib/crypto";
-import { BadRequestError, ForbiddenRequestError, NotFoundError, PermissionBoundaryError } from "@app/lib/errors";
+import {
+  BadRequestError,
+  ForbiddenRequestError,
+  NotFoundError,
+  PermissionBoundaryError,
+  UnauthorizedError
+} from "@app/lib/errors";
 import { extractIPDetails, isValidIpOrCidr } from "@app/lib/ip";

 import { ActorType, AuthTokenType } from "../auth/auth-type";
+import { TIdentityDALFactory } from "../identity/identity-dal";
 import { TIdentityAccessTokenDALFactory } from "../identity-access-token/identity-access-token-dal";
 import { TIdentityAccessTokenJwtPayload } from "../identity-access-token/identity-access-token-types";
 import { TMembershipIdentityDALFactory } from "../membership-identity/membership-identity-dal";
@@ -32,11 +39,12 @@ import {
 } from "./identity-token-auth-types";

 type TIdentityTokenAuthServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityTokenAuthDAL: Pick<
    TIdentityTokenAuthDALFactory,
    "transaction" | "create" | "findOne" | "updateById" | "delete"
  >;
-  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "updateById" | "getIdentityById">;
+  membershipIdentityDAL: Pick<TMembershipIdentityDALFactory, "findOne" | "update" | "getIdentityById">;
  identityAccessTokenDAL: Pick<
    TIdentityAccessTokenDALFactory,
    "create" | "find" | "update" | "findById" | "findOne" | "updateById" | "delete"
@@ -49,8 +57,8 @@ type TIdentityTokenAuthServiceFactoryDep = {
 export type TIdentityTokenAuthServiceFactory = ReturnType<typeof identityTokenAuthServiceFactory>;

 export const identityTokenAuthServiceFactory = ({
+  identityDAL,
  identityTokenAuthDAL,
-  // identityDAL,
  membershipIdentityDAL,
  identityAccessTokenDAL,
  permissionService,
@@ -400,13 +408,13 @@ export const identityTokenAuthServiceFactory = ({

    const identityTokenAuth = await identityTokenAuthDAL.findOne({ identityId });

+    const identity = await identityDAL.findById(identityTokenAuth.identityId);
+    if (!identity) throw new UnauthorizedError({ message: "Identity not found" });
+
    const identityAccessToken = await identityTokenAuthDAL.transaction(async (tx) => {
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
-        {
-          lastLoginAuthMethod: IdentityAuthMethod.TOKEN_AUTH,
-          lastLoginTime: new Date()
-        },
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
+        { lastLoginAuthMethod: IdentityAuthMethod.TOKEN_AUTH, lastLoginTime: new Date() },
        tx
      );
      const newToken = await identityAccessTokenDAL.create(
@@ -441,7 +449,7 @@ export const identityTokenAuthServiceFactory = ({
          }
    );

-    return { accessToken, identityTokenAuth, identityAccessToken, identityMembershipOrg };
+    return { accessToken, identityTokenAuth, identityAccessToken, identity };
  };

  const getTokenAuthTokens = async ({
--- a/backend/src/services/identity-ua/identity-ua-service.ts
+++ b/backend/src/services/identity-ua/identity-ua-service.ts
@@ -41,8 +41,10 @@ import {
  TRevokeUaDTO,
  TUpdateUaDTO
 } from "./identity-ua-types";
+import { TIdentityDALFactory } from "../identity/identity-dal";

 type TIdentityUaServiceFactoryDep = {
+  identityDAL: Pick<TIdentityDALFactory, "findById">;
  identityUaDAL: TIdentityUaDALFactory;
  identityUaClientSecretDAL: TIdentityUaClientSecretDALFactory;
  identityAccessTokenDAL: TIdentityAccessTokenDALFactory;
@@ -71,7 +73,8 @@ export const identityUaServiceFactory = ({
  permissionService,
  licenseService,
  orgDAL,
-  keyStore
+  keyStore,
+  identityDAL
 }: TIdentityUaServiceFactoryDep) => {
  const login = async (clientId: string, clientSecret: string, ip: string) => {
    const identityUa = await identityUaDAL.findOne({ clientId });
@@ -101,16 +104,6 @@ export const identityUaServiceFactory = ({
      });
    }

-    const identityMembershipOrg = await membershipIdentityDAL.findOne({
-      actorIdentityId: identityUa.identityId,
-      scope: AccessScope.Organization
-    });
-    if (!identityMembershipOrg) {
-      throw new UnauthorizedError({
-        message: "Invalid credentials"
-      });
-    }
-
    const clientSecretPrefix = clientSecret.slice(0, 4);
    const clientSecretInfo = await identityUaClientSecretDAL.find({
      identityUAId: identityUa.id,
@@ -228,10 +221,11 @@ export const identityUaServiceFactory = ({
            accessTokenMaxTTL: 1000000000
          };

+    const identity = await identityDAL.findById(identityUa.identityId);
    const identityAccessToken = await identityUaDAL.transaction(async (tx) => {
      const uaClientSecretDoc = await identityUaClientSecretDAL.incrementUsage(validClientSecretInfo!.id, tx);
-      await membershipIdentityDAL.updateById(
-        identityMembershipOrg.id,
+      await membershipIdentityDAL.update(
+        { scope: AccessScope.Organization, scopeOrgId: identity.orgId, actorIdentityId: identity.id },
        {
          lastLoginAuthMethod: IdentityAuthMethod.UNIVERSAL_AUTH,
          lastLoginTime: new Date()
@@ -277,7 +271,7 @@ export const identityUaServiceFactory = ({
      identityUa,
      validClientSecretInfo,
      identityAccessToken,
-      identityMembershipOrg,
+      identity,
      ...accessTokenTTLParams
    };
  };