From 4d6f8a612817b2dc6d476f36f508cadc60aeeed5 Mon Sep 17 00:00:00 2001
From: Sheen Capadngan <sheensantoscapadngan@gmail.com>
Date: Fri, 3 Oct 2025 23:37:15 +0800
Subject: [PATCH] misc: more security improvements

---
 .../src/ee/routes/v1/pam-session-router.ts    |  2 +-
 .../services/pam-session/pam-session-dal.ts   | 21 ++++++++++-
 .../pam-session/pam-session-service.ts        | 37 +++++++++++++------
 3 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/backend/src/ee/routes/v1/pam-session-router.ts b/backend/src/ee/routes/v1/pam-session-router.ts
index 6e60704da..c353fddfa 100644
--- a/backend/src/ee/routes/v1/pam-session-router.ts
+++ b/backend/src/ee/routes/v1/pam-session-router.ts
@@ -124,7 +124,7 @@ export const registerPamSessionRouter = async (server: FastifyZodProvider) => {
         })
       }
     },
-    onRequest: verifyAuth([AuthMode.IDENTITY_ACCESS_TOKEN]),
+    onRequest: verifyAuth([AuthMode.JWT, AuthMode.IDENTITY_ACCESS_TOKEN]),
     handler: async (req) => {
       const { session, projectId } = await server.services.pamSession.endSessionById(
         req.params.sessionId,
diff --git a/backend/src/ee/services/pam-session/pam-session-dal.ts b/backend/src/ee/services/pam-session/pam-session-dal.ts
index 6957b1fde..f8b3a3393 100644
--- a/backend/src/ee/services/pam-session/pam-session-dal.ts
+++ b/backend/src/ee/services/pam-session/pam-session-dal.ts
@@ -1,9 +1,26 @@
+import { Knex } from "knex";
+
 import { TDbClient } from "@app/db";
 import { TableName } from "@app/db/schemas";
-import { ormify } from "@app/lib/knex";
+import { ormify, selectAllTableCols } from "@app/lib/knex";
 
 export type TPamSessionDALFactory = ReturnType<typeof pamSessionDALFactory>;
 export const pamSessionDALFactory = (db: TDbClient) => {
   const orm = ormify(db, TableName.PamSession);
-  return { ...orm };
+
+  const findById = async (id: string, tx?: Knex) => {
+    const session = await (tx || db.replicaNode())(TableName.PamSession)
+      .leftJoin(TableName.PamAccount, `${TableName.PamSession}.accountId`, `${TableName.PamAccount}.id`)
+      .leftJoin(TableName.PamResource, `${TableName.PamAccount}.resourceId`, `${TableName.PamResource}.id`)
+      .leftJoin(TableName.GatewayV2, `${TableName.PamResource}.gatewayId`, `${TableName.GatewayV2}.id`)
+      .select(selectAllTableCols(TableName.PamSession))
+      .select(db.ref("name").withSchema(TableName.GatewayV2).as("gatewayName"))
+      .select(db.ref("identityId").withSchema(TableName.GatewayV2).as("gatewayIdentityId"))
+      .where(`${TableName.PamSession}.id`, id)
+      .first();
+
+    return session;
+  };
+
+  return { ...orm, findById };
 };
diff --git a/backend/src/ee/services/pam-session/pam-session-service.ts b/backend/src/ee/services/pam-session/pam-session-service.ts
index 07185bec2..c17144a3d 100644
--- a/backend/src/ee/services/pam-session/pam-session-service.ts
+++ b/backend/src/ee/services/pam-session/pam-session-service.ts
@@ -95,6 +95,10 @@ export const pamSessionServiceFactory = ({
     const session = await pamSessionDAL.findById(sessionId);
     if (!session) throw new NotFoundError({ message: `Session with ID '${sessionId}' not found` });
 
+    if (session.status !== PamSessionStatus.Active) {
+      throw new BadRequestError({ message: "Cannot update logs for sessions that are not active" });
+    }
+
     const project = await projectDAL.findById(session.projectId);
     if (!project) throw new NotFoundError({ message: `Project with ID '${session.projectId}' not found` });
 
@@ -111,6 +115,10 @@ export const pamSessionServiceFactory = ({
       OrgPermissionSubjects.Gateway
     );
 
+    if (session.gatewayIdentityId !== actor.id) {
+      throw new ForbiddenRequestError({ message: "Identity does not have access to update logs for this session" });
+    }
+
     const { encryptor } = await kmsService.createCipherPairWithDataKey({
       type: KmsDataKey.SecretManager,
       projectId: session.projectId
@@ -128,11 +136,6 @@ export const pamSessionServiceFactory = ({
   };
 
   const endSessionById = async (sessionId: string, actor: OrgServiceActor) => {
-    // To be hit by gateways only
-    if (actor.type !== ActorType.IDENTITY) {
-      throw new ForbiddenRequestError({ message: "Only gateways can perform this action" });
-    }
-
     const session = await pamSessionDAL.findById(sessionId);
     if (!session) throw new NotFoundError({ message: `Session with ID '${sessionId}' not found` });
 
@@ -147,10 +150,22 @@ export const pamSessionServiceFactory = ({
       actor.orgId
     );
 
-    ForbiddenError.from(permission).throwUnlessCan(
-      OrgPermissionGatewayActions.CreateGateways,
-      OrgPermissionSubjects.Gateway
-    );
+    if (actor.type === ActorType.IDENTITY) {
+      ForbiddenError.from(permission).throwUnlessCan(
+        OrgPermissionGatewayActions.CreateGateways,
+        OrgPermissionSubjects.Gateway
+      );
+
+      if (session.gatewayIdentityId !== actor.id) {
+        throw new ForbiddenRequestError({ message: "Identity does not have access to end this session" });
+      }
+    } else if (actor.type === ActorType.USER) {
+      if (session.userId !== actor.id) {
+        throw new ForbiddenRequestError({ message: "You are not authorized to end this session" });
+      }
+    } else {
+      throw new ForbiddenRequestError({ message: "Only identities and users can perform this action" });
+    }
 
     if (session.status === PamSessionStatus.Ended) {
       return {
@@ -159,8 +174,8 @@ export const pamSessionServiceFactory = ({
       };
     }
 
-    if (session.status !== PamSessionStatus.Active) {
-      throw new BadRequestError({ message: "Cannot end sessions that are not active" });
+    if (session.status !== PamSessionStatus.Active && session.status !== PamSessionStatus.Starting) {
+      throw new BadRequestError({ message: "Cannot end sessions that are not active or starting" });
     }
 
     const updatedSession = await pamSessionDAL.updateById(sessionId, {