From bdcc2386185d1d44da42114ce19ce66827b5d57e Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Fri, 1 Aug 2025 11:38:01 -0300 Subject: [PATCH] refactor: Improve cookie security and centralized utility (#9240) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ✨ (authContext.tsx): Add setCookieWithOptions function to set cookies with specific options for better security and control 📝 (use-post-refresh-access.ts): Update cookies.set calls to use setCookieWithOptions function for consistent cookie settings ♻️ (utils.ts): Refactor setCookieWithOptions function to include httpOnly option and update sameSite values to lowercase for consistency * 📝 (frontend): add useGetCookieAuth and useSetCookieAuth hooks for managing cookies in auth context 🔧 (frontend): refactor authStore to use new cookie hooks for managing access token and api key cookies 🔧 (frontend): refactor use-post-logout and use-post-refresh-access to use new cookie hooks for cookie management * 📝 (frontend): Remove redundant useGetCookieAuth hook and use direct access to cookies in authStore and related components 🔧 (utils): Refactor getAuthCookie and setAuthCookie functions to use react-cookie directly for better code organization and readability --- src/frontend/src/contexts/authContext.tsx | 15 +- .../auth/__tests__/use-post-logout.test.ts | 167 ++++++++++++++ .../__tests__/use-post-refresh-access.test.ts | 150 +++++++++++++ .../API/queries/auth/use-post-logout.ts | 3 +- .../queries/auth/use-post-refresh-access.ts | 5 +- .../__tests__/use-get-cookie-auth.test.ts | 109 +++++++++ .../__tests__/use-set-cookie-auth.test.ts | 130 +++++++++++ .../src/stores/__tests__/authStore.test.ts | 206 ++++++++++++++++++ src/frontend/src/stores/authStore.ts | 7 +- src/frontend/src/utils/utils.ts | 24 +- 10 files changed, 802 insertions(+), 14 deletions(-) create mode 100644 src/frontend/src/controllers/API/queries/auth/__tests__/use-post-logout.test.ts create mode 100644 src/frontend/src/controllers/API/queries/auth/__tests__/use-post-refresh-access.test.ts create mode 100644 src/frontend/src/shared/hooks/__tests__/use-get-cookie-auth.test.ts create mode 100644 src/frontend/src/shared/hooks/__tests__/use-set-cookie-auth.test.ts create mode 100644 src/frontend/src/stores/__tests__/authStore.test.ts diff --git a/src/frontend/src/contexts/authContext.tsx b/src/frontend/src/contexts/authContext.tsx index fbe7372d9..81d893255 100644 --- a/src/frontend/src/contexts/authContext.tsx +++ b/src/frontend/src/contexts/authContext.tsx @@ -10,6 +10,7 @@ import { useGetUserData } from "@/controllers/API/queries/auth"; import { useGetGlobalVariablesMutation } from "@/controllers/API/queries/variables/use-get-mutation-global-variables"; import useAuthStore from "@/stores/authStore"; import { setLocalStorage } from "@/utils/local-storage-util"; +import { getAuthCookie, setAuthCookie } from "@/utils/utils"; import { useStoreStore } from "../stores/storeStore"; import type { Users } from "../types/api"; import type { AuthContextType } from "../types/contexts/auth"; @@ -31,11 +32,11 @@ export const AuthContext = createContext(initialValue); export function AuthProvider({ children }): React.ReactElement { const cookies = new Cookies(); const [accessToken, setAccessToken] = useState( - cookies.get(LANGFLOW_ACCESS_TOKEN) ?? null, + getAuthCookie(cookies, LANGFLOW_ACCESS_TOKEN) ?? null, ); const [userData, setUserData] = useState(null); const [apiKey, setApiKey] = useState( - cookies.get(LANGFLOW_API_TOKEN), + getAuthCookie(cookies, LANGFLOW_API_TOKEN), ); const checkHasStore = useStoreStore((state) => state.checkHasStore); @@ -46,14 +47,14 @@ export function AuthProvider({ children }): React.ReactElement { const { mutate: mutateGetGlobalVariables } = useGetGlobalVariablesMutation(); useEffect(() => { - const storedAccessToken = cookies.get(LANGFLOW_ACCESS_TOKEN); + const storedAccessToken = getAuthCookie(cookies, LANGFLOW_ACCESS_TOKEN); if (storedAccessToken) { setAccessToken(storedAccessToken); } }, []); useEffect(() => { - const apiKey = cookies.get(LANGFLOW_API_TOKEN); + const apiKey = getAuthCookie(cookies, LANGFLOW_API_TOKEN); if (apiKey) { setApiKey(apiKey); } @@ -82,12 +83,12 @@ export function AuthProvider({ children }): React.ReactElement { autoLogin: string, refreshToken?: string, ) { - cookies.set(LANGFLOW_ACCESS_TOKEN, newAccessToken, { path: "/" }); - cookies.set(LANGFLOW_AUTO_LOGIN_OPTION, autoLogin, { path: "/" }); + setAuthCookie(cookies, LANGFLOW_ACCESS_TOKEN, newAccessToken); + setAuthCookie(cookies, LANGFLOW_AUTO_LOGIN_OPTION, autoLogin); setLocalStorage(LANGFLOW_ACCESS_TOKEN, newAccessToken); if (refreshToken) { - cookies.set(LANGFLOW_REFRESH_TOKEN, refreshToken, { path: "/" }); + setAuthCookie(cookies, LANGFLOW_REFRESH_TOKEN, refreshToken); } setAccessToken(newAccessToken); setIsAuthenticated(true); diff --git a/src/frontend/src/controllers/API/queries/auth/__tests__/use-post-logout.test.ts b/src/frontend/src/controllers/API/queries/auth/__tests__/use-post-logout.test.ts new file mode 100644 index 000000000..35b98fbf5 --- /dev/null +++ b/src/frontend/src/controllers/API/queries/auth/__tests__/use-post-logout.test.ts @@ -0,0 +1,167 @@ +// Logout functionality tests + +// Mock all dependencies before imports +const mockLogout = jest.fn(); +const mockResetFlowState = jest.fn(); +const mockResetFlowsManagerStore = jest.fn(); +const mockResetFolderStore = jest.fn(); +const mockQueryClient = { invalidateQueries: jest.fn() }; +const mockGetAuthCookie = jest.fn(); +const mockApiPost = jest.fn(); + +jest.mock("@/stores/authStore", () => { + const mockState = { autoLogin: false }; + const mockStore = jest.fn((selector) => { + if (selector.toString().includes("logout")) return mockLogout; + return false; + }); + mockStore.getState = jest.fn(() => mockState); + return mockStore; +}); + +jest.mock("@/stores/flowStore", () => { + const mockStore = jest.fn(); + mockStore.getState = jest.fn(() => ({ resetFlowState: mockResetFlowState })); + return mockStore; +}); + +jest.mock("@/stores/flowsManagerStore", () => { + const mockStore = jest.fn(); + mockStore.getState = jest.fn(() => ({ + resetStore: mockResetFlowsManagerStore, + })); + return mockStore; +}); + +jest.mock("@/stores/foldersStore", () => ({ + useFolderStore: { + getState: jest.fn(() => ({ resetStore: mockResetFolderStore })), + }, +})); + +jest.mock("@/utils/utils", () => ({ + getAuthCookie: mockGetAuthCookie, +})); + +jest.mock("@/controllers/API/api", () => ({ + api: { + post: mockApiPost, + }, +})); + +jest.mock("@/controllers/API/services/request-processor", () => ({ + UseRequestProcessor: jest.fn(() => ({ + mutate: jest.fn((key, fn, options) => ({ + mutate: async () => { + try { + await fn(); + if (options?.onSuccess) options.onSuccess(); + } catch (error) { + if (options?.onError) options.onError(error); + throw error; + } + }, + })), + queryClient: mockQueryClient, + })), +})); + +jest.mock("react-cookie", () => ({ + Cookies: jest.fn().mockImplementation(() => ({})), +})); + +jest.mock("@/constants/constants", () => ({ + ...jest.requireActual("@/constants/constants"), + IS_AUTO_LOGIN: false, // Override to disable auto login for testing + LANGFLOW_AUTO_LOGIN_OPTION: "auto_login_lf", +})); + +jest.mock("@/controllers/API/helpers/constants", () => ({ + getURL: jest.fn((key) => `/api/v1/${key.toLowerCase()}`), +})); + +import { useLogout } from "../use-post-logout"; + +describe("logout functionality", () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetAuthCookie.mockReturnValue(null); + }); + + describe("logout behavior with auto login disabled", () => { + it("should call API logout when auto login is disabled", async () => { + mockGetAuthCookie.mockReturnValue(null); // Not "auto", so autoLogin is false + mockApiPost.mockResolvedValue({ data: { success: true } }); + + const logoutMutation = useLogout(); + await logoutMutation.mutate(); + + expect(mockApiPost).toHaveBeenCalledWith( + expect.stringContaining("logout"), + ); + }); + + it("should reset all stores on successful logout", async () => { + mockGetAuthCookie.mockReturnValue(null); // Not "auto", so autoLogin is false + mockApiPost.mockResolvedValue({ data: { success: true } }); + + const logoutMutation = useLogout(); + await logoutMutation.mutate(); + + expect(mockLogout).toHaveBeenCalled(); + expect(mockResetFlowState).toHaveBeenCalled(); + expect(mockResetFlowsManagerStore).toHaveBeenCalled(); + expect(mockResetFolderStore).toHaveBeenCalled(); + }); + + it("should invalidate queries on successful logout", async () => { + mockGetAuthCookie.mockReturnValue(null); // Not "auto", so autoLogin is false + mockApiPost.mockResolvedValue({ data: { success: true } }); + + const logoutMutation = useLogout(); + await logoutMutation.mutate(); + + expect(mockQueryClient.invalidateQueries).toHaveBeenCalledWith({ + queryKey: ["useGetRefreshFlowsQuery"], + }); + expect(mockQueryClient.invalidateQueries).toHaveBeenCalledWith({ + queryKey: ["useGetFolders"], + }); + expect(mockQueryClient.invalidateQueries).toHaveBeenCalledWith({ + queryKey: ["useGetFolder"], + }); + }); + }); + + describe("logout behavior with auto login enabled", () => { + it("should skip API call when auto login is enabled via cookie", async () => { + mockGetAuthCookie.mockReturnValue("auto"); + + const logoutMutation = useLogout(); + await logoutMutation.mutate(); + + expect(mockApiPost).not.toHaveBeenCalled(); + }); + + it("should still reset stores even when skipping API call", async () => { + mockGetAuthCookie.mockReturnValue("auto"); + + const logoutMutation = useLogout(); + await logoutMutation.mutate(); + + expect(mockLogout).toHaveBeenCalled(); + expect(mockResetFlowState).toHaveBeenCalled(); + }); + }); + + describe("error handling", () => { + it("should handle API errors gracefully", async () => { + mockGetAuthCookie.mockReturnValue(null); // Not "auto", so autoLogin is false + const mockError = new Error("API Error"); + mockApiPost.mockRejectedValue(mockError); + + const logoutMutation = useLogout(); + await expect(logoutMutation.mutate()).rejects.toThrow("API Error"); + }); + }); +}); diff --git a/src/frontend/src/controllers/API/queries/auth/__tests__/use-post-refresh-access.test.ts b/src/frontend/src/controllers/API/queries/auth/__tests__/use-post-refresh-access.test.ts new file mode 100644 index 000000000..c2cdeef1f --- /dev/null +++ b/src/frontend/src/controllers/API/queries/auth/__tests__/use-post-refresh-access.test.ts @@ -0,0 +1,150 @@ +// Refresh token functionality tests + +// Mock all dependencies before imports +jest.mock("@/utils/utils", () => ({ + setAuthCookie: jest.fn(), +})); + +jest.mock( + "@/stores/authStore", + () => jest.fn((selector) => false), // autoLogin = false +); + +jest.mock("@/controllers/API/api", () => ({ + api: { + post: jest.fn(), + }, +})); + +jest.mock("@/controllers/API/services/request-processor", () => ({ + UseRequestProcessor: jest.fn(() => ({ + mutate: jest.fn((key, fn, options) => ({ + mutate: async () => { + return await fn(); + }, + })), + })), +})); + +jest.mock("react-cookie", () => ({ + Cookies: jest.fn().mockImplementation(() => ({})), +})); + +import { useRefreshAccessToken } from "../use-post-refresh-access"; + +const mockSetAuthCookie = require("@/utils/utils").setAuthCookie; +const mockApiPost = require("@/controllers/API/api").api.post; + +describe("refresh token functionality", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe("successful token refresh", () => { + it("should call refresh API and set new refresh token cookie", async () => { + const mockRefreshResponse = { + access_token: "new-access-token", + refresh_token: "new-refresh-token", + token_type: "bearer", + }; + + mockApiPost.mockResolvedValue({ data: mockRefreshResponse }); + + const refreshMutation = useRefreshAccessToken(); + const result = await refreshMutation.mutate(); + + expect(mockApiPost).toHaveBeenCalledWith( + expect.stringContaining("refresh"), + ); + expect(mockSetAuthCookie).toHaveBeenCalledWith( + expect.any(Object), // cookies instance + "refresh_token_lf", + "new-refresh-token", + ); + expect(result).toEqual(mockRefreshResponse); + }); + + it("should return the refresh response data", async () => { + const mockRefreshResponse = { + access_token: "new-access-token-123", + refresh_token: "new-refresh-token-456", + token_type: "bearer", + }; + + mockApiPost.mockResolvedValue({ data: mockRefreshResponse }); + + const refreshMutation = useRefreshAccessToken(); + const result = await refreshMutation.mutate(); + + expect(result).toEqual(mockRefreshResponse); + }); + }); + + describe("error handling", () => { + it("should throw error when refresh API fails", async () => { + const mockError = new Error("Refresh failed"); + mockApiPost.mockRejectedValue(mockError); + + const refreshMutation = useRefreshAccessToken(); + await expect(refreshMutation.mutate()).rejects.toThrow("Refresh failed"); + }); + + it("should not set cookie when API fails", async () => { + const mockError = new Error("API Error"); + mockApiPost.mockRejectedValue(mockError); + + const refreshMutation = useRefreshAccessToken(); + + try { + await refreshMutation.mutate(); + } catch (error) { + // Expected to throw + } + + expect(mockSetAuthCookie).not.toHaveBeenCalled(); + }); + }); + + describe("cookie management", () => { + it("should use useSetCookieAuth hook for setting refresh token", async () => { + const mockRefreshResponse = { + access_token: "access-token", + refresh_token: "refresh-token-xyz", + token_type: "bearer", + }; + + mockApiPost.mockResolvedValue({ data: mockRefreshResponse }); + + const refreshMutation = useRefreshAccessToken(); + await refreshMutation.mutate(); + + expect(mockSetAuthCookie).toHaveBeenCalledTimes(1); + expect(mockSetAuthCookie).toHaveBeenCalledWith( + expect.any(Object), // cookies instance + "refresh_token_lf", + "refresh-token-xyz", + ); + }); + + it("should set refresh token cookie before returning response", async () => { + const mockRefreshResponse = { + access_token: "access-token", + refresh_token: "refresh-token-abc", + token_type: "bearer", + }; + + mockApiPost.mockResolvedValue({ data: mockRefreshResponse }); + + const refreshMutation = useRefreshAccessToken(); + const response = await refreshMutation.mutate(); + + // Verify cookie was set before response was returned + expect(mockSetAuthCookie).toHaveBeenCalledWith( + expect.any(Object), // cookies instance + "refresh_token_lf", + "refresh-token-abc", + ); + expect(response).toEqual(mockRefreshResponse); + }); + }); +}); diff --git a/src/frontend/src/controllers/API/queries/auth/use-post-logout.ts b/src/frontend/src/controllers/API/queries/auth/use-post-logout.ts index 0b5ce0014..d63a7749a 100644 --- a/src/frontend/src/controllers/API/queries/auth/use-post-logout.ts +++ b/src/frontend/src/controllers/API/queries/auth/use-post-logout.ts @@ -8,6 +8,7 @@ import useFlowStore from "@/stores/flowStore"; import useFlowsManagerStore from "@/stores/flowsManagerStore"; import { useFolderStore } from "@/stores/foldersStore"; import type { useMutationFunctionType } from "@/types/api"; +import { getAuthCookie } from "@/utils/utils"; import { api } from "../../api"; import { getURL } from "../../helpers/constants"; import { UseRequestProcessor } from "../../services/request-processor"; @@ -23,7 +24,7 @@ export const useLogout: useMutationFunctionType = ( async function logoutUser(): Promise { const autoLogin = useAuthStore.getState().autoLogin || - cookies.get(LANGFLOW_AUTO_LOGIN_OPTION) === "auto" || + getAuthCookie(cookies, LANGFLOW_AUTO_LOGIN_OPTION) === "auto" || isAutoLoginEnv; if (autoLogin) { diff --git a/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts b/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts index 01843d01c..ef7e560d8 100644 --- a/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts +++ b/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts @@ -2,6 +2,7 @@ import { Cookies } from "react-cookie"; import { IS_AUTO_LOGIN, LANGFLOW_REFRESH_TOKEN } from "@/constants/constants"; import useAuthStore from "@/stores/authStore"; import type { useMutationFunctionType } from "@/types/api"; +import { setAuthCookie } from "@/utils/utils"; import { api } from "../../api"; import { getURL } from "../../helpers/constants"; import { UseRequestProcessor } from "../../services/request-processor"; @@ -18,12 +19,12 @@ export const useRefreshAccessToken: useMutationFunctionType< IRefreshAccessToken > = (options?) => { const { mutate } = UseRequestProcessor(); - const cookies = new Cookies(); const autoLogin = useAuthStore((state) => state.autoLogin); async function refreshAccess(): Promise { const res = await api.post(`${getURL("REFRESH")}`); - cookies.set(LANGFLOW_REFRESH_TOKEN, res.data.refresh_token, { path: "/" }); + const cookies = new Cookies(); + setAuthCookie(cookies, LANGFLOW_REFRESH_TOKEN, res.data.refresh_token); return res.data; } diff --git a/src/frontend/src/shared/hooks/__tests__/use-get-cookie-auth.test.ts b/src/frontend/src/shared/hooks/__tests__/use-get-cookie-auth.test.ts new file mode 100644 index 000000000..5e72246bc --- /dev/null +++ b/src/frontend/src/shared/hooks/__tests__/use-get-cookie-auth.test.ts @@ -0,0 +1,109 @@ +import { Cookies } from "react-cookie"; + +// Mock all complex dependencies to avoid import issues +jest.mock("@/stores/authStore", () => ({ + __esModule: true, + default: jest.fn(() => ({})), +})); + +jest.mock("@/stores/darkStore", () => ({ + useDarkStore: { + getState: () => ({ refreshStars: jest.fn() }), + setState: jest.fn(), + subscribe: jest.fn(), + destroy: jest.fn(), + }, +})); + +jest.mock("@/stores/alertStore", () => ({ + __esModule: true, + default: jest.fn(() => ({})), +})); + +jest.mock("@/utils/styleUtils", () => ({})); + +jest.mock( + "@/components/core/parameterRenderComponent/components/tableComponent/components/tableAutoCellRender", + () => () => null, +); + +jest.mock( + "@/components/core/parameterRenderComponent/components/tableComponent/components/tableDropdownCellEditor", + () => () => null, +); + +// Jest can't find this module to mock it, let's skip this mock + +// Jest can't find this module either + +// Mock react-cookie +jest.mock("react-cookie"); + +import { getAuthCookie } from "@/utils/utils"; + +describe("getAuthCookie", () => { + let mockCookies: jest.Mocked; + + beforeEach(() => { + mockCookies = { + get: jest.fn(), + set: jest.fn(), + remove: jest.fn(), + } as any; + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it("should return the cookie value when cookie exists", () => { + const mockTokenValue = "test-access-token"; + mockCookies.get.mockReturnValue(mockTokenValue); + + const result = getAuthCookie(mockCookies, "access_token"); + + expect(mockCookies.get).toHaveBeenCalledWith("access_token"); + expect(result).toBe(mockTokenValue); + }); + + it("should return undefined when cookie does not exist", () => { + mockCookies.get.mockReturnValue(undefined); + + const result = getAuthCookie(mockCookies, "nonexistent_token"); + + expect(mockCookies.get).toHaveBeenCalledWith("nonexistent_token"); + expect(result).toBeUndefined(); + }); + + it("should handle different token names", () => { + const testCases = [ + { tokenName: "access_token_lf", value: "access-123" }, + { tokenName: "refresh_token_lf", value: "refresh-456" }, + { tokenName: "apikey_tkn_lflw", value: "api-789" }, + ]; + + testCases.forEach(({ tokenName, value }) => { + mockCookies.get.mockReturnValue(value); + + const result = getAuthCookie(mockCookies, tokenName); + + expect(mockCookies.get).toHaveBeenCalledWith(tokenName); + expect(result).toBe(value); + }); + }); + + it("should handle empty string token names", () => { + const result = getAuthCookie(mockCookies, ""); + + expect(mockCookies.get).toHaveBeenCalledWith(""); + }); + + it("should handle null values from cookies", () => { + mockCookies.get.mockReturnValue(null); + + const result = getAuthCookie(mockCookies, "test_token"); + + expect(mockCookies.get).toHaveBeenCalledWith("test_token"); + expect(result).toBeNull(); + }); +}); diff --git a/src/frontend/src/shared/hooks/__tests__/use-set-cookie-auth.test.ts b/src/frontend/src/shared/hooks/__tests__/use-set-cookie-auth.test.ts new file mode 100644 index 000000000..9e6970479 --- /dev/null +++ b/src/frontend/src/shared/hooks/__tests__/use-set-cookie-auth.test.ts @@ -0,0 +1,130 @@ +import { Cookies } from "react-cookie"; + +// Mock all complex dependencies to avoid import issues +jest.mock("@/stores/authStore", () => ({ + __esModule: true, + default: jest.fn(() => ({})), +})); + +jest.mock("@/stores/darkStore", () => ({ + useDarkStore: { + getState: () => ({ refreshStars: jest.fn() }), + setState: jest.fn(), + subscribe: jest.fn(), + destroy: jest.fn(), + }, +})); + +jest.mock("@/stores/alertStore", () => ({ + __esModule: true, + default: jest.fn(() => ({})), +})); + +jest.mock("@/utils/styleUtils", () => ({})); + +jest.mock( + "@/components/core/parameterRenderComponent/components/tableComponent/components/tableAutoCellRender", + () => () => null, +); + +jest.mock( + "@/components/core/parameterRenderComponent/components/tableComponent/components/tableDropdownCellEditor", + () => () => null, +); + +// Jest can't find this module to mock it, let's skip this mock + +// Jest can't find this module either + +// Mock react-cookie +jest.mock("react-cookie"); + +import { setAuthCookie } from "@/utils/utils"; + +describe("setAuthCookie", () => { + let mockCookies: jest.Mocked; + + beforeEach(() => { + mockCookies = { + get: jest.fn(), + set: jest.fn(), + remove: jest.fn(), + } as any; + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it("should set a cookie with correct options", () => { + const tokenName = "access_token_lf"; + const tokenValue = "test-access-token"; + + setAuthCookie(mockCookies, tokenName, tokenValue); + + expect(mockCookies.set).toHaveBeenCalledWith(tokenName, tokenValue, { + path: "/", + secure: true, + sameSite: "strict", + }); + }); + + it("should handle different token types", () => { + const testCases = [ + { tokenName: "access_token_lf", value: "access-123" }, + { tokenName: "refresh_token_lf", value: "refresh-456" }, + { tokenName: "apikey_tkn_lflw", value: "api-789" }, + ]; + + testCases.forEach(({ tokenName, value }) => { + setAuthCookie(mockCookies, tokenName, value); + + expect(mockCookies.set).toHaveBeenCalledWith(tokenName, value, { + path: "/", + secure: true, + sameSite: "strict", + }); + }); + }); + + it("should handle empty string values", () => { + const tokenName = "test_token"; + const tokenValue = ""; + + setAuthCookie(mockCookies, tokenName, tokenValue); + + expect(mockCookies.set).toHaveBeenCalledWith(tokenName, tokenValue, { + path: "/", + secure: true, + sameSite: "strict", + }); + }); + + it("should use correct cookie options for security", () => { + setAuthCookie(mockCookies, "test_token", "test_value"); + + const cookieOptions = mockCookies.set.mock.calls[0][2]; + + expect(cookieOptions).toEqual({ + path: "/", + secure: true, + sameSite: "strict", + }); + + // Ensure httpOnly is NOT set (removed from utils.ts) + expect(cookieOptions).not.toHaveProperty("httpOnly"); + }); + + it("should handle special characters in token values", () => { + const tokenName = "test_token"; + const tokenValue = "token-with-special-chars_@#$%"; + + setAuthCookie(mockCookies, tokenName, tokenValue); + + expect(mockCookies.set).toHaveBeenCalledWith(tokenName, tokenValue, { + path: "/", + secure: true, + sameSite: "strict", + }); + }); +}); diff --git a/src/frontend/src/stores/__tests__/authStore.test.ts b/src/frontend/src/stores/__tests__/authStore.test.ts new file mode 100644 index 000000000..1305a5434 --- /dev/null +++ b/src/frontend/src/stores/__tests__/authStore.test.ts @@ -0,0 +1,206 @@ +import { act, renderHook } from "@testing-library/react"; + +// Mock the darkStore to avoid import.meta issues +jest.mock("../darkStore", () => ({ + useDarkStore: { + getState: () => ({ refreshStars: jest.fn() }), + setState: jest.fn(), + subscribe: jest.fn(), + destroy: jest.fn(), + }, +})); + +// Mock all complex dependencies to avoid import issues +jest.mock("@/stores/alertStore", () => ({ + __esModule: true, + default: jest.fn(() => ({})), +})); + +jest.mock("@/utils/styleUtils", () => ({})); + +jest.mock( + "@/components/core/parameterRenderComponent/components/tableComponent/components/tableAutoCellRender", + () => () => null, +); + +jest.mock( + "@/components/core/parameterRenderComponent/components/tableComponent/components/tableDropdownCellEditor", + () => () => null, +); + +// Jest can't find this module to mock it, let's skip this mock + +// Jest can't find this module either + +import useAuthStore from "../authStore"; + +// We can't easily mock the cookie hook at initialization time, so we'll test actual behavior +describe("useAuthStore", () => { + describe("state management", () => { + it("should update isAdmin state", () => { + const { result } = renderHook(() => useAuthStore()); + + act(() => { + result.current.setIsAdmin(true); + }); + + expect(result.current.isAdmin).toBe(true); + + act(() => { + result.current.setIsAdmin(false); + }); + + expect(result.current.isAdmin).toBe(false); + }); + + it("should update isAuthenticated state", () => { + const { result } = renderHook(() => useAuthStore()); + + act(() => { + result.current.setIsAuthenticated(true); + }); + + expect(result.current.isAuthenticated).toBe(true); + + act(() => { + result.current.setIsAuthenticated(false); + }); + + expect(result.current.isAuthenticated).toBe(false); + }); + + it("should update accessToken state", () => { + const { result } = renderHook(() => useAuthStore()); + + act(() => { + result.current.setAccessToken("new-access-token"); + }); + + expect(result.current.accessToken).toBe("new-access-token"); + + act(() => { + result.current.setAccessToken(null); + }); + + expect(result.current.accessToken).toBeNull(); + }); + + it("should update userData state", () => { + const { result } = renderHook(() => useAuthStore()); + const mockUserData = { + id: "123", + username: "testuser", + is_superuser: false, + is_active: true, + profile_image: "", + create_at: new Date(), + updated_at: new Date(), + }; + + act(() => { + result.current.setUserData(mockUserData); + }); + + expect(result.current.userData).toEqual(mockUserData); + + act(() => { + result.current.setUserData(null); + }); + + expect(result.current.userData).toBeNull(); + }); + + it("should update autoLogin state", () => { + const { result } = renderHook(() => useAuthStore()); + + act(() => { + result.current.setAutoLogin(true); + }); + + expect(result.current.autoLogin).toBe(true); + + act(() => { + result.current.setAutoLogin(false); + }); + + expect(result.current.autoLogin).toBe(false); + }); + + it("should update apiKey state", () => { + const { result } = renderHook(() => useAuthStore()); + + act(() => { + result.current.setApiKey("new-api-key"); + }); + + expect(result.current.apiKey).toBe("new-api-key"); + + act(() => { + result.current.setApiKey(null); + }); + + expect(result.current.apiKey).toBeNull(); + }); + + it("should update authenticationErrorCount state", () => { + const { result } = renderHook(() => useAuthStore()); + + act(() => { + result.current.setAuthenticationErrorCount(5); + }); + + expect(result.current.authenticationErrorCount).toBe(5); + + act(() => { + result.current.setAuthenticationErrorCount(0); + }); + + expect(result.current.authenticationErrorCount).toBe(0); + }); + }); + + describe("logout function", () => { + it("should reset auth-related state on logout", async () => { + const { result } = renderHook(() => useAuthStore()); + + // Set up some state first + act(() => { + result.current.setIsAuthenticated(true); + result.current.setIsAdmin(true); + result.current.setAccessToken("access-token"); + result.current.setApiKey("api-key"); + result.current.setUserData({ + id: "123", + username: "test", + is_superuser: true, + is_active: true, + profile_image: "", + create_at: new Date(), + updated_at: new Date(), + }); + result.current.setAutoLogin(true); + }); + + // Verify state is set + expect(result.current.isAuthenticated).toBe(true); + expect(result.current.isAdmin).toBe(true); + expect(result.current.accessToken).toBe("access-token"); + expect(result.current.apiKey).toBe("api-key"); + expect(result.current.userData).toBeTruthy(); + expect(result.current.autoLogin).toBe(true); + + // Perform logout + await act(async () => { + await result.current.logout(); + }); + + // Verify state is reset + expect(result.current.isAuthenticated).toBe(false); + expect(result.current.isAdmin).toBe(false); + expect(result.current.accessToken).toBeNull(); + expect(result.current.apiKey).toBeNull(); + expect(result.current.userData).toBeNull(); + expect(result.current.autoLogin).toBe(false); + }); + }); +}); diff --git a/src/frontend/src/stores/authStore.ts b/src/frontend/src/stores/authStore.ts index 4a281e444..86bdcfd56 100644 --- a/src/frontend/src/stores/authStore.ts +++ b/src/frontend/src/stores/authStore.ts @@ -2,7 +2,10 @@ import { Cookies } from "react-cookie"; import { create } from "zustand"; -import { LANGFLOW_ACCESS_TOKEN } from "@/constants/constants"; +import { + LANGFLOW_ACCESS_TOKEN, + LANGFLOW_API_TOKEN, +} from "@/constants/constants"; import type { AuthStoreType } from "@/types/zustand/auth"; const cookies = new Cookies(); @@ -12,7 +15,7 @@ const useAuthStore = create((set, get) => ({ accessToken: cookies.get(LANGFLOW_ACCESS_TOKEN) ?? null, userData: null, autoLogin: null, - apiKey: cookies.get("apikey_tkn_lflw"), + apiKey: cookies.get(LANGFLOW_API_TOKEN), authenticationErrorCount: 0, setIsAdmin: (isAdmin) => set({ isAdmin }), diff --git a/src/frontend/src/utils/utils.ts b/src/frontend/src/utils/utils.ts index 1096163c3..d9fcc37f9 100644 --- a/src/frontend/src/utils/utils.ts +++ b/src/frontend/src/utils/utils.ts @@ -6,6 +6,7 @@ import TableDropdownCellEditor from "@/components/core/parameterRenderComponent/ import useAlertStore from "@/stores/alertStore"; import { type ColumnField, FormatterType } from "@/types/utils/functions"; import "moment-timezone"; +import { Cookies } from "react-cookie"; import { twMerge } from "tailwind-merge"; import { DRAG_EVENTS_CUSTOM_TYPESS, @@ -536,7 +537,9 @@ export function brokenEdgeMessage({ field: string; }; }) { - return `${source.nodeDisplayName}${source.outputDisplayName ? " | " + source.outputDisplayName : ""} -> ${target.displayName}${target.field ? " | " + target.field : ""}`; + return `${source.nodeDisplayName}${ + source.outputDisplayName ? " | " + source.outputDisplayName : "" + } -> ${target.displayName}${target.field ? " | " + target.field : ""}`; } export function FormatColumns(columns: ColumnField[]): ColDef[] { if (!columns) return []; @@ -823,7 +826,8 @@ export interface CookieOptions { maxAge?: number; expires?: Date; secure?: boolean; - sameSite?: "Strict" | "Lax" | "None"; + sameSite?: "strict" | "lax" | "none"; + httpOnly?: boolean; } /** @@ -1002,3 +1006,19 @@ export const stripReleaseStageFromVersion = (version: string): string => { } return version; }; + +export const getAuthCookie = (cookies: Cookies, tokenName: string) => { + return cookies.get(tokenName); +}; + +export const setAuthCookie = ( + cookies: Cookies, + tokenName: string, + value: string, +) => { + cookies.set(tokenName, value, { + path: "/", + secure: true, + sameSite: "strict", + }); +};