From e6c0bbb25b1a64c4ab2d91fd9f78afc18f1e0fad Mon Sep 17 00:00:00 2001
From: Daniel Hougaard <daniel@infisical.com>
Date: Mon, 9 Sep 2024 23:15:58 +0400
Subject: [PATCH] fix: stale session after logging into CLI

---
 backend/src/server/routes/index.ts               |  2 +-
 backend/src/server/routes/v3/login-router.ts     |  5 +++--
 backend/src/services/auth/auth-login-service.ts  |  7 +------
 frontend/src/hooks/api/auth/queries.tsx          | 14 ++++++++++----
 frontend/src/pages/login/select-organization.tsx |  5 ++++-
 5 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/backend/src/server/routes/index.ts b/backend/src/server/routes/index.ts
index f33456bd5..762e22e54 100644
--- a/backend/src/server/routes/index.ts
+++ b/backend/src/server/routes/index.ts
@@ -468,7 +468,7 @@ export const registerRoutes = async (
     projectMembershipDAL
   });
 
-  const loginService = authLoginServiceFactory({ userDAL, smtpService, tokenService, orgDAL, tokenDAL: authTokenDAL });
+  const loginService = authLoginServiceFactory({ userDAL, smtpService, tokenService, orgDAL });
   const passwordService = authPaswordServiceFactory({
     tokenService,
     smtpService,
diff --git a/backend/src/server/routes/v3/login-router.ts b/backend/src/server/routes/v3/login-router.ts
index 61a0c74e5..dfacae8a8 100644
--- a/backend/src/server/routes/v3/login-router.ts
+++ b/backend/src/server/routes/v3/login-router.ts
@@ -42,7 +42,8 @@ export const registerLoginRouter = async (server: FastifyZodProvider) => {
     },
     schema: {
       body: z.object({
-        organizationId: z.string().trim()
+        organizationId: z.string().trim(),
+        customUserAgent: z.enum(["cli"]).optional()
       }),
       response: {
         200: z.object({
@@ -53,7 +54,7 @@ export const registerLoginRouter = async (server: FastifyZodProvider) => {
     handler: async (req, res) => {
       const cfg = getConfig();
       const tokens = await server.services.login.selectOrganization({
-        userAgent: req.headers["user-agent"],
+        userAgent: req.body.customUserAgent ?? req.headers["user-agent"],
         authJwtToken: req.headers.authorization,
         organizationId: req.body.organizationId,
         ipAddress: req.realIp
diff --git a/backend/src/services/auth/auth-login-service.ts b/backend/src/services/auth/auth-login-service.ts
index 91bf40198..fb38c024e 100644
--- a/backend/src/services/auth/auth-login-service.ts
+++ b/backend/src/services/auth/auth-login-service.ts
@@ -12,7 +12,6 @@ import { BadRequestError, DatabaseError, UnauthorizedError } from "@app/lib/erro
 import { logger } from "@app/lib/logger";
 import { getServerCfg } from "@app/services/super-admin/super-admin-service";
 
-import { TTokenDALFactory } from "../auth-token/auth-token-dal";
 import { TAuthTokenServiceFactory } from "../auth-token/auth-token-service";
 import { TokenType } from "../auth-token/auth-token-types";
 import { TOrgDALFactory } from "../org/org-dal";
@@ -34,7 +33,6 @@ type TAuthLoginServiceFactoryDep = {
   orgDAL: TOrgDALFactory;
   tokenService: TAuthTokenServiceFactory;
   smtpService: TSmtpService;
-  tokenDAL: TTokenDALFactory;
 };
 
 export type TAuthLoginFactory = ReturnType<typeof authLoginServiceFactory>;
@@ -42,8 +40,7 @@ export const authLoginServiceFactory = ({
   userDAL,
   tokenService,
   smtpService,
-  orgDAL,
-  tokenDAL
+  orgDAL
 }: TAuthLoginServiceFactoryDep) => {
   /*
    * Private
@@ -376,8 +373,6 @@ export const authLoginServiceFactory = ({
       });
     }
 
-    await tokenDAL.incrementTokenSessionVersion(user.id, decodedToken.tokenVersionId);
-
     const tokens = await generateUserTokens({
       authMethod: decodedToken.authMethod,
       user,
diff --git a/frontend/src/hooks/api/auth/queries.tsx b/frontend/src/hooks/api/auth/queries.tsx
index fcec4f2ef..411aebae8 100644
--- a/frontend/src/hooks/api/auth/queries.tsx
+++ b/frontend/src/hooks/api/auth/queries.tsx
@@ -60,7 +60,10 @@ export const useLogin1 = () => {
   });
 };
 
-export const selectOrganization = async (data: { organizationId: string }) => {
+export const selectOrganization = async (data: {
+  organizationId: string;
+  customUserAgent?: string;
+}) => {
   const { data: res } = await apiRequest.post<{ token: string }>(
     "/api/v3/auth/select-organization",
     data
@@ -71,11 +74,14 @@ export const selectOrganization = async (data: { organizationId: string }) => {
 export const useSelectOrganization = () => {
   const queryClient = useQueryClient();
   return useMutation({
-    mutationFn: async (details: { organizationId: string }) => {
+    mutationFn: async (details: { organizationId: string; customUserAgent?: string }) => {
       const data = await selectOrganization(details);
 
-      SecurityClient.setToken(data.token);
-      SecurityClient.setProviderAuthToken("");
+      // If a custom user agent is set, then this session is meant for another consuming application, not the web application.
+      if (!details.customUserAgent) {
+        SecurityClient.setToken(data.token);
+        SecurityClient.setProviderAuthToken("");
+      }
 
       return data;
     },
diff --git a/frontend/src/pages/login/select-organization.tsx b/frontend/src/pages/login/select-organization.tsx
index cf953664d..32081af9f 100644
--- a/frontend/src/pages/login/select-organization.tsx
+++ b/frontend/src/pages/login/select-organization.tsx
@@ -68,7 +68,10 @@ export default function LoginPage() {
         return;
       }
 
-      const { token } = await selectOrg.mutateAsync({ organizationId: organization.id });
+      const { token } = await selectOrg.mutateAsync({
+        organizationId: organization.id,
+        customUserAgent: callbackPort ? "cli" : undefined
+      });
 
       if (callbackPort) {
         const privateKey = localStorage.getItem("PRIVATE_KEY");