diff --git a/airflow/api_fastapi/auth/managers/base_auth_manager.py b/airflow/api_fastapi/auth/managers/base_auth_manager.py index e1ac86e090ba5..c8714e4103460 100644 --- a/airflow/api_fastapi/auth/managers/base_auth_manager.py +++ b/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. diff --git a/airflow/api_fastapi/auth/managers/simple/routes/login.py b/airflow/api_fastapi/auth/managers/simple/routes/login.py index da53e1f82b83f..b637a12373be3 100644 --- a/airflow/api_fastapi/auth/managers/simple/routes/login.py +++ b/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( diff --git a/airflow/api_fastapi/auth/managers/simple/ui/package.json b/airflow/api_fastapi/auth/managers/simple/ui/package.json index f4d766d9f0820..a26b9aad6b254 100644 --- a/airflow/api_fastapi/auth/managers/simple/ui/package.json +++ b/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", diff --git a/airflow/api_fastapi/auth/managers/simple/ui/pnpm-lock.yaml b/airflow/api_fastapi/auth/managers/simple/ui/pnpm-lock.yaml index bb2dd4fd45f63..d0c2f4685710d 100644 --- a/airflow/api_fastapi/auth/managers/simple/ui/pnpm-lock.yaml +++ b/airflow/api_fastapi/auth/managers/simple/ui/pnpm-lock.yaml @@ -10,7 +10,7 @@ importers: dependencies: '@chakra-ui/react': specifier: ^3.1.1 - version: 3.3.3(@emotion/react@11.14.0(react@18.3.1))(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + version: 3.3.3(@emotion/react@11.14.0(@types/react@19.0.10)(react@18.3.1))(react-dom@18.3.1(react@18.3.1))(react@18.3.1) '@tanstack/react-query': specifier: ^5.52.1 version: 5.64.1(react@18.3.1) @@ -20,6 +20,9 @@ importers: react: specifier: ^18.3.1 version: 18.3.1 + react-cookie: + specifier: ^7.0.0 + version: 7.2.2(react@18.3.1) react-dom: specifier: ^18.3.1 version: 18.3.1(react@18.3.1) @@ -38,7 +41,7 @@ importers: version: 6.6.3 '@testing-library/react': specifier: ^16.0.0 - version: 16.2.0(@testing-library/dom@10.4.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + version: 16.2.0(@testing-library/dom@10.4.0)(@types/react@19.0.10)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) '@vitejs/plugin-react-swc': specifier: ^3.7.0 version: 3.7.2(@swc/helpers@0.5.15)(vite@6.2.1(jiti@1.21.7)) @@ -663,15 +666,24 @@ packages: '@types/aria-query@5.0.4': resolution: {integrity: sha512-rfT93uj5s0PRL7EzccGMs3brplhcrghnDoV26NqKhCAS1hVo+WdNsPvE/yb6ilfr5hi2MEk6d5EWJTKdxg8jVw==} + '@types/cookie@0.6.0': + resolution: {integrity: sha512-4Kh9a6B2bQciAhf7FSuMRRkUWecJgJu9nPnx3yzpsfXX/c50REIqpHY4C82bXP90qrLtXtkDxTZosYO3UpOwlA==} + '@types/estree@1.0.6': resolution: {integrity: sha512-AYnb1nQyY49te+VRAVgmzfcgjYS91mY5P0TKUDCLEM+gNnA+3T6rWITXRLYCpahpqSQbN5cE+gHpnPyXjHWxcw==} + '@types/hoist-non-react-statics@3.3.6': + resolution: {integrity: sha512-lPByRJUer/iN/xa4qpyL0qmL11DqNW81iU/IG1S3uvRUq4oKagz8VCxZjiWkumgt66YT3vOdDgZ0o32sGKtCEw==} + '@types/json-schema@7.0.15': resolution: {integrity: sha512-5+fP8P8MFNC+AyZCDxrB2pkZFPGzqQWUzpSeuuVLvm8VMcorNYavBqoFcxK8bQz4Qsbn4oUEEem4wDLfcysGHA==} '@types/parse-json@4.0.2': resolution: {integrity: sha512-dISoDXWWQwUquiKsyZ4Ng+HX2KsPL7LyHKHQwgGFEA3IaKac4Obd+h2a/a6waisAoepJlBcx9paWqjA8/HVjCw==} + '@types/react@19.0.10': + resolution: {integrity: sha512-JuRQ9KXLEjaUNjTWpzuR231Z2WpIwczOkBEIvbHNCzQefFIT0L8IqE6NV6ULLyC1SI/i234JnDoMkfg+RjQj2g==} + '@vitejs/plugin-react-swc@3.7.2': resolution: {integrity: sha512-y0byko2b2tSVVf5Gpng1eEhX1OvPC7x8yns1Fx8jDzlJp4LS6CMkCPfLw47cjyoMrshQDoQw4qcgjsU9VvlCew==} peerDependencies: @@ -1070,6 +1082,10 @@ packages: convert-source-map@1.9.0: resolution: {integrity: sha512-ASFBup0Mz1uyiIjANan1jzLQami9z1PoYSZCiiYW2FczPbenXc45FZdBZLzOT+r6+iciuEModtmCti+hjaAk0A==} + cookie@0.7.2: + resolution: {integrity: sha512-yki5XnKuf750l50uGTllt6kKILY4nQ1eNIQatoXEByZ5dWgnKqbnqmTrBE5B4N7lrMJKQ2ytWMiTO2o0v6Ew/w==} + engines: {node: '>= 0.6'} + cosmiconfig@7.1.0: resolution: {integrity: sha512-AdmX6xUzdNASswsFtmwSt7Vj8po9IuqXm0UXz7QKPuEUmPB4XyjGfaAr2PSuELMwkRMVH1EpIkX5bTZGRB3eCA==} engines: {node: '>=10'} @@ -1702,6 +1718,11 @@ packages: rc9@2.1.2: resolution: {integrity: sha512-btXCnMmRIBINM2LDZoEmOogIZU7Qe7zn4BpomSKZ/ykbLObuBdvG+mFq11DL6fjH1DRwHhrlgtYWG96bJiC7Cg==} + react-cookie@7.2.2: + resolution: {integrity: sha512-e+hi6axHcw9VODoeVu8WyMWyoosa1pzpyjfvrLdF7CexfU+WSGZdDuRfHa4RJgTpfv3ZjdIpHE14HpYBieHFhg==} + peerDependencies: + react: '>= 16.3.0' + react-dom@18.3.1: resolution: {integrity: sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==} peerDependencies: @@ -1892,6 +1913,9 @@ packages: engines: {node: '>=0.8.0'} hasBin: true + universal-cookie@7.2.2: + resolution: {integrity: sha512-fMiOcS3TmzP2x5QV26pIH3mvhexLIT0HmPa3V7Q7knRfT9HG6kTwq02HZGLPw0sAOXrAmotElGRvTLCMbJsvxQ==} + uqr@0.1.2: resolution: {integrity: sha512-MJu7ypHq6QasgF5YRTjqscSzQp/W11zoUk6kvmlH+fmWEs63Y0Eib13hYFwAzagRJcVY8WVnlV+eBDUGMJ5IbA==} @@ -2157,11 +2181,11 @@ snapshots: '@babel/helper-string-parser': 7.25.9 '@babel/helper-validator-identifier': 7.25.9 - '@chakra-ui/react@3.3.3(@emotion/react@11.14.0(react@18.3.1))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': + '@chakra-ui/react@3.3.3(@emotion/react@11.14.0(@types/react@19.0.10)(react@18.3.1))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': dependencies: '@ark-ui/react': 4.8.0(react-dom@18.3.1(react@18.3.1))(react@18.3.1) '@emotion/is-prop-valid': 1.3.1 - '@emotion/react': 11.14.0(react@18.3.1) + '@emotion/react': 11.14.0(@types/react@19.0.10)(react@18.3.1) '@emotion/serialize': 1.3.3 '@emotion/use-insertion-effect-with-fallbacks': 1.2.0(react@18.3.1) '@emotion/utils': 1.4.2 @@ -2202,7 +2226,7 @@ snapshots: '@emotion/memoize@0.9.0': {} - '@emotion/react@11.14.0(react@18.3.1)': + '@emotion/react@11.14.0(@types/react@19.0.10)(react@18.3.1)': dependencies: '@babel/runtime': 7.26.9 '@emotion/babel-plugin': 11.13.5 @@ -2213,6 +2237,8 @@ snapshots: '@emotion/weak-memoize': 0.4.0 hoist-non-react-statics: 3.3.2 react: 18.3.1 + optionalDependencies: + '@types/react': 19.0.10 transitivePeerDependencies: - supports-color @@ -2585,12 +2611,14 @@ snapshots: lodash: 4.17.21 redent: 3.0.0 - '@testing-library/react@16.2.0(@testing-library/dom@10.4.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': + '@testing-library/react@16.2.0(@testing-library/dom@10.4.0)(@types/react@19.0.10)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': dependencies: '@babel/runtime': 7.26.0 '@testing-library/dom': 10.4.0 react: 18.3.1 react-dom: 18.3.1(react@18.3.1) + optionalDependencies: + '@types/react': 19.0.10 '@ts-morph/common@0.23.0': dependencies: @@ -2601,12 +2629,23 @@ snapshots: '@types/aria-query@5.0.4': {} + '@types/cookie@0.6.0': {} + '@types/estree@1.0.6': {} + '@types/hoist-non-react-statics@3.3.6': + dependencies: + '@types/react': 19.0.10 + hoist-non-react-statics: 3.3.2 + '@types/json-schema@7.0.15': {} '@types/parse-json@4.0.2': {} + '@types/react@19.0.10': + dependencies: + csstype: 3.1.3 + '@vitejs/plugin-react-swc@3.7.2(@swc/helpers@0.5.15)(vite@6.2.1(jiti@1.21.7))': dependencies: '@swc/core': 1.10.7(@swc/helpers@0.5.15) @@ -3286,6 +3325,8 @@ snapshots: convert-source-map@1.9.0: {} + cookie@0.7.2: {} + cosmiconfig@7.1.0: dependencies: '@types/parse-json': 4.0.2 @@ -3915,6 +3956,13 @@ snapshots: defu: 6.1.4 destr: 2.0.3 + react-cookie@7.2.2(react@18.3.1): + dependencies: + '@types/hoist-non-react-statics': 3.3.6 + hoist-non-react-statics: 3.3.2 + react: 18.3.1 + universal-cookie: 7.2.2 + react-dom@18.3.1(react@18.3.1): dependencies: loose-envify: 1.4.0 @@ -4096,6 +4144,11 @@ snapshots: uglify-js@3.19.3: optional: true + universal-cookie@7.2.2: + dependencies: + '@types/cookie': 0.6.0 + cookie: 0.7.2 + uqr@0.1.2: {} uri-js@4.4.1: diff --git a/airflow/api_fastapi/auth/managers/simple/ui/src/login/Login.tsx b/airflow/api_fastapi/auth/managers/simple/ui/src/login/Login.tsx index 915e70b55c5cc..e0321ad08e9ad 100644 --- a/airflow/api_fastapi/auth/managers/simple/ui/src/login/Login.tsx +++ b/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; diff --git a/airflow/ui/src/utils/tokenHandler.test.ts b/airflow/ui/src/utils/tokenHandler.test.ts index b0073f9560db0..8359bc193a1c5 100644 --- a/airflow/ui/src/utils/tokenHandler.test.ts +++ b/airflow/ui/src/utils/tokenHandler.test.ts @@ -17,21 +17,23 @@ * 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 = {}; @@ -39,14 +41,10 @@ describe.each([ 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="); }); }); diff --git a/airflow/ui/src/utils/tokenHandler.ts b/airflow/ui/src/utils/tokenHandler.ts index ab267cf96ecf6..b257dd06c08e6 100644 --- a/airflow/ui/src/utils/tokenHandler.ts +++ b/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}`; } diff --git a/docker_tests/test_docker_compose_quick_start.py b/docker_tests/test_docker_compose_quick_start.py index ccf4ef18fcd2e..452ad41aa8e2b 100644 --- a/docker_tests/test_docker_compose_quick_start.py +++ b/docker_tests/test_docker_compose_quick_start.py @@ -18,12 +18,10 @@ import json import os -import re import shlex from pprint import pprint from shutil import copyfile from time import sleep -from urllib.parse import parse_qs, urlparse import pytest import requests @@ -61,32 +59,17 @@ def get_jwt_token() -> str: """ # get csrf token from login page session = requests.Session() - get_login_form_response = session.get(f"http://{DOCKER_COMPOSE_HOST_PORT}/auth/login") - csrf_token = re.search( - r'', - get_login_form_response.text, - ) - assert csrf_token, "Failed to get csrf token from login page" - csrf_token_str = csrf_token.group(1) - assert csrf_token_str, "Failed to get csrf token from login page" - # login with form data + url = f"http://{DOCKER_COMPOSE_HOST_PORT}/auth/token" login_response = session.post( - f"http://{DOCKER_COMPOSE_HOST_PORT}/auth/login", - data={ + url, + json={ "username": AIRFLOW_WWW_USER_USERNAME, "password": AIRFLOW_WWW_USER_PASSWORD, - "csrf_token": csrf_token_str, }, ) - redirect_url = login_response.url - # ensure redirect_url is a string - redirect_url_str = str(redirect_url) if redirect_url is not None else "" - assert "/?token" in redirect_url_str, f"Login failed with redirect url {redirect_url_str}" - parsed_url = urlparse(redirect_url_str) - query_params = parse_qs(str(parsed_url.query)) - jwt_token_list = query_params.get("token") - jwt_token = jwt_token_list[0] if jwt_token_list else None - assert jwt_token, f"Failed to get JWT token from redirect url {redirect_url_str}" + jwt_token = login_response.json().get("jwt_token") + + assert jwt_token, f"Failed to get JWT token from redirect url {url} with status code {login_response}" return jwt_token diff --git a/docs/apache-airflow/core-concepts/auth-manager/index.rst b/docs/apache-airflow/core-concepts/auth-manager/index.rst index dc973f1d5e7b5..170a819985f0d 100644 --- a/docs/apache-airflow/core-concepts/auth-manager/index.rst +++ b/docs/apache-airflow/core-concepts/auth-manager/index.rst @@ -92,13 +92,30 @@ Some reasons you may want to write a custom auth manager include: * You'd like to use an auth manager that leverages an identity provider from your preferred cloud provider. * You have a private user management tool that is only available to you or your organization. - Authentication related BaseAuthManager methods ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ * ``get_user``: Return the signed-in user. * ``get_url_login``: Return the URL the user is redirected to for signing in. +JWT token management by auth managers +------------------------------------- +The auth manager is responsible of creating the JWT token and pass it to Airflow UI. The protocol to exchange the JWT +token between the auth manager and Airflow UI is using cookies. The auth manager needs to save the JWT token in a +cookie named ``_token`` before redirecting to the Airflow UI. The Airflow UI will then read the cookie, save it and +delete the cookie. + +.. code-block:: python + + from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN + + response = RedirectResponse(url="/") + response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=True) + return response + +.. note:: + Do not set the cookie parameter ``httponly`` to ``True``. Airflow UI needs to access the JWT token from the cookie. + Authorization related BaseAuthManager methods ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/kubernetes_tests/test_base.py b/kubernetes_tests/test_base.py index 8cad42de29618..03e361f0eaae2 100644 --- a/kubernetes_tests/test_base.py +++ b/kubernetes_tests/test_base.py @@ -25,7 +25,6 @@ from datetime import datetime, timezone from pathlib import Path from subprocess import check_call, check_output -from urllib.parse import parse_qs, urlparse import pytest import requests @@ -145,13 +144,9 @@ def _get_jwt_token(username: str, password: str) -> str: Note: API server is still using FAB Auth Manager. Steps: - 1. Get the login page to get the csrf token - - The csrf token is in the hidden input field with id "csrf_token" - 2. Login with the username and password + 1. Login with the username and password - Must use the same session to keep the csrf token session - 3. Extract the JWT token from the redirect url - - Expected to have a connection error - - The redirect url should have the JWT token as a query parameter + 2. Extract the JWT token from the auth/token url :param session: The session to use for the request :param username: The username to use for the login @@ -163,28 +158,14 @@ def _get_jwt_token(username: str, password: str) -> str: session = requests.Session() session.mount("http://", HTTPAdapter(max_retries=retry)) session.mount("https://", HTTPAdapter(max_retries=retry)) - get_login_form_response = session.get(f"http://{KUBERNETES_HOST_PORT}/auth/login") - csrf_token = re.search( - r'', - get_login_form_response.text, - ) - assert csrf_token, "Failed to get csrf token from login page" - csrf_token_str = csrf_token.group(1) - assert csrf_token_str, "Failed to get csrf token from login page" - # login with form data + url = f"http://{KUBERNETES_HOST_PORT}/auth/token" login_response = session.post( - f"http://{KUBERNETES_HOST_PORT}/auth/login", - data={"username": username, "password": password, "csrf_token": csrf_token_str}, + url, + json={"username": username, "password": password}, ) - redirect_url = login_response.url - # ensure redirect_url is a string - redirect_url_str = str(redirect_url) if redirect_url is not None else "" - assert "/?token" in redirect_url_str, f"Login failed with redirect url {redirect_url_str}" - parsed_url = urlparse(redirect_url_str) - query_params = parse_qs(str(parsed_url.query)) - jwt_token_list = query_params.get("token") - jwt_token = jwt_token_list[0] if jwt_token_list else None - assert jwt_token, f"Failed to get JWT token from redirect url {redirect_url_str}" + jwt_token = login_response.json().get("jwt_token") + + assert jwt_token, f"Failed to get JWT token from redirect url {url} with status code {login_response}" return jwt_token def _get_session_with_retries(self): diff --git a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/router/login.py b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/router/login.py index 293ed79b4a0cf..d00ef2e7febdc 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/router/login.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/router/login.py @@ -19,7 +19,6 @@ import logging from typing import Any -from urllib.parse import urljoin import anyio from fastapi import HTTPException, Request @@ -27,6 +26,7 @@ from starlette.responses import RedirectResponse from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX, get_auth_manager +from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN from airflow.api_fastapi.common.router import AirflowRouter from airflow.configuration import conf from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME @@ -80,8 +80,11 @@ def login_callback(request: Request): username=saml_auth.get_nameid(), email=attributes["email"][0] if "email" in attributes else None, ) - url = urljoin(conf.get("api", "base_url"), f"?token={get_auth_manager().generate_jwt(user)}") - return RedirectResponse(url=url, status_code=303) + url = conf.get("api", "base_url") + token = get_auth_manager().generate_jwt(user) + response = RedirectResponse(url=url, status_code=303) + response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=True) + return response def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth: diff --git a/providers/amazon/tests/unit/amazon/aws/auth_manager/router/test_login.py b/providers/amazon/tests/unit/amazon/aws/auth_manager/router/test_login.py index af763444cae29..07c61e32f6723 100644 --- a/providers/amazon/tests/unit/amazon/aws/auth_manager/router/test_login.py +++ b/providers/amazon/tests/unit/amazon/aws/auth_manager/router/test_login.py @@ -119,7 +119,8 @@ def test_login_callback_successful(self): ) assert response.status_code == 303 assert "location" in response.headers - assert response.headers["location"].startswith("http://localhost:8080/?token=") + assert "_token" in response.cookies + assert response.headers["location"].startswith("http://localhost:8080/") def test_login_callback_unsuccessful(self): with conf_vars( diff --git a/providers/fab/src/airflow/providers/fab/www/views.py b/providers/fab/src/airflow/providers/fab/www/views.py index 6c2cf076be242..43503502ff5c2 100644 --- a/providers/fab/src/airflow/providers/fab/www/views.py +++ b/providers/fab/src/airflow/providers/fab/www/views.py @@ -23,6 +23,7 @@ from flask import ( g, + make_response, redirect, render_template, request, @@ -31,6 +32,7 @@ from flask_appbuilder import IndexView, expose 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.configuration import conf from airflow.utils.net import get_hostname from airflow.version import version @@ -70,7 +72,10 @@ class FabIndexView(IndexView): def index(self): if g.user is not None and g.user.is_authenticated: token = get_auth_manager().generate_jwt(g.user) - return redirect(urljoin(conf.get("api", "base_url"), f"?token={token}"), code=302) + response = make_response(redirect(f"{conf.get('api', 'base_url')}", code=302)) + response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=True) + + return response else: return redirect(conf.get("api", "base_url"), code=302) diff --git a/tests/api_fastapi/auth/managers/simple/routes/test_login.py b/tests/api_fastapi/auth/managers/simple/routes/test_login.py index fdde46df07e64..70a39814f2db3 100644 --- a/tests/api_fastapi/auth/managers/simple/routes/test_login.py +++ b/tests/api_fastapi/auth/managers/simple/routes/test_login.py @@ -62,7 +62,7 @@ def test_create_token_all_admins(self, test_client): response = test_client.get("/auth/token", follow_redirects=False) assert response.status_code == 307 assert "location" in response.headers - assert response.headers["location"].startswith("http://localhost:8080/?token=") + assert response.cookies.get("_token") is not None def test_create_token_all_admins_config_disabled(self, test_client): response = test_client.get("/auth/token")