Repository: PrefectHQ/prefect

Issue:
Have `create_csrf_token` signal that CSRF protection is disabled.



Discussion:

Resolving patch:
diff --git a/src/prefect/client/base.py b/src/prefect/client/base.py
--- a/src/prefect/client/base.py
+++ b/src/prefect/client/base.py
@@ -361,9 +361,16 @@ async def _add_csrf_headers(self, request: Request):
             try:
                 token_response = await self.send(token_request)
             except PrefectHTTPStatusError as exc:
-                if exc.response.status_code == status.HTTP_404_NOT_FOUND:
-                    # The token endpoint is not available, likely an older
-                    # server version; disable CSRF support
+                old_server = exc.response.status_code == status.HTTP_404_NOT_FOUND
+                unconfigured_server = (
+                    exc.response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+                    and "CSRF protection is disabled." in exc.response.text
+                )
+
+                if old_server or unconfigured_server:
+                    # The token endpoint is either unavailable, suggesting an
+                    # older server, or CSRF protection is disabled. In either
+                    # case we should disable CSRF support.
                     self.enable_csrf_support = False
                     return
 
diff --git a/src/prefect/server/api/csrf_token.py b/src/prefect/server/api/csrf_token.py
--- a/src/prefect/server/api/csrf_token.py
+++ b/src/prefect/server/api/csrf_token.py
@@ -1,10 +1,12 @@
-from prefect._vendor.fastapi import Depends, Query
+from prefect._vendor.fastapi import Depends, Query, status
+from prefect._vendor.starlette.exceptions import HTTPException
 
 from prefect.logging import get_logger
 from prefect.server import models, schemas
 from prefect.server.database.dependencies import provide_database_interface
 from prefect.server.database.interface import PrefectDBInterface
 from prefect.server.utilities.server import PrefectRouter
+from prefect.settings import PREFECT_SERVER_CSRF_PROTECTION_ENABLED
 
 logger = get_logger("server.api")
 
@@ -17,7 +19,16 @@ async def create_csrf_token(
     client: str = Query(..., description="The client to create a CSRF token for"),
 ) -> schemas.core.CsrfToken:
     """Create or update a CSRF token for a client"""
+    if PREFECT_SERVER_CSRF_PROTECTION_ENABLED.value() is False:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="CSRF protection is disabled.",
+        )
+
     async with db.session_context(begin_transaction=True) as session:
-        return await models.csrf_token.create_or_update_csrf_token(
+        token = await models.csrf_token.create_or_update_csrf_token(
             session=session, client=client
         )
+        await models.csrf_token.delete_expired_tokens(session=session)
+
+    return token


Test patch:
diff --git a/tests/client/test_base_client.py b/tests/client/test_base_client.py
--- a/tests/client/test_base_client.py
+++ b/tests/client/test_base_client.py
@@ -57,6 +57,12 @@
     request=Request("a test request", "fake.url/fake/route"),
 )
 
+RESPONSE_CSRF_DISABLED = Response(
+    status.HTTP_422_UNPROCESSABLE_ENTITY,
+    json={"detail": "CSRF protection is disabled."},
+    request=Request("a test request", "fake.url/fake/route"),
+)
+
 RESPONSE_INVALID_TOKEN = Response(
     status_code=status.HTTP_403_FORBIDDEN,
     json={"detail": "Invalid CSRF token or client identifier."},
@@ -681,3 +687,12 @@ async def test_disables_csrf_support_404_token_endpoint(self):
             assert client.enable_csrf_support is True
             await client.post(url="fake.url/fake/route")
             assert client.enable_csrf_support is False
+
+    async def test_disables_csrf_support_422_csrf_disabled(self):
+        async with mocked_client(responses=[RESPONSE_CSRF_DISABLED, RESPONSE_200]) as (
+            client,
+            send,
+        ):
+            assert client.enable_csrf_support is True
+            await client.post(url="fake.url/fake/route")
+            assert client.enable_csrf_support is False
diff --git a/tests/server/api/test_csrf_token.py b/tests/server/api/test_csrf_token.py
--- a/tests/server/api/test_csrf_token.py
+++ b/tests/server/api/test_csrf_token.py
@@ -1,7 +1,19 @@
+from datetime import datetime, timedelta, timezone
+
+import pytest
+import sqlalchemy as sa
 from httpx import AsyncClient
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from prefect.server import models, schemas
+from prefect.server.database.interface import PrefectDBInterface
+from prefect.settings import PREFECT_SERVER_CSRF_PROTECTION_ENABLED, temporary_settings
+
+
+@pytest.fixture(autouse=True)
+def enable_csrf_protection():
+    with temporary_settings({PREFECT_SERVER_CSRF_PROTECTION_ENABLED: True}):
+        yield
 
 
 async def test_can_get_csrf_token(client: AsyncClient, session: AsyncSession):
@@ -13,6 +25,7 @@ async def test_can_get_csrf_token(client: AsyncClient, session: AsyncSession):
         session=session, client="client123"
     )
 
+    assert db_token
     assert token.client == "client123"
     assert token.token == db_token.token
 
@@ -31,3 +44,41 @@ async def test_client_param_required(client: AsyncClient):
         "exception_message": "Invalid request received.",
         "request_body": None,
     }
+
+
+async def test_422_when_csrf_protection_disabled(client: AsyncClient):
+    with temporary_settings({PREFECT_SERVER_CSRF_PROTECTION_ENABLED: False}):
+        response = await client.get("/csrf-token?client=client123")
+        assert response.status_code == 422
+        assert response.json() == {"detail": "CSRF protection is disabled."}
+
+
+async def test_deletes_expired_tokens(
+    db: PrefectDBInterface, session: AsyncSession, client: AsyncClient
+):
+    # Create some tokens
+    for i in range(5):
+        await models.csrf_token.create_or_update_csrf_token(
+            session=session, client=f"client{i}"
+        )
+
+    # Update some of them to be expired
+    await session.execute(
+        sa.update(db.CsrfToken)
+        .where(db.CsrfToken.client.in_(["client0", "client1"]))
+        .values(expiration=datetime.now(timezone.utc) - timedelta(days=1))
+    )
+
+    await session.commit()
+
+    response = await client.get("/csrf-token?client=client123")
+    assert response.status_code == 200
+
+    all_tokens = (await session.execute(sa.select(db.CsrfToken))).scalars().all()
+
+    # Five tokens were created upfront, one was created from the client.get
+    # request, and two of those were expired. (5 + 1) - 2 = 4
+    assert len(all_tokens) == 4
+
+    assert "client0" not in [t.client for t in all_tokens]
+    assert "client1" not in [t.client for t in all_tokens]
diff --git a/tests/server/models/test_csrf_token.py b/tests/server/models/test_csrf_token.py
--- a/tests/server/models/test_csrf_token.py
+++ b/tests/server/models/test_csrf_token.py
@@ -124,19 +124,8 @@ async def test_can_delete_expired_tokens(
 
         await models.csrf_token.delete_expired_tokens(session=session)
 
-        removed = [
-            await models.csrf_token.read_token_for_client(
-                session=session, client=f"client{i}"
-            )
-            for i in range(2)
-        ]
-
-        existing = [
-            await models.csrf_token.read_token_for_client(
-                session=session, client=f"client{i}"
-            )
-            for i in range(2, 5)
-        ]
+        all_tokens = (await session.execute(sa.select(db.CsrfToken))).scalars().all()
 
-        assert removed == [None, None]
-        assert all(token is not None for token in existing)
+        assert len(all_tokens) == 3
+        assert "client0" not in [t.client for t in all_tokens]
+        assert "client1" not in [t.client for t in all_tokens]