AIP-38 Move token handling to axios interceptor

airflow/ui/src/layouts/BaseLayout.tsx

@@ -17,38 +17,20 @@
  * under the License.
  */
 import { Box } from "@chakra-ui/react";
-import { useEffect, type PropsWithChildren } from "react";
-import { Outlet, useSearchParams } from "react-router-dom";
-import { useLocalStorage } from "usehooks-ts";
+import type { PropsWithChildren } from "react";
+import { Outlet } from "react-router-dom";
 
 import { useConfig } from "src/queries/useConfig";
 
 import { Nav } from "./Nav";
 
-export const TOKEN_STORAGE_KEY = "token";
-
-export const TOKEN_QUERY_PARAM_NAME = "token";
-
 export const BaseLayout = ({ children }: PropsWithChildren) => {
   const instanceName = useConfig("instance_name");
 
   if (typeof instanceName === "string") {
     document.title = instanceName;
   }
 
-  const [searchParams, setSearchParams] = useSearchParams();
-  const paramToken = searchParams.get(TOKEN_QUERY_PARAM_NAME);
-
-  const [, setToken] = useLocalStorage<string | null>(TOKEN_STORAGE_KEY, paramToken);
-
-  useEffect(() => {
-    if (paramToken !== null) {
-      setToken(paramToken);
-      searchParams.delete(TOKEN_QUERY_PARAM_NAME);
-      setSearchParams(searchParams);
-    }
-  }, [paramToken, searchParams, setSearchParams, setToken]);
-
   return (
     <>
       <Nav />

airflow/ui/src/main.tsx

@@ -27,9 +27,9 @@ import { ColorModeProvider } from "src/context/colorMode";
 import { TimezoneProvider } from "src/context/timezone";
 import { router } from "src/router";
 
-import { TOKEN_STORAGE_KEY } from "./layouts/BaseLayout";
 import { queryClient } from "./queryClient";
 import { system } from "./theme";
+import { tokenHandler } from "./utils/tokenHandler";
 
 // redirect to login page if the API responds with unauthorized or forbidden errors
 axios.interceptors.response.use(
@@ -46,17 +46,7 @@ axios.interceptors.response.use(
   },
 );
 
-axios.interceptors.request.use((config) => {
-  const token: string | null = localStorage.getItem(TOKEN_STORAGE_KEY);
-
-  if (token !== null) {
-    // usehooks-ts stores a JSON.stringified version of values, we cannot use usehooks-ts here because we are outside of
-    // a react component. Therefore using bare localStorage.getItem and manually parsing the value.
-    config.headers.Authorization = `Bearer ${JSON.parse(token)}`;
-  }
-
-  return config;
-});
+axios.interceptors.request.use(tokenHandler);
 
 createRoot(document.querySelector("#root") as HTMLDivElement).render(
   <StrictMode>

airflow/ui/src/layouts/BaseLayout.test.tsx → airflow/ui/src/utils/tokenHandler.test.ts

@@ -16,45 +16,41 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { render } from "@testing-library/react";
-import { URLSearchParams } from "node:url";
-import * as reactRouterDom from "react-router-dom";
+import type { InternalAxiosRequestConfig } from "axios";
 import { afterEach, describe, it, vi, expect } from "vitest";
 
-import { Wrapper } from "src/utils/Wrapper";
-
-import { BaseLayout, TOKEN_QUERY_PARAM_NAME, TOKEN_STORAGE_KEY } from "./BaseLayout";
+import { TOKEN_QUERY_PARAM_NAME, TOKEN_STORAGE_KEY, tokenHandler } from "./tokenHandler";
 
 describe.each([
   { searchParams: new URLSearchParams({ token: "something" }) },
   { searchParams: new URLSearchParams({ param2: "someParam2", token: "else" }) },
   { searchParams: new URLSearchParams({}) },
-])("BaseLayout", ({ searchParams }) => {
+])("TokenFlow Interceptor", ({ searchParams }) => {
   it("Should read from the SearchParams, persist to the localStorage and remove from the SearchParams", () => {
-    const useSearchParamMock = vi.spyOn(reactRouterDom, "useSearchParams");
-
-    const setSearchParamsMock = vi.fn();
-
     const token = searchParams.get(TOKEN_QUERY_PARAM_NAME);
 
-    useSearchParamMock.mockImplementation(() => [searchParams, setSearchParamsMock]);
-
     const setItemMock = vi.spyOn(localStorage, "setItem");
 
-    render(<BaseLayout />, { wrapper: Wrapper });
+    vi.stubGlobal("location", { search: searchParams.toString() });
+
+    const headers = {};
+
+    const config = { headers };
 
-    expect(useSearchParamMock).toHaveBeenCalled();
+    const { headers: updatedHeaders } = tokenHandler(config as InternalAxiosRequestConfig);
 
     if (token === null) {
       expect(setItemMock).toHaveBeenCalledTimes(0);
     } else {
       expect(setItemMock).toHaveBeenCalledOnce();
-      expect(setItemMock).toHaveBeenCalledWith(TOKEN_STORAGE_KEY, JSON.stringify(token));
+      expect(setItemMock).toHaveBeenCalledWith(TOKEN_STORAGE_KEY, token);
       expect(searchParams).not.to.contains.keys(TOKEN_QUERY_PARAM_NAME);
+      expect(updatedHeaders).toEqual({ Authorization: `Bearer ${token}` });
     }
   });
 });
 
 afterEach(() => {
   vi.restoreAllMocks();
+  vi.unstubAllGlobals();
 });

airflow/ui/src/utils/tokenHandler.ts

@@ -0,0 +1,45 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import type { InternalAxiosRequestConfig } from "axios";
+
+export const TOKEN_STORAGE_KEY = "token";
+export const TOKEN_QUERY_PARAM_NAME = "token";
+
+export const tokenHandler = (config: InternalAxiosRequestConfig) => {
+  const searchParams = new URLSearchParams(globalThis.location.search);
+
+  const tokenUrl = searchParams.get(TOKEN_QUERY_PARAM_NAME);
+
+  let token: string | null;
+
+  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;
+  }
+
+  if (token !== null) {
+    config.headers.Authorization = `Bearer ${token}`;
+  }
+
+  return config;
+};

providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py

@@ -408,7 +408,7 @@ def security_manager(self) -> FabAirflowSecurityManagerOverride:
 
     def get_url_login(self, **kwargs) -> str:
         """Return the login page url."""
-        return f"{self.apiserver_endpoint}/auth/login"
+        return f"{self.apiserver_endpoint}/auth/login/"
 
     def get_url_logout(self):
         """Return the logout page url."""

providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py

@@ -556,7 +556,7 @@ class TestSecurityManager:
 
     def test_get_url_login(self, auth_manager):
         result = auth_manager.get_url_login()
-        assert result == "http://localhost:8080/auth/login"
+        assert result == "http://localhost:8080/auth/login/"
 
     @pytest.mark.db_test
     def test_get_url_logout_when_auth_view_not_defined(self, auth_manager_with_appbuilder):