Adds a state param into keycloak login (#64114)

Author: aritra24

providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py

@@ -19,6 +19,7 @@
 
 import json
 import logging
+import secrets
 from typing import cast
 from urllib.parse import quote, urljoin
 
@@ -53,15 +54,23 @@ class AuthManagerRefreshTokenExpiredException(Exception):  # type: ignore[no-red
 login_router = AirflowRouter(tags=["KeycloakAuthManagerLogin"])
 
 COOKIE_NAME_ID_TOKEN = "_id_token"
+COOKIE_NAME_OAUTH_STATE = "_oauth_state"
 
 
 @login_router.get("/login")
 def login(request: Request) -> RedirectResponse:
     """Initiate the authentication."""
     client = KeycloakAuthManager.get_keycloak_client()
     redirect_uri = request.url_for("login_callback")
-    auth_url = client.auth_url(redirect_uri=str(redirect_uri), scope="openid")
-    return RedirectResponse(auth_url)
+    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=""))
+    cookie_path = get_cookie_path()
+    response.set_cookie(
+        COOKIE_NAME_OAUTH_STATE, state, max_age=300, path=cookie_path, httponly=True, secure=secure
+    )
+    return response
 
 
 @login_router.get("/login_callback")
@@ -70,6 +79,10 @@ def login_callback(request: Request):
     code = request.query_params.get("code")
     if not code:
         return HTMLResponse("Missing code", status_code=400)
+    state_q = request.query_params.get("state", "")
+    state_c = request.cookies.get(COOKIE_NAME_OAUTH_STATE, "")
+    if not state_q or not state_c or not secrets.compare_digest(state_q, state_c):
+        return HTMLResponse("Invalid OAuth state parameter", status_code=403)
 
     client = KeycloakAuthManager.get_keycloak_client()
     redirect_uri = request.url_for("login_callback")

providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py

@@ -40,6 +40,7 @@ def test_login(self, mock_get_keycloak_client, client):
     def test_login_callback(self, mock_get_keycloak_client, mock_get_auth_manager, client):
         code = "code"
         token = "token"
+        state = "state"
         mock_keycloak_client = Mock()
         mock_keycloak_client.token.return_value = {
             "access_token": "access_token",
@@ -55,7 +56,9 @@ def test_login_callback(self, mock_get_keycloak_client, mock_get_auth_manager, c
         mock_get_auth_manager.return_value = mock_auth_manager
         mock_auth_manager.generate_jwt.return_value = token
         response = client.get(
-            AUTH_MANAGER_FASTAPI_APP_PREFIX + f"/login_callback?code={code}", follow_redirects=False
+            AUTH_MANAGER_FASTAPI_APP_PREFIX + f"/login_callback?code={code}&state={state}",
+            follow_redirects=False,
+            cookies={"_oauth_state": state},
         )
         mock_keycloak_client.token.assert_called_once_with(
             grant_type="authorization_code",