diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py index 33ae5ca781a0e..9f2c80fe46a3d 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py @@ -65,7 +65,7 @@ def login(request: Request) -> RedirectResponse: state = secrets.token_urlsafe(32) auth_url = client.auth_url(redirect_uri=str(redirect_uri), scope="openid", state=state) response = RedirectResponse(auth_url) - secure = bool(conf.get("api", "ssl_cert", fallback="")) + secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback="")) cookie_path = get_cookie_path() response.set_cookie( COOKIE_NAME_OAUTH_STATE, state, max_age=300, path=cookie_path, httponly=True, secure=secure @@ -105,7 +105,7 @@ def login_callback(request: Request): token = get_auth_manager().generate_jwt(user) response = RedirectResponse(url=conf.get("api", "base_url", fallback="/"), status_code=303) - secure = bool(conf.get("api", "ssl_cert", fallback="")) + secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback="")) # In Airflow 3.1.1 authentication changes, front-end no longer handle the token # See https://github.com/apache/airflow/pull/55506 cookie_path = get_cookie_path() diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py index c4529150a3ba6..4f20715cc7675 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py +++ b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py @@ -78,6 +78,49 @@ def test_login_callback(self, mock_get_keycloak_client, mock_get_auth_manager, c assert response.cookies["_token"] == token assert response.cookies["_id_token"] == "id_token" + @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client") + def test_login_sets_secure_state_cookie_behind_tls_proxy(self, mock_get_keycloak_client, client): + mock_keycloak_client = Mock() + mock_keycloak_client.auth_url.return_value = "redirect_url" + mock_get_keycloak_client.return_value = mock_keycloak_client + response = client.get( + "https://testserver" + AUTH_MANAGER_FASTAPI_APP_PREFIX + "/login", + follow_redirects=False, + ) + set_cookies = response.headers.get_list("set-cookie") + assert any("_oauth_state" in c and "Secure" in c for c in set_cookies) + + @patch("airflow.providers.keycloak.auth_manager.routes.login.get_auth_manager") + @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client") + def test_login_callback_sets_secure_cookies_behind_tls_proxy( + self, mock_get_keycloak_client, mock_get_auth_manager, client + ): + state = "state" + mock_keycloak_client = Mock() + mock_keycloak_client.token.return_value = { + "access_token": "access_token", + "refresh_token": "refresh_token", + "id_token": "id_token", + } + mock_keycloak_client.userinfo.return_value = { + "sub": "sub", + "preferred_username": "preferred_username", + } + mock_get_keycloak_client.return_value = mock_keycloak_client + mock_auth_manager = Mock() + mock_get_auth_manager.return_value = mock_auth_manager + mock_auth_manager.generate_jwt.return_value = "token" + response = client.get( + "https://testserver" + + AUTH_MANAGER_FASTAPI_APP_PREFIX + + f"/login_callback?code=code&state={state}", + follow_redirects=False, + cookies={"_oauth_state": state}, + ) + set_cookies = response.headers.get_list("set-cookie") + assert any("_token=" in c and "Secure" in c for c in set_cookies) + assert any("_id_token=" in c and "Secure" in c for c in set_cookies) + def test_login_callback_without_code(self, client): response = client.get(AUTH_MANAGER_FASTAPI_APP_PREFIX + "/login_callback") assert response.status_code == 400