Set JWT token to localStorage from cookies

airflow/api_fastapi/auth/managers/base_auth_manager.py

@@ -72,6 +72,9 @@
 T = TypeVar("T", bound=BaseUser)
 
 
+COOKIE_NAME_JWT_TOKEN = "_token"
+
+
 class BaseAuthManager(Generic[T], LoggingMixin, metaclass=ABCMeta):
     """
     Class to derive in order to implement concrete auth managers.

airflow/api_fastapi/auth/managers/simple/routes/login.py

@@ -17,12 +17,11 @@
 
 from __future__ import annotations
 
-from urllib.parse import urljoin
-
 from fastapi import HTTPException, status
 from starlette.responses import RedirectResponse
 
 from airflow.api_fastapi.app import get_auth_manager
+from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
 from airflow.api_fastapi.auth.managers.simple.datamodels.login import LoginBody, LoginResponse
 from airflow.api_fastapi.auth.managers.simple.services.login import SimpleAuthManagerLogin
 from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser
@@ -62,8 +61,14 @@ def create_token_all_admins() -> RedirectResponse:
         username="Anonymous",
         role="ADMIN",
     )
-    url = urljoin(conf.get("api", "base_url"), f"?token={get_auth_manager().generate_jwt(user)}")
-    return RedirectResponse(url=url)
+
+    response = RedirectResponse(url=conf.get("api", "base_url"))
+    response.set_cookie(
+        COOKIE_NAME_JWT_TOKEN,
+        get_auth_manager().generate_jwt(user),
+        secure=True,
+    )
+    return response
 
 
 @login_router.post(

airflow/api_fastapi/auth/managers/simple/ui/package.json

@@ -18,7 +18,8 @@
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "react-hook-form": "^7.20.0",
-    "react-router-dom": "^6.26.2"
+    "react-router-dom": "^6.26.2",
+    "react-cookie": "^7.0.0"
   },
   "devDependencies": {
     "@7nohe/openapi-react-query-codegen": "^1.6.0",

airflow/api_fastapi/auth/managers/simple/ui/src/login/Login.tsx

@@ -25,6 +25,7 @@ import {LoginForm} from "src/login/LoginForm";
 import type {ApiError} from "openapi-gen/requests/core/ApiError";
 import type {LoginResponse, HTTPExceptionResponse, HTTPValidationError} from "openapi-gen/requests/types.gen";
 import { useSearchParams } from "react-router-dom";
+import { useCookies } from 'react-cookie';
 
 export type LoginBody = {
     username: string; password: string;
@@ -36,11 +37,15 @@ type ExpandedApiError = {
 
 export const Login = () => {
     const [searchParams, setSearchParams] = useSearchParams();
+    const [cookies, setCookie] = useCookies(['_token']);
 
     const onSuccess = (data: LoginResponse) => {
         // Redirect to appropriate page with the token
         const next = searchParams.get("next")
-        globalThis.location.replace(`${next ?? ""}?token=${data.jwt_token}`);
+
+        setCookie('_token', data.jwt_token, {path: "/", secure: true});
+
+        globalThis.location.replace(`${next ?? ""}`);
     }
     const {createToken, error: err, isPending, setError} = useCreateToken({onSuccess});
     const error = err as ExpandedApiError;

airflow/ui/src/utils/tokenHandler.test.ts

@@ -17,36 +17,34 @@
  * under the License.
  */
 import type { InternalAxiosRequestConfig } from "axios";
-import { afterEach, describe, it, vi, expect } from "vitest";
+import { afterEach, describe, it, vi, expect, beforeAll } from "vitest";
 
-import { TOKEN_QUERY_PARAM_NAME, TOKEN_STORAGE_KEY, tokenHandler } from "./tokenHandler";
+import { TOKEN_STORAGE_KEY, tokenHandler } from "./tokenHandler";
 
-describe.each([
-  { searchParams: new URLSearchParams({ token: "something" }) },
-  { searchParams: new URLSearchParams({ param2: "someParam2", token: "else" }) },
-  { searchParams: new URLSearchParams({}) },
-])("TokenFlow Interceptor", ({ searchParams }) => {
-  it("Should read from the SearchParams, persist to the localStorage and remove from the SearchParams", () => {
-    const token = searchParams.get(TOKEN_QUERY_PARAM_NAME);
+describe("TokenFlow Interceptor", () => {
+  beforeAll(() => {
+    Object.defineProperty(document, "cookie", {
+      writable: true,
+    });
+  });
 
-    const setItemMock = vi.spyOn(localStorage, "setItem");
+  it("Should read from the cookie, persist to the localStorage and remove from the cookie", () => {
+    const token = "test-token";
 
-    vi.stubGlobal("location", { search: searchParams.toString() });
+    document.cookie = `_token=${token};`;
+
+    const setItemMock = vi.spyOn(localStorage, "setItem");
 
     const headers = {};
 
     const config = { headers };
 
     const { headers: updatedHeaders } = tokenHandler(config as InternalAxiosRequestConfig);
 
-    if (token === null) {
-      expect(setItemMock).toHaveBeenCalledTimes(0);
-    } else {
-      expect(setItemMock).toHaveBeenCalledOnce();
-      expect(setItemMock).toHaveBeenCalledWith(TOKEN_STORAGE_KEY, token);
-      expect(searchParams).not.to.contains.keys(TOKEN_QUERY_PARAM_NAME);
-      expect(updatedHeaders).toEqual({ Authorization: `Bearer ${token}` });
-    }
+    expect(setItemMock).toHaveBeenCalledOnce();
+    expect(setItemMock).toHaveBeenCalledWith(TOKEN_STORAGE_KEY, token);
+    expect(updatedHeaders).toEqual({ Authorization: `Bearer ${token}` });
+    expect(document.cookie).toContain("_token=; expires=");
   });
 });
 

airflow/ui/src/utils/tokenHandler.ts

@@ -19,25 +19,29 @@
 import type { InternalAxiosRequestConfig } from "axios";
 
 export const TOKEN_STORAGE_KEY = "token";
-export const TOKEN_QUERY_PARAM_NAME = "token";
+const getTokenFromCookies = (): string | undefined => {
+  const cookies = document.cookie.split(";");
 
-export const tokenHandler = (config: InternalAxiosRequestConfig) => {
-  const searchParams = new URLSearchParams(globalThis.location.search);
-
-  const tokenUrl = searchParams.get(TOKEN_QUERY_PARAM_NAME);
+  for (const cookie of cookies) {
+    if (cookie.startsWith("_token=")) {
+      const [, token] = cookie.split("=");
 
-  let token: string | null;
+      if (token !== undefined) {
+        localStorage.setItem(TOKEN_STORAGE_KEY, token);
+        document.cookie = "_token=; expires=Thu, 01 Jan 2000 00:00:00 UTC; path=/;";
 
-  if (tokenUrl === null) {
-    token = localStorage.getItem(TOKEN_STORAGE_KEY);
-  } else {
-    localStorage.setItem(TOKEN_QUERY_PARAM_NAME, tokenUrl);
-    searchParams.delete(TOKEN_QUERY_PARAM_NAME);
-    globalThis.location.search = searchParams.toString();
-    token = tokenUrl;
+        return token;
+      }
+    }
   }
 
-  if (token !== null) {
+  return undefined;
+};
+
+export const tokenHandler = (config: InternalAxiosRequestConfig) => {
+  const token = localStorage.getItem(TOKEN_STORAGE_KEY) ?? getTokenFromCookies();
+
+  if (token !== undefined) {
     config.headers.Authorization = `Bearer ${token}`;
   }