set Secure flag on keycloak login cookies behind a TLS proxy

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()

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