diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index 76a3129d6d..5178aca079 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -8,13 +8,14 @@ import useAuthorized from "./useAuthorized"; // Unmock useAuthorized to test the actual implementation vi.unmock("@/app/(dashboard)/hooks/useAuthorized"); -const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock } = vi.hoisted(() => ({ +const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock, buildLoginUrlWithReturnMock } = vi.hoisted(() => ({ replaceMock: vi.fn(), clearTokenCookiesMock: vi.fn(), getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), getUiConfigMock: vi.fn(), decodeTokenMock: vi.fn(), checkTokenValidityMock: vi.fn(), + buildLoginUrlWithReturnMock: vi.fn((baseUrl: string) => baseUrl), })); vi.mock("next/navigation", () => ({ @@ -49,6 +50,14 @@ vi.mock("@/utils/jwtUtils", async (importOriginal) => { }; }); +vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + buildLoginUrlWithReturn: buildLoginUrlWithReturnMock, + storeReturnUrl: vi.fn(), + }; +}); const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -81,6 +90,7 @@ describe("useAuthorized", () => { getUiConfigMock.mockReset(); decodeTokenMock.mockReset(); checkTokenValidityMock.mockReset(); + buildLoginUrlWithReturnMock.mockClear(); clearCookie(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 0b60971c1e..8f8c403a4e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -3,39 +3,12 @@ import { getProxyBaseUrl } from "@/components/networking"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { checkTokenValidity, decodeToken } from "@/utils/jwtUtils"; +import { buildLoginUrlWithReturn, storeReturnUrl } from "@/utils/returnUrlUtils"; import { useRouter } from "next/navigation"; -import { useEffect, useMemo } from "react"; +import { useCallback, useEffect, useMemo } from "react"; +import { formatUserRole } from "@/utils/roles"; import { useUIConfig } from "./uiConfig/useUIConfig"; -function formatUserRole(userRole: string) { - if (!userRole) { - return "Undefined Role"; - } - switch (userRole.toLowerCase()) { - case "app_owner": - return "App Owner"; - case "demo_app_owner": - return "App Owner"; - case "app_admin": - return "Admin"; - case "proxy_admin": - return "Admin"; - case "proxy_admin_viewer": - return "Admin Viewer"; - case "org_admin": - return "Org Admin"; - case "internal_user": - return "Internal User"; - case "internal_user_viewer": - case "internal_viewer": // TODO:remove if deprecated - return "Internal Viewer"; - case "app_user": - return "App User"; - default: - return "Unknown Role"; - } -} - const useAuthorized = () => { const router = useRouter(); const { data: uiConfig, isLoading: isUIConfigLoading } = useUIConfig(); @@ -47,6 +20,14 @@ const useAuthorized = () => { const isLoading = isUIConfigLoading; const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled; + // Helper function to redirect to login while preserving the current URL + const redirectToLogin = useCallback(() => { + storeReturnUrl(); + const baseLoginUrl = `${getProxyBaseUrl()}/ui/login`; + const loginUrlWithReturn = buildLoginUrlWithReturn(baseLoginUrl); + router.replace(loginUrlWithReturn); + }, [router]); + // Single useEffect for all redirect logic useEffect(() => { if (isLoading) return; @@ -55,9 +36,9 @@ const useAuthorized = () => { if (token) { clearTokenCookies(); } - router.replace(`${getProxyBaseUrl()}/ui/login`); + redirectToLogin(); } - }, [isLoading, isAuthorized, token, router]); + }, [isLoading, isAuthorized, token, redirectToLogin]); return { isLoading, diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 2a2915ea35..d54140c5a2 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -6,6 +6,7 @@ import LoadingScreen from "@/components/common_components/LoadingScreen"; import { getProxyBaseUrl } from "@/components/networking"; import { getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; +import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd"; import { useRouter } from "next/navigation"; @@ -32,12 +33,24 @@ function LoginPageContent() { const rawToken = getCookie("token"); if (rawToken && !isJwtExpired(rawToken)) { - router.replace(`${getProxyBaseUrl()}/ui`); + // User already logged in - redirect to return URL or default + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + router.replace(returnUrl); + } else { + router.replace(`${getProxyBaseUrl()}/ui`); + } return; } if (uiConfig && uiConfig.auto_redirect_to_sso) { - router.push(`${getProxyBaseUrl()}/sso/key/generate`); + // For SSO, pass the return URL to the SSO endpoint + const returnUrl = getReturnUrl(); + let ssoUrl = `${getProxyBaseUrl()}/sso/key/generate`; + if (returnUrl && isValidReturnUrl(returnUrl)) { + ssoUrl += `?redirect_to=${encodeURIComponent(returnUrl)}`; + } + router.push(ssoUrl); return; } @@ -49,7 +62,13 @@ function LoginPageContent() { { username, password }, { onSuccess: (data) => { - router.push(data.redirect_url); + // Check if we have a return URL to use instead of the default redirect + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + router.push(returnUrl); + } else { + router.push(data.redirect_url); + } }, }, ); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index b927f312df..bddcb0ab59 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -23,7 +23,7 @@ import Navbar from "@/components/navbar"; import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName, getInProductNudgesCall } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; -import { fetchUserModels } from "@/components/organisms/create_key_button"; +import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button"; import Organizations, { fetchOrganizations } from "@/components/organizations"; import PassThroughSettings from "@/components/pass_through_settings"; import PromptsPanel from "@/components/prompts"; @@ -44,10 +44,12 @@ import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; import { isJwtExpired } from "@/utils/jwtUtils"; -import { isAdminRole } from "@/utils/roles"; +import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; +import { formatUserRole, isAdminRole } from "@/utils/roles"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { jwtDecode } from "jwt-decode"; import { useSearchParams } from "next/navigation"; -import { Suspense, useEffect, useState } from "react"; +import { Suspense, useEffect, useMemo, useRef, useState } from "react"; import { ConfigProvider, theme } from "antd"; function getCookie(name: string) { @@ -67,35 +69,6 @@ function deleteCookie(name: string, path = "/") { document.cookie = `${name}=; Max-Age=0; Path=${path}`; } -function formatUserRole(userRole: string) { - if (!userRole) { - return "Undefined Role"; - } - switch (userRole.toLowerCase()) { - case "app_owner": - return "App Owner"; - case "demo_app_owner": - return "App Owner"; - case "app_admin": - return "Admin"; - case "proxy_admin": - return "Admin"; - case "proxy_admin_viewer": - return "Admin Viewer"; - case "org_admin": - return "Org Admin"; - case "internal_user": - return "Internal User"; - case "internal_user_viewer": - case "internal_viewer": // TODO:remove if deprecated - return "Internal Viewer"; - case "app_user": - return "App User"; - default: - return "Unknown Role"; - } -} - interface ProxySettings { PROXY_BASE_URL: string; PROXY_LOGOUT_URL: string; @@ -141,6 +114,58 @@ function CreateKeyPageContent() { const invitation_id = searchParams.get("invitation_id"); + // Parse URL query parameters for pre-filling the create key form + // Includes validation to prevent injection and DoS attacks + const autoOpenCreate = searchParams.get("create") === "true"; + const prefillData: CreateKeyPrefillData | undefined = useMemo(() => { + if (!autoOpenCreate) return undefined; + + const ownedBy = searchParams.get("owned_by"); + const teamId = searchParams.get("team_id"); + const keyAlias = searchParams.get("key_alias"); + const modelsParam = searchParams.get("models"); + const keyType = searchParams.get("key_type"); + + // Only return prefill data if at least one field is provided + if (!ownedBy && !teamId && !keyAlias && !modelsParam && !keyType) { + return undefined; + } + + // Validate owned_by against allowed values + const validOwnedByValues = ["you", "service_account", "another_user"]; + const validatedOwnedBy = ownedBy && validOwnedByValues.includes(ownedBy) + ? (ownedBy as CreateKeyPrefillData["owned_by"]) + : undefined; + + // Validate key_type against allowed values + const validKeyTypes = ["default", "llm_api", "management"]; + const validatedKeyType = keyType && validKeyTypes.includes(keyType) + ? (keyType as CreateKeyPrefillData["key_type"]) + : undefined; + + // Sanitize key_alias (limit length, trim whitespace) + const sanitizedKeyAlias = keyAlias + ? keyAlias.trim().slice(0, 256) // Reasonable max length + : undefined; + + // Sanitize models (limit array size and individual model name length) + const sanitizedModels = modelsParam + ? modelsParam + .split(",") + .slice(0, 100) // Limit number of models to prevent DoS + .map(m => m.trim().slice(0, 256)) // Limit individual model name length + .filter(m => m.length > 0) // Remove empty strings + : undefined; + + return { + owned_by: validatedOwnedBy, + team_id: teamId?.trim() || undefined, + key_alias: sanitizedKeyAlias, + models: sanitizedModels && sanitizedModels.length > 0 ? sanitizedModels : undefined, + key_type: validatedKeyType, + }; + }, [searchParams, autoOpenCreate]); + // Get page from URL, default to 'api-keys' if not present const [page, setPage] = useState(() => { return searchParams.get("page") || "api-keys"; @@ -161,6 +186,9 @@ function CreateKeyPageContent() { const [accessToken, setAccessToken] = useState(null); const [sidebarCollapsed, setSidebarCollapsed] = useState(false); + // Track if we've already attempted a return URL redirect to prevent race conditions + const hasAttemptedReturnRedirectRef = useRef(false); + const toggleSidebar = () => { setSidebarCollapsed(!sidebarCollapsed); }; @@ -205,12 +233,48 @@ function CreateKeyPageContent() { useEffect(() => { if (redirectToLogin) { + // Store the current URL so we can redirect back after login + storeReturnUrl(); + // Build login URL with return URL parameter + const baseLoginUrl = (proxyBaseUrl || "") + "/ui/login"; + const dest = buildLoginUrlWithReturn(baseLoginUrl); // Replace instead of assigning to avoid back-button loops - const dest = (proxyBaseUrl || "") + "/ui/login"; window.location.replace(dest); } }, [redirectToLogin]); + // Check for a stored return URL after successful authentication + // This handles the case where user comes back from SSO and we need to redirect to the original URL + useEffect(() => { + // Skip if still loading, no token, or we've already attempted a redirect + if (authLoading || !token || hasAttemptedReturnRedirectRef.current) { + return; + } + + // Mark that we've attempted the redirect to prevent race conditions + // This prevents duplicate redirects if token changes (e.g., refresh) + hasAttemptedReturnRedirectRef.current = true; + + // Check for a stored return URL + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + const currentUrl = window.location.href; + const normalizedReturnUrl = normalizeUrlForCompare(returnUrl); + const normalizedCurrentUrl = normalizeUrlForCompare(currentUrl); + // Only redirect if the return URL is different from the current URL + // This prevents infinite redirect loops + if (normalizedReturnUrl !== normalizedCurrentUrl) { + window.location.replace(returnUrl); + } + } + }, [authLoading, token]); + + useEffect(() => { + if (!token) { + hasAttemptedReturnRedirectRef.current = false; + } + }, [token]); + useEffect(() => { if (!token) { return; @@ -407,9 +471,8 @@ function CreateKeyPageContent() { />
- -
- + +
{page == "api-keys" ? ( ) : page == "models" ? ( { - const fn = vi.fn().mockResolvedValue({ +const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall } = vi.hoisted(() => { + const formStateRef = { current: {} as Record }; + const mockKeyCreateCall = vi.fn().mockResolvedValue({ key: "test-api-key", soft_budget: null, }); - return { mockKeyCreateCall: fn }; + const formMock = { + setFieldsValue: vi.fn((values: Record) => { + Object.assign(formStateRef.current, values); + }), + setFieldValue: vi.fn((name: string, value: any) => { + formStateRef.current[name] = value; + }), + getFieldValue: vi.fn((name: string) => formStateRef.current[name]), + resetFields: vi.fn(() => { + formStateRef.current = {}; + }), + }; + const radioGroupValueRef = { current: null as string | null }; + return { + formMock, + setFieldsValueMock: formMock.setFieldsValue, + radioGroupValueRef, + formStateRef, + mockKeyCreateCall, + }; +}); + +const defaultAuthorizedState = { + accessToken: "test-token", + userId: "test-user-id", + userRole: "Admin", + premiumUser: false, +}; + +let authorizedState = { ...defaultAuthorizedState }; + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => authorizedState, +})); + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + keyKeys: { + lists: () => ["keys"], + }, +})); + +vi.mock("@ant-design/icons", () => ({ + InfoCircleOutlined: () => null, +})); + +vi.mock("react-copy-to-clipboard", () => ({ + CopyToClipboard: ({ children }: { children: any }) => children, +})); + +vi.mock("@tremor/react", () => { + const React = require("react"); + const Stub = ({ children }: { children?: any }) => React.createElement("div", null, children); + const Button = ({ children, ...props }: { children?: any }) => + React.createElement("button", props, children); + const TextInput = (props: any) => React.createElement("input", props); + + return { + Accordion: Stub, + AccordionBody: Stub, + AccordionHeader: Stub, + Button, + Col: Stub, + Grid: Stub, + Text: Stub, + TextInput, + Title: Stub, + }; +}); + +vi.mock("antd", () => { + const React = require("react"); + + const getValueFromEvent = (event: any) => { + if (event?.target) { + if (event.target.type === "checkbox") { + return event.target.checked; + } + return event.target.value; + } + return event; + }; + + const Form = ({ children, onFinish, ...props }: { children?: any; onFinish?: (values: Record) => void }) => + React.createElement( + "form", + { + ...props, + onSubmit: (event: Event) => { + event.preventDefault(); + onFinish?.({ ...formStateRef.current }); + }, + }, + children, + ); + + Form.Item = ({ children, name }: { children?: any; name?: string }) => { + if (!name || !React.isValidElement(children)) { + return React.createElement(React.Fragment, null, children); + } + + return React.cloneElement(children, { + value: formStateRef.current[name], + onChange: (event: any) => { + formStateRef.current[name] = getValueFromEvent(event); + }, + }); + }; + + Form.useForm = () => [formMock]; + + const Select = ({ children, onChange, ...props }: { children?: any; onChange?: (value: string) => void }) => + React.createElement( + "select", + { + ...props, + onChange: (event: any) => onChange?.(event.target.value), + }, + children, + ); + + Select.Option = ({ children, ...props }: { children?: any }) => + React.createElement("option", props, children); + + const Input = (props: any) => React.createElement("input", props); + Input.Password = (props: any) => React.createElement("input", { ...props, type: "password" }); + Input.TextArea = (props: any) => React.createElement("textarea", props); + + const Modal = ({ children, open }: { children?: any; open?: boolean }) => + open ? React.createElement("div", null, children) : null; + + const Radio = ({ children, ...props }: { children?: any }) => + React.createElement("div", props, children); + + Radio.Group = ({ children, value }: { children?: any; value?: string }) => { + radioGroupValueRef.current = value ?? null; + return React.createElement("div", null, children); + }; + + const Switch = (props: any) => React.createElement("input", { ...props, type: "checkbox" }); + const Tag = ({ children }: { children?: any }) => React.createElement("span", null, children); + const Tooltip = ({ children }: { children?: any }) => React.createElement(React.Fragment, null, children); + + const Button = ({ children, htmlType, ...props }: { children?: any; htmlType?: string }) => + React.createElement("button", { ...props, type: htmlType ?? props.type }, children); + + return { + Button, + Form, + Input, + message: { + success: vi.fn(), + error: vi.fn(), + warning: vi.fn(), + info: vi.fn(), + }, + Modal, + Radio, + Select, + Switch, + Tag, + Tooltip, + }; }); vi.mock("../networking", () => ({ keyCreateCall: mockKeyCreateCall, - modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }] }), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }), getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }), @@ -42,6 +204,25 @@ vi.mock("../molecules/notifications_manager", () => ({ }, })); +vi.mock("../agent_management/AgentSelector", () => ({ default: () => null })); +vi.mock("../common_components/budget_duration_dropdown", () => ({ default: () => null })); +vi.mock("../common_components/check_openapi_schema", () => ({ default: () => null })); +vi.mock("../common_components/KeyLifecycleSettings", () => ({ default: () => null })); +vi.mock("../common_components/ModelAliasManager", () => ({ default: () => null })); +vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () => null })); +vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null })); +vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null })); +vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null })); +vi.mock("../common_components/team_dropdown", () => ({ default: () => null })); +vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null })); +vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null })); +vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); +vi.mock("../shared/numerical_input", () => ({ default: () => null })); +vi.mock("../vector_store_management/VectorStoreSelector", () => ({ default: () => null })); +vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({ + getModelDisplayName: (model: string) => model, +})); + vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); @@ -61,7 +242,7 @@ vi.mock("../common_components/AccessGroupSelector", () => ({ onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + onChange={(event) => onChange?.(event.target.value ? event.target.value.split(",").map((v) => v.trim()) : [])} /> ), })); @@ -69,14 +250,19 @@ vi.mock("../common_components/AccessGroupSelector", () => ({ describe("CreateKey", () => { const defaultProps = { team: null, - data: [], teams: [], + data: [], addKey: vi.fn(), }; beforeEach(() => { vi.clearAllMocks(); - localStorage.clear(); + if (typeof window !== "undefined" && window.localStorage && typeof window.localStorage.clear === "function") { + window.localStorage.clear(); + } + authorizedState = { ...defaultAuthorizedState }; + radioGroupValueRef.current = null; + formStateRef.current = {}; mockKeyCreateCall.mockResolvedValue({ key: "test-api-key", soft_budget: null, @@ -96,26 +282,8 @@ describe("CreateKey", () => { }); await waitFor(() => { - expect(screen.getByText("Key Type")).toBeInTheDocument(); - }); - - // Open the Key Type dropdown - const keyTypeSection = screen.getByText("Key Type").closest(".ant-form-item")!; - const selectElement = keyTypeSection.querySelector(".ant-select-selector")!; - act(() => { - fireEvent.mouseDown(selectElement); - }); - - await waitFor(() => { - // Verify "AI APIs" appears as an option - const options = document.querySelectorAll(".ant-select-item-option"); - const optionTexts = Array.from(options).map((el) => el.textContent); - const hasAIAPIs = optionTexts.some((text) => text?.includes("AI APIs")); - expect(hasAIAPIs).toBe(true); - - // Verify old "LLM API" label does NOT appear - const hasLLMAPI = optionTexts.some((text) => text?.includes("LLM API")); - expect(hasLLMAPI).toBe(false); + expect(screen.getByText("AI APIs")).toBeInTheDocument(); + expect(screen.queryByText("LLM API")).not.toBeInTheDocument(); }); }); @@ -126,46 +294,118 @@ describe("CreateKey", () => { fireEvent.click(screen.getByRole("button", { name: /create new key/i })); }); - await waitFor(() => { - expect(screen.getByLabelText(/key name/i)).toBeInTheDocument(); - }); - - fireEvent.change(screen.getByLabelText(/key name/i), { target: { value: "Test Key" } }); - - const optionalSettingsAccordion = screen.getByText("Optional Settings"); - act(() => { - fireEvent.click(optionalSettingsAccordion); - }); - await waitFor(() => { expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); }); - fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + act(() => { + fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + formMock.setFieldValue("key_alias", "Test Key"); + }); - const modelsCombobox = screen.getAllByRole("combobox").find((el) => el.closest('[class*="ant-form-item"]')?.textContent?.includes("Models")) || - screen.getAllByRole("combobox")[1]; - if (modelsCombobox) { - act(() => fireEvent.mouseDown(modelsCombobox)); - await waitFor(() => { - const allTeamModels = [...document.body.querySelectorAll(".ant-select-item")].find( - (el) => el.textContent?.includes("All Team Models"), - ); - if (allTeamModels) fireEvent.click(allTeamModels); - }); - } + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create key/i })); + }); - const createButton = screen.getByRole("button", { name: /create key/i }); - act(() => fireEvent.click(createButton)); + await waitFor(() => { + expect(mockKeyCreateCall).toHaveBeenCalled(); + const formValues = mockKeyCreateCall.mock.calls[0][2]; + expect(formValues).toHaveProperty("access_group_ids"); + expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]); + }); + }); - await waitFor( - () => { - expect(mockKeyCreateCall).toHaveBeenCalled(); - const formValues = mockKeyCreateCall.mock.calls[0][2]; - expect(formValues).toHaveProperty("access_group_ids"); - expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]); - }, - { timeout: 15000 }, + it("should prefill models when provided without team_id", async () => { + renderWithProviders( + , ); - }, { timeout: 30000 }); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ models: ["gpt-4"] }); + }); + }); + + it("should prefill team_id when it exists in teams", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ team_id: "team-1" }); + }); + }); + + it("should ignore team_id when it does not exist in teams", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" }); + }); + + expect(setFieldsValueMock).not.toHaveBeenCalledWith({ team_id: "team-404" }); + }); + + it('should fall back to "you" when owned_by is another_user for non-admin', async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" }); + }); + + expect(radioGroupValueRef.current).toBe("you"); + }); + + it("should apply owned_by another_user for admin", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(radioGroupValueRef.current).toBe("another_user"); + }); + }); + + it("should prefill key_type when provided", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index c54725bd5a..78ff01d641 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -49,11 +49,24 @@ import { simplifyKeyGenerateError } from "./utils"; const { Option } = Select; +/** + * Interface for pre-filling the create key form from URL parameters + */ +export interface CreateKeyPrefillData { + owned_by?: "you" | "service_account" | "another_user"; + team_id?: string; + key_alias?: string; + models?: string[]; + key_type?: "default" | "llm_api" | "management"; +} + interface CreateKeyProps { team: Team | null; data: any[] | null; teams: Team[] | null; addKey: (data: any) => void; + autoOpenCreate?: boolean; + prefillData?: CreateKeyPrefillData; } interface User { @@ -144,7 +157,7 @@ export const fetchUserModels = async ( * Please contribute to the new refactor. * ───────────────────────────────────────────────────────────────────────── */ -const CreateKey: React.FC = ({ team, teams, data, addKey }) => { +const CreateKey: React.FC = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); const { data: projects, isLoading: isProjectsLoading } = useProjects(); @@ -159,6 +172,8 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { const [modelsToPick, setModelsToPick] = useState([]); const [keyOwner, setKeyOwner] = useState("you"); const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data)); + const [hasPrefilled, setHasPrefilled] = useState(false); + const [pendingPrefillModels, setPendingPrefillModels] = useState(null); const [guardrailsList, setGuardrailsList] = useState([]); const [policiesList, setPoliciesList] = useState([]); const [promptsList, setPromptsList] = useState([]); @@ -284,6 +299,55 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { fetchPossibleRoles(); }, [accessToken]); + // Auto-open modal and prefill form from URL params (deep link). + // Guarded by write access so we don't open for read-only users. + useEffect(() => { + if (autoOpenCreate && !hasPrefilled && teams && userRole && rolesWithWriteAccess.includes(userRole)) { + // Open the modal + setIsModalVisible(true); + setHasPrefilled(true); + + // Apply prefill data if provided + if (prefillData) { + // Set key owner (owned_by) - validate that "another_user" is only allowed for Admin + if (prefillData.owned_by) { + if (prefillData.owned_by === "another_user" && userRole !== "Admin") { + // Ignore invalid owned_by for non-admin users, fall back to default + setKeyOwner("you"); + } else { + setKeyOwner(prefillData.owned_by); + } + } + + // Set team - find the team by ID and set it (only if team exists in user's teams) + if (prefillData.team_id) { + const selectedTeam = teams?.find((t) => t.team_id === prefillData.team_id) || null; + if (selectedTeam) { + setSelectedCreateKeyTeam(selectedTeam); + form.setFieldsValue({ team_id: prefillData.team_id }); + } + // Silently ignore invalid team_id - don't prefill with a team user doesn't have access to + } + + // Set key alias + if (prefillData.key_alias) { + form.setFieldsValue({ key_alias: prefillData.key_alias }); + } + + // Defer model selection until we load the allowed model list. + if (prefillData.models && prefillData.models.length > 0) { + setPendingPrefillModels(prefillData.models); + } + + // Set key type + if (prefillData.key_type) { + setKeyType(prefillData.key_type); + form.setFieldsValue({ key_type: prefillData.key_type }); + } + } + } + }, [autoOpenCreate, prefillData, teams, hasPrefilled, form, userRole]); + // Check if team selection is required const isTeamSelectionRequired = modelsToPick.includes("no-default-models"); const isFormDisabled = isTeamSelectionRequired && !selectedCreateKeyTeam; @@ -477,6 +541,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { NotificationsManager.success("Virtual Key copied to clipboard"); }; + // Fetch available models when team or auth changes. + // Note: Model prefill from URL params is handled by the useEffect below, which + // watches for pendingPrefillModels + modelsToPick to both be populated. useEffect(() => { if (selectedProjectId) { // When a project is selected, use the project's models @@ -492,9 +559,28 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { setModelsToPick(allModels); }); } - form.setFieldValue("models", []); - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [selectedCreateKeyTeam, selectedProjectId, accessToken, userID, userRole]); + // Only clear models if we don't have pending prefill models + if (!pendingPrefillModels) { + form.setFieldValue("models", []); + } + }, [selectedCreateKeyTeam, selectedProjectId, accessToken, userID, userRole, form]); + + // Apply deferred model prefill once the available model list arrives. + // This handles timing where prefill data arrives before or after models are fetched. + useEffect(() => { + if (!pendingPrefillModels || pendingPrefillModels.length === 0) { + return; + } + if (!modelsToPick || modelsToPick.length === 0) { + return; + } + + const validModels = pendingPrefillModels.filter((model) => modelsToPick.includes(model)); + if (validModels.length > 0) { + form.setFieldsValue({ models: validModels }); + } + setPendingPrefillModels(null); + }, [pendingPrefillModels, modelsToPick, form]); // Sync team when project is selected but teams loaded later (race condition) useEffect(() => { diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index f0efe6bbaa..dc8d2d9cdd 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -16,7 +16,7 @@ import { Organization, userInfoCall, } from "./networking"; -import CreateKey from "./organisms/create_key_button"; +import CreateKey, { CreateKeyPrefillData } from "./organisms/create_key_button"; import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable"; export interface ProxySettings { @@ -55,6 +55,8 @@ interface UserDashboardProps { organizations: Organization[] | null; addKey: (data: any) => void; createClicked: boolean; + autoOpenCreate?: boolean; + prefillData?: CreateKeyPrefillData; } type TeamInterface = { @@ -77,6 +79,8 @@ const UserDashboard: React.FC = ({ organizations, addKey, createClicked, + autoOpenCreate, + prefillData, }) => { const [userSpendData, setUserSpendData] = useState(null); const [currentOrg, setCurrentOrg] = useState(null); @@ -352,6 +356,8 @@ const UserDashboard: React.FC = ({ teams={teams as Team[]} data={keys} addKey={addKey} + autoOpenCreate={autoOpenCreate} + prefillData={prefillData} /> diff --git a/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts b/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts new file mode 100644 index 0000000000..3c09e55014 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts @@ -0,0 +1,383 @@ +import { + buildLoginUrlWithReturn, + clearStoredReturnUrl, + consumeReturnUrl, + getCurrentUrl, + getReturnUrl, + getReturnUrlFromParams, + getStoredReturnUrl, + isValidReturnUrl, + storeReturnUrl, +} from "./returnUrlUtils"; + +describe("returnUrlUtils", () => { + const originalLocation = window.location; + + beforeEach(() => { + // Clear cookies before each test + document.cookie.split(";").forEach((c) => { + document.cookie = c + .replace(/^ +/, "") + .replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/"); + }); + + // Reset location mock + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?page=api-keys", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?page=api-keys", + }, + writable: true, + }); + }); + + afterEach(() => { + // Restore original location + Object.defineProperty(window, "location", { + value: originalLocation, + writable: true, + }); + }); + + describe("getCurrentUrl", () => { + it("should return the current URL", () => { + const url = getCurrentUrl(); + expect(url).toBe("http://localhost:3000/ui?page=api-keys"); + }); + }); + + describe("storeReturnUrl and getStoredReturnUrl", () => { + it("should store and retrieve the return URL from cookie", () => { + storeReturnUrl(); + const storedUrl = getStoredReturnUrl(); + expect(storedUrl).toBe("http://localhost:3000/ui?page=api-keys"); + }); + + it("should return null if no URL is stored", () => { + const storedUrl = getStoredReturnUrl(); + expect(storedUrl).toBeNull(); + }); + }); + + describe("clearStoredReturnUrl", () => { + it("should clear the stored return URL", () => { + storeReturnUrl(); + expect(getStoredReturnUrl()).not.toBeNull(); + + clearStoredReturnUrl(); + expect(getStoredReturnUrl()).toBeNull(); + }); + }); + + describe("getReturnUrlFromParams", () => { + it("should return the redirect_to parameter from URL", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fcreate%3Dtrue", + }, + writable: true, + }); + + const returnUrl = getReturnUrlFromParams(); + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + }); + + it("should return null if redirect_to parameter is not present", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "?page=api-keys", + }, + writable: true, + }); + + const returnUrl = getReturnUrlFromParams(); + expect(returnUrl).toBeNull(); + }); + }); + + describe("buildLoginUrlWithReturn", () => { + it("should build login URL with return URL parameter", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui?create=true&team_id=123", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn("/ui/login"); + expect(loginUrl).toBe( + "/ui/login?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fcreate%3Dtrue%26team_id%3D123" + ); + }); + + it("should not add return URL if already on login page", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui/login", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn("/ui/login"); + expect(loginUrl).toBe("/ui/login"); + }); + + it("should handle login URL with existing query parameters", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui?page=api-keys", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn("/ui/login?foo=bar"); + expect(loginUrl).toContain("&redirect_to="); + }); + }); + + describe("getReturnUrl", () => { + it("should prefer URL params over cookie", () => { + // Store a URL in cookie + storeReturnUrl(); + + // Set a different URL in the params + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fpage%3Dteams", + }, + writable: true, + }); + + const returnUrl = getReturnUrl(); + expect(returnUrl).toBe("http://localhost:3000/ui?page=teams"); + }); + + it("should fall back to cookie if no URL param", () => { + // Store a URL in cookie first + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?create=true", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?create=true", + }, + writable: true, + }); + storeReturnUrl(); + + // Clear the URL params + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "", + }, + writable: true, + }); + + const returnUrl = getReturnUrl(); + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + }); + + it("should return null if no return URL found", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "", + }, + writable: true, + }); + + const returnUrl = getReturnUrl(); + expect(returnUrl).toBeNull(); + }); + }); + + describe("isValidReturnUrl", () => { + it("should validate relative URLs starting with /", () => { + expect(isValidReturnUrl("/ui?page=api-keys")).toBe(true); + expect(isValidReturnUrl("/ui/teams")).toBe(true); + }); + + it("should reject protocol-relative URLs", () => { + expect(isValidReturnUrl("//evil.com")).toBe(false); + }); + + it("should validate same-hostname URLs (even with different ports) in dev", () => { + // Same hostname, same port + expect(isValidReturnUrl("http://localhost:3000/ui?page=teams")).toBe(true); + // Same hostname, different port (important for dev environments) + expect(isValidReturnUrl("http://localhost:4000/ui?page=teams")).toBe(true); + }); + + it("should reject different-hostname URLs", () => { + expect(isValidReturnUrl("http://evil.com/ui")).toBe(false); + expect(isValidReturnUrl("https://google.com")).toBe(false); + }); + + it("should reject empty URLs", () => { + expect(isValidReturnUrl("")).toBe(false); + }); + + it("should reject invalid URLs", () => { + expect(isValidReturnUrl("not-a-url")).toBe(false); + }); + + it("should reject XSS attempts with javascript: protocol", () => { + expect(isValidReturnUrl('javascript:alert("xss")')).toBe(false); + expect(isValidReturnUrl("javascript:void(0)")).toBe(false); + }); + + it("should reject data: URLs", () => { + expect(isValidReturnUrl("data:text/html,")).toBe(false); + }); + + it("should allow 127.x.x.x addresses in dev environment", () => { + Object.defineProperty(window, "location", { + value: { + href: "http://127.0.0.1:3000/ui", + origin: "http://127.0.0.1:3000", + hostname: "127.0.0.1", + protocol: "http:", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + expect(isValidReturnUrl("http://127.0.0.1:4000/ui")).toBe(true); + }); + + it("should allow .local domains in dev environment", () => { + Object.defineProperty(window, "location", { + value: { + href: "http://myapp.local:3000/ui", + origin: "http://myapp.local:3000", + hostname: "myapp.local", + protocol: "http:", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + // Same hostname with different port should be allowed in dev + expect(isValidReturnUrl("http://myapp.local:4000/ui")).toBe(true); + }); + + it("should reject cross-port redirects in production environment", () => { + // Simulate production environment + Object.defineProperty(window, "location", { + value: { + href: "https://app.example.com/ui", + origin: "https://app.example.com", + hostname: "app.example.com", + protocol: "https:", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + // Same origin should work + expect(isValidReturnUrl("https://app.example.com/ui?page=teams")).toBe(true); + // Different port should be rejected in production + expect(isValidReturnUrl("https://app.example.com:8080/ui")).toBe(false); + // Different hostname should be rejected + expect(isValidReturnUrl("https://evil.com/ui")).toBe(false); + }); + }); + + describe("consumeReturnUrl", () => { + it("should return and clear the stored return URL", () => { + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?create=true", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?create=true", + }, + writable: true, + }); + storeReturnUrl(); + + // Clear the URL params for the consume call + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui/login", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui/login", + search: "", + }, + writable: true, + }); + + const returnUrl = consumeReturnUrl(); + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + expect(getStoredReturnUrl()).toBeNull(); + }); + + it("should return null for invalid return URLs (different hostname)", () => { + // Manually set an invalid URL in cookie + document.cookie = "litellm_return_url=" + encodeURIComponent("http://evil.com/phishing") + "; path=/"; + + const returnUrl = consumeReturnUrl(); + expect(returnUrl).toBeNull(); + }); + + it("should allow URLs with different ports on same hostname", () => { + // Store URL with port 3000 + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?create=true", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?create=true", + }, + writable: true, + }); + storeReturnUrl(); + + // Now we're on port 4000 + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:4000/ui", + origin: "http://localhost:4000", + hostname: "localhost", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + const returnUrl = consumeReturnUrl(); + // Should be valid because same hostname (localhost) + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + }); + + it("should return null if no return URL found", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "", + }, + writable: true, + }); + + const returnUrl = consumeReturnUrl(); + expect(returnUrl).toBeNull(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/returnUrlUtils.ts b/ui/litellm-dashboard/src/utils/returnUrlUtils.ts new file mode 100644 index 0000000000..76562a0122 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/returnUrlUtils.ts @@ -0,0 +1,304 @@ +/** + * Utility functions for managing return URLs during authentication flows. + * + * When a user is redirected to login, we store the original URL so they can be + * redirected back after successful authentication. + * + * NOTE: We use cookies instead of sessionStorage because the SSO flow may cross + * different ports (e.g., localhost:3000 -> localhost:4000), and sessionStorage + * is not shared across different origins. Cookies on the same hostname are shared + * across different ports. + */ + +const RETURN_URL_COOKIE_NAME = "litellm_return_url"; +const RETURN_URL_PARAM = "redirect_to"; + +/** + * Gets the current URL with all query parameters. + * Returns null if running on server-side. + */ +export function getCurrentUrl(): string | null { + if (typeof window === "undefined") { + return null; + } + return window.location.href; +} + +/** + * Sets a cookie with the given name and value. + * Automatically adds Secure flag when running over HTTPS. + */ +function setCookie(name: string, value: string, maxAgeSeconds: number = 300): void { + if (typeof document === "undefined") { + return; + } + // Set cookie with path=/ so it's available across all paths + // Use SameSite=Lax to allow the cookie to be sent on navigation from external sites (SSO redirect) + // Add Secure flag when running over HTTPS to prevent cookie from being sent over unencrypted connections + const isSecure = typeof window !== "undefined" && window.location.protocol === "https:"; + const secureFlag = isSecure ? "; Secure" : ""; + document.cookie = `${name}=${encodeURIComponent(value)}; path=/; max-age=${maxAgeSeconds}; SameSite=Lax${secureFlag}`; +} + +/** + * Gets a cookie value by name. + */ +function getCookie(name: string): string | null { + if (typeof document === "undefined") { + return null; + } + const match = document.cookie.match(new RegExp(`(^| )${name}=([^;]+)`)); + if (match) { + try { + return decodeURIComponent(match[2]); + } catch { + return match[2]; + } + } + return null; +} + +/** + * Deletes a cookie by name. + */ +function deleteCookie(name: string): void { + if (typeof document === "undefined") { + return; + } + document.cookie = `${name}=; path=/; max-age=0`; +} + +/** + * Stores the current URL in a cookie before redirecting to login. + * This allows us to redirect the user back to their original destination after login. + * Cookie expires in 5 minutes (300 seconds). + */ +export function storeReturnUrl(): void { + if (typeof window === "undefined") { + return; + } + + const currentUrl = getCurrentUrl(); + if (currentUrl) { + setCookie(RETURN_URL_COOKIE_NAME, currentUrl, 300); + } +} + +/** + * Retrieves the stored return URL from the cookie. + * Returns null if no return URL is stored or if running on server-side. + */ +export function getStoredReturnUrl(): string | null { + if (typeof window === "undefined") { + return null; + } + return getCookie(RETURN_URL_COOKIE_NAME); +} + +/** + * Clears the stored return URL from the cookie. + * Should be called after redirecting to the return URL. + */ +export function clearStoredReturnUrl(): void { + if (typeof window === "undefined") { + return; + } + + try { + deleteCookie(RETURN_URL_COOKIE_NAME); + } catch (error) { + console.error("Failed to clear return URL cookie:", error); + } +} + +/** + * Gets the return URL from URL query parameters. + * Used when the return URL is passed via query string to the login page. + */ +export function getReturnUrlFromParams(): string | null { + if (typeof window === "undefined") { + return null; + } + + const searchParams = new URLSearchParams(window.location.search); + return searchParams.get(RETURN_URL_PARAM); +} + +/** + * Builds a login URL with the return URL as a query parameter. + * + * @param baseLoginUrl - The base login URL (e.g., "/ui/login") + * @param returnUrl - The URL to redirect to after login (defaults to current URL) + */ +export function buildLoginUrlWithReturn(baseLoginUrl: string, returnUrl?: string): string { + const url = returnUrl || getCurrentUrl(); + + if (!url) { + return baseLoginUrl; + } + + // Don't add return URL if we're already on the login page + if (url.includes("/login")) { + return baseLoginUrl; + } + + const separator = baseLoginUrl.includes("?") ? "&" : "?"; + return `${baseLoginUrl}${separator}${RETURN_URL_PARAM}=${encodeURIComponent(url)}`; +} + +/** + * Gets the best return URL to use after login. + * Priority: + * 1. URL query parameter (redirect_to) + * 2. Cookie + * 3. null (caller should use default) + */ +export function getReturnUrl(): string | null { + // First check URL params + const paramUrl = getReturnUrlFromParams(); + if (paramUrl) { + return paramUrl; + } + + // Then check cookie + const storedUrl = getStoredReturnUrl(); + if (storedUrl) { + return storedUrl; + } + + return null; +} + +/** + * Checks if we're running in a development environment. + * Returns true for localhost, 127.0.0.1, IPv6 localhost, or .local domains. + * This determines whether cross-port redirects are allowed (dev only). + */ +function isDevEnvironment(): boolean { + if (typeof window === "undefined") { + return false; + } + const hostname = window.location.hostname; + return ( + hostname === "localhost" || + hostname === "127.0.0.1" || + hostname === "::1" || + hostname.startsWith("127.") || // Full IPv4 loopback range (127.0.0.0/8) + hostname.endsWith(".local") // Common dev domain suffix + ); +} + +/** + * Validates a return URL to prevent open redirect attacks. + * - Always allows relative URLs (starting with / but not //) + * - In dev (localhost): allows same hostname with any port + * - In production: requires exact origin match (protocol + hostname + port) + * + * @param url - The URL to validate + * @returns true if the URL is safe to redirect to + */ +export function isValidReturnUrl(url: string): boolean { + if (!url) { + return false; + } + + // Allow relative URLs + if (url.startsWith("/") && !url.startsWith("//")) { + return true; + } + + // For absolute URLs, validate against current origin + if (typeof window === "undefined") { + return false; + } + + try { + const returnUrlObj = new URL(url); + const currentHostname = window.location.hostname; + + // Hostname must always match + if (returnUrlObj.hostname !== currentHostname) { + return false; + } + + // In dev environments (localhost), allow any port on the same hostname + // This supports SSO flows that cross ports (e.g., localhost:3000 -> localhost:4000) + if (isDevEnvironment()) { + return true; + } + + // In production, require exact origin match (protocol + hostname + port) + return returnUrlObj.origin === window.location.origin; + } catch { + // Invalid URL + return false; + } +} + +export function normalizeUrlForCompare(url: string): string { + if (typeof window === "undefined") { + return url; + } + + try { + const parsed = new URL(url, window.location.origin); + let pathname = parsed.pathname; + if (pathname.length > 1 && pathname.endsWith("/")) { + pathname = pathname.slice(0, -1); + } + + const params = new URLSearchParams(parsed.search); + const sortedParams = new URLSearchParams(); + Array.from(params.entries()) + .sort(([a], [b]) => a.localeCompare(b)) + .forEach(([key, value]) => { + sortedParams.append(key, value); + }); + + const search = sortedParams.toString(); + const hash = parsed.hash || ""; + return `${parsed.origin}${pathname}${search ? `?${search}` : ""}${hash}`; + } catch { + return url; + } +} + +/** + * Gets and clears the return URL in one operation. + * Returns the validated return URL or null if invalid/not found. + * + * Priority: + * 1. If redirect_to param is valid, use it and clear cookie + * 2. If redirect_to param is invalid/missing, check cookie + * 3. Only clear cookie when we have a valid URL to return + */ +export function consumeReturnUrl(): string | null { + // Check URL param first + const paramUrl = getReturnUrlFromParams(); + if (paramUrl) { + if (isValidReturnUrl(paramUrl)) { + clearStoredReturnUrl(); + return paramUrl; + } + // Log rejected URLs in development for debugging + if (isDevEnvironment()) { + console.warn("[returnUrlUtils] Invalid return URL in params rejected:", paramUrl); + } + } + + // Fall back to cookie + const storedUrl = getStoredReturnUrl(); + if (storedUrl) { + if (isValidReturnUrl(storedUrl)) { + clearStoredReturnUrl(); + return storedUrl; + } + // Log rejected URLs in development for debugging + if (isDevEnvironment()) { + console.warn("[returnUrlUtils] Invalid return URL in cookie rejected:", storedUrl); + } + } + + // No valid URL found - don't clear cookie (nothing to clear or already invalid) + return null; +} diff --git a/ui/litellm-dashboard/src/utils/roles.ts b/ui/litellm-dashboard/src/utils/roles.ts index 580b4568c5..608a54ae14 100644 --- a/ui/litellm-dashboard/src/utils/roles.ts +++ b/ui/litellm-dashboard/src/utils/roles.ts @@ -31,3 +31,32 @@ export const isUserTeamAdminForSingleTeam = (teamMemberWithRoles: Member[] | nul } return teamMemberWithRoles.some((member) => member.user_id === userID && member.role === "admin"); }; + +export const formatUserRole = (userRole: string): string => { + if (!userRole) { + return "Undefined Role"; + } + switch (userRole.toLowerCase()) { + case "app_owner": + return "App Owner"; + case "demo_app_owner": + return "App Owner"; + case "app_admin": + return "Admin"; + case "proxy_admin": + return "Admin"; + case "proxy_admin_viewer": + return "Admin Viewer"; + case "org_admin": + return "Org Admin"; + case "internal_user": + return "Internal User"; + case "internal_user_viewer": + case "internal_viewer": // TODO:remove if deprecated + return "Internal Viewer"; + case "app_user": + return "App User"; + default: + return "Unknown Role"; + } +}; diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index c3c5ae5923..8b05def9ba 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -5,12 +5,13 @@ import { vi, describe, it, beforeEach, afterEach, expect } from "vitest"; /** ---------------------------- * Hoisted helpers for mocks (required by Vitest) * --------------------------- */ -const { stub, jwtDecodeMock } = vi.hoisted(() => { +const { stub, jwtDecodeMock, consumeReturnUrlMock } = vi.hoisted(() => { const React = require("react"); const stub = (name: string) => () => React.createElement("div", { "data-testid": name }); return { stub, jwtDecodeMock: vi.fn(), + consumeReturnUrlMock: vi.fn(), }; }); @@ -84,6 +85,14 @@ vi.mock("jwt-decode", () => ({ jwtDecode: (token: string) => jwtDecodeMock(token), })); +vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + consumeReturnUrl: consumeReturnUrlMock, + }; +}); + // Super-light stubs for all heavy components so rendering doesn't explode vi.mock("@/components/navbar", () => ({ default: stub("navbar") })); vi.mock("@/components/user_dashboard", () => ({ default: stub("user-dashboard") })); @@ -152,6 +161,7 @@ beforeEach(() => { // Fresh module state & DOM vi.clearAllMocks(); clearAllCookies(); + consumeReturnUrlMock.mockReturnValue(null); // Make location.replace spy-able to validate redirect delete (window as any).location; @@ -191,9 +201,11 @@ describe("CreateKeyPage auth behavior", () => { // Act render(); - // Assert: we eventually redirect to SSO login (single replace, not assign/href) + // Assert: we eventually redirect to SSO login with return URL (single replace, not assign/href) await waitFor(() => { - expect(window.location.replace).toHaveBeenCalledWith("https://example.com/ui/login"); + expect(window.location.replace).toHaveBeenCalledWith( + expect.stringContaining("https://example.com/ui/login?redirect_to=") + ); }); // And we attempted to clear the cookie (defensive deletion) @@ -235,4 +247,41 @@ describe("CreateKeyPage auth behavior", () => { expect(screen.getByTestId("navbar")).toBeInTheDocument(); }); }); + + it("should not redirect when return URL only differs by query order", async () => { + setCookie("token=validtoken"); + + jwtDecodeMock.mockImplementation((tok: string) => { + expect(tok).toBe("validtoken"); + return { + exp: Math.floor(Date.now() / 1000) + 60 * 60, + key: "accessKey-123", + user_role: "app_user", + user_email: "user@example.com", + login_method: "username_password", + premium_user: false, + auth_header_name: "x-litellm-auth", + user_id: "u_123", + }; + }); + + // Current URL has params in a different order + delete (window as any).location; + (window as any).location = { + ...originalLocation, + href: "http://localhost/ui?b=2&a=1", + origin: "http://localhost", + assign: vi.fn(), + replace: vi.fn(), + }; + + // Return URL has the same params in a different order + consumeReturnUrlMock.mockReturnValue("http://localhost/ui?a=1&b=2"); + + render(); + + await waitFor(() => { + expect(window.location.replace).not.toHaveBeenCalled(); + }); + }); });