CVE-2023-6022: Cross-Site Request Forgery (CSRF) (CWE-352)
value: Cross-Site Request Forgery (CSRF) in GitHub repository prefecthq/prefect prior to 2.16.5.
CWE description: The web application does not, or can not, sufficiently verify whether a well-formed, valid, consistent request was intentionally provided by the user who submitted the request.
CVSS v3 base score: 8.8
Published: 2023-11-16T17:15Z
Repository: https://github.com/prefecthq/prefect
Language: Python
Files changed: src/prefect/client/base.py, src/prefect/client/cloud.py, src/prefect/client/orchestration.py, src/prefect/client/schemas/objects.py, src/prefect/settings.py, tests/cli/test_worker.py, tests/client/test_base_client.py, tests/client/test_prefect_client.py, tests/test_deployments.py
Commit date: 2024-03-18 12:29:27 -0400
Commit message: Add CSRF support to orchestration client (#12314)


Diff: diff --git a/src/prefect/client/base.py b/src/prefect/client/base.py
index a61b10b0d..ffd32326c 100644
--- a/src/prefect/client/base.py
+++ b/src/prefect/client/base.py
@@ -1,32 +1,35 @@
 import copy
 import sys
 import threading
+import uuid
 from collections import defaultdict
 from contextlib import asynccontextmanager
-from functools import partial
+from datetime import datetime, timezone
 from typing import (
     Any,
     AsyncGenerator,
     Awaitable,
     Callable,
     Dict,
     MutableMapping,
+    Optional,
     Protocol,
     Set,
     Tuple,
     Type,
     runtime_checkable,
 )
 
 import anyio
 import httpx
 from asgi_lifespan import LifespanManager
-from httpx import HTTPStatusError, Response
+from httpx import HTTPStatusError, Request, Response
 from prefect._vendor.starlette import status
 from typing_extensions import Self
 
+from prefect.client.schemas.objects import CsrfToken
 from prefect.exceptions import PrefectHTTPStatusError
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_CLIENT_MAX_RETRIES,
     PREFECT_CLIENT_RETRY_EXTRA_CODES,
@@ -186,13 +189,24 @@ class PrefectHttpxClient(httpx.AsyncClient):
 
     For more details on rate limit headers, see:
     [Configuring Cloudflare Rate Limiting](https://support.cloudflare.com/hc/en-us/articles/115001635128-Configuring-Rate-Limiting-from-UI)
     """
 
+    def __init__(self, *args, enable_csrf_support: bool = False, **kwargs):
+        self.enable_csrf_support: bool = enable_csrf_support
+        self.csrf_token: Optional[str] = None
+        self.csrf_token_expiration: Optional[datetime] = None
+        self.csrf_client_id: uuid.UUID = uuid.uuid4()
+
+        super().__init__(*args, **kwargs)
+
     async def _send_with_retry(
         self,
-        request: Callable,
+        request: Request,
+        send: Callable[[Request], Awaitable[Response]],
+        send_args: Tuple,
+        send_kwargs: Dict,
         retry_codes: Set[int] = set(),
         retry_exceptions: Tuple[Exception, ...] = tuple(),
     ):
         """
         Send a request and retry it if it fails.
@@ -205,25 +219,38 @@ class PrefectHttpxClient(httpx.AsyncClient):
         exponential backoff if a retry header is not provided.
         """
         try_count = 0
         response = None
 
+        is_change_request = request.method.lower() in {"post", "put", "patch", "delete"}
+
+        if self.enable_csrf_support and is_change_request:
+            await self._add_csrf_headers(request=request)
+
         while try_count <= PREFECT_CLIENT_MAX_RETRIES.value():
             try_count += 1
             retry_seconds = None
             exc_info = None
 
             try:
-                response = await request()
+                response = await send(request, *send_args, **send_kwargs)
             except retry_exceptions:  # type: ignore
                 if try_count > PREFECT_CLIENT_MAX_RETRIES.value():
                     raise
                 # Otherwise, we will ignore this error but capture the info for logging
                 exc_info = sys.exc_info()
             else:
-                # We got a response; return immediately if it is not retryable
-                if response.status_code not in retry_codes:
+                # We got a response; check if it's a CSRF error, otherwise
+                # return immediately if it is not retryable
+                if (
+                    response.status_code == status.HTTP_403_FORBIDDEN
+                    and "Invalid CSRF token" in response.text
+                ):
+                    # We got a CSRF error, clear the token and try again
+                    self.csrf_token = None
+                    await self._add_csrf_headers(request)
+                elif response.status_code not in retry_codes:
                     return response
 
                 if "Retry-After" in response.headers:
                     retry_seconds = float(response.headers["Retry-After"])
 
@@ -266,23 +293,28 @@ class PrefectHttpxClient(httpx.AsyncClient):
         ), "Retry handling ended without response or exception"
 
         # We ran out of retries, return the failed response
         return response
 
-    async def send(self, *args, **kwargs) -> Response:
+    async def send(self, request: Request, *args, **kwargs) -> Response:
         """
         Send a request with automatic retry behavior for the following status codes:
 
+        - 403 Forbidden, if the request failed due to CSRF protection
+        - 408 Request Timeout
         - 429 CloudFlare-style rate limiting
         - 502 Bad Gateway
         - 503 Service unavailable
+        - Any additional status codes provided in `PREFECT_CLIENT_RETRY_EXTRA_CODES`
         """
 
-        api_request = partial(super().send, *args, **kwargs)
-
+        super_send = super().send
         response = await self._send_with_retry(
-            request=api_request,
+            request=request,
+            send=super_send,
+            send_args=args,
+            send_kwargs=kwargs,
             retry_codes={
                 status.HTTP_429_TOO_MANY_REQUESTS,
                 status.HTTP_503_SERVICE_UNAVAILABLE,
                 status.HTTP_502_BAD_GATEWAY,
                 status.HTTP_408_REQUEST_TIMEOUT,
@@ -310,5 +342,36 @@ class PrefectHttpxClient(httpx.AsyncClient):
         # NOTE: We may want to remove this and handle responses per route in the
         #       `PrefectClient`
         response.raise_for_status()
 
         return response
+
+    async def _add_csrf_headers(self, request: Request):
+        now = datetime.now(timezone.utc)
+
+        if not self.enable_csrf_support:
+            return
+
+        if not self.csrf_token or (
+            self.csrf_token_expiration and now > self.csrf_token_expiration
+        ):
+            token_request = self.build_request(
+                "GET", f"/csrf-token?client={self.csrf_client_id}"
+            )
+
+            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
+                    self.enable_csrf_support = False
+                    return
+
+                raise
+
+            token: CsrfToken = CsrfToken.parse_obj(token_response.json())
+            self.csrf_token = token.token
+            self.csrf_token_expiration = token.expiration
+
+        request.headers["Prefect-Csrf-Token"] = self.csrf_token
+        request.headers["Prefect-Csrf-Client"] = str(self.csrf_client_id)
diff --git a/src/prefect/client/cloud.py b/src/prefect/client/cloud.py
index 151f18447..479edf5b7 100644
--- a/src/prefect/client/cloud.py
+++ b/src/prefect/client/cloud.py
@@ -70,11 +70,11 @@ class CloudClient:
         httpx_settings["headers"].setdefault("Authorization", f"Bearer {api_key}")
 
         httpx_settings.setdefault("base_url", host)
         if not PREFECT_UNIT_TEST_MODE.value():
             httpx_settings.setdefault("follow_redirects", True)
-        self._client = PrefectHttpxClient(**httpx_settings)
+        self._client = PrefectHttpxClient(**httpx_settings, enable_csrf_support=False)
 
     async def api_healthcheck(self):
         """
         Attempts to connect to the Cloud API and raises the encountered exception if not
         successful.
diff --git a/src/prefect/client/orchestration.py b/src/prefect/client/orchestration.py
index e2821364b..8f766ed53 100644
--- a/src/prefect/client/orchestration.py
+++ b/src/prefect/client/orchestration.py
@@ -134,10 +134,11 @@ from prefect.settings import (
     PREFECT_API_ENABLE_HTTP2,
     PREFECT_API_KEY,
     PREFECT_API_REQUEST_TIMEOUT,
     PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
     PREFECT_API_URL,
+    PREFECT_CLIENT_CSRF_SUPPORT_ENABLED,
     PREFECT_CLOUD_API_URL,
     PREFECT_UNIT_TEST_MODE,
 )
 from prefect.utilities.collections import AutoEnum
 
@@ -314,11 +315,19 @@ class PrefectClient:
             ),
         )
 
         if not PREFECT_UNIT_TEST_MODE:
             httpx_settings.setdefault("follow_redirects", True)
-        self._client = PrefectHttpxClient(**httpx_settings)
+
+        enable_csrf_support = (
+            self.server_type != ServerType.CLOUD
+            and PREFECT_CLIENT_CSRF_SUPPORT_ENABLED.value()
+        )
+
+        self._client = PrefectHttpxClient(
+            **httpx_settings, enable_csrf_support=enable_csrf_support
+        )
         self._loop = None
 
         # See https://www.python-httpx.org/advanced/#custom-transports
         #
         # If we're using an HTTP/S client (not the ephemeral client), adjust the
diff --git a/src/prefect/client/schemas/objects.py b/src/prefect/client/schemas/objects.py
index b60258092..cc545109c 100644
--- a/src/prefect/client/schemas/objects.py
+++ b/src/prefect/client/schemas/objects.py
@@ -1630,5 +1630,18 @@ class GlobalConcurrencyLimit(ObjectBaseModel):
         description=(
             "Controls the rate at which slots are released when the concurrency limit"
             " is used as a rate limit."
         ),
     )
+
+
+class CsrfToken(ObjectBaseModel):
+    token: str = Field(
+        default=...,
+        description="The CSRF token",
+    )
+    client: str = Field(
+        default=..., description="The client id associated with the CSRF token"
+    )
+    expiration: DateTimeTZ = Field(
+        default=..., description="The expiration time of the CSRF token"
+    )
diff --git a/src/prefect/settings.py b/src/prefect/settings.py
index 31142f8c5..70fbb8521 100644
--- a/src/prefect/settings.py
+++ b/src/prefect/settings.py
@@ -655,10 +655,25 @@ PREFECT_CLIENT_RETRY_EXTRA_CODES = Setting(
 A comma-separated list of extra HTTP status codes to retry on. Defaults to an empty string.
 429, 502 and 503 are always retried. Please note that not all routes are idempotent and retrying
 may result in unexpected behavior.
 """
 
+PREFECT_CLIENT_CSRF_SUPPORT_ENABLED = Setting(bool, default=True)
+"""
+Determines if CSRF token handling is active in the Prefect client for API
+requests.
+
+When enabled (`True`), the client automatically manages CSRF tokens by
+retrieving, storing, and including them in applicable state-changing requests
+(POST, PUT, PATCH, DELETE) to the API.
+
+Disabling this setting (`False`) means the client will not handle CSRF tokens,
+which might be suitable for environments where CSRF protection is disabled.
+
+Defaults to `True`, ensuring CSRF protection is enabled by default.
+"""
+
 PREFECT_CLOUD_API_URL = Setting(
     str,
     default="https://api.prefect.cloud/api",
     value_callback=check_for_deprecated_cloud_url,
 )
@@ -1206,15 +1221,35 @@ greater than the load balancer's idle timeout.
 Note this setting only applies when calling `prefect server start`; if hosting the
 API with another tool you will need to configure this there instead.
 """
 
 PREFECT_SERVER_CSRF_PROTECTION_ENABLED = Setting(bool, default=False)
-"""Whether or not to protect the API from CSRF attacks. Experimental and
-currently defaults to `False`."""
+"""
+Controls the activation of CSRF protection for the Prefect server API.
+
+When enabled (`True`), the server enforces CSRF validation checks on incoming
+state-changing requests (POST, PUT, PATCH, DELETE), requiring a valid CSRF
+token to be included in the request headers or body. This adds a layer of
+security by preventing unauthorized or malicious sites from making requests on
+behalf of authenticated users.
+
+It is recommended to enable this setting in production environments where the
+API is exposed to web clients to safeguard against CSRF attacks.
+
+Note: Enabling this setting requires corresponding support in the client for
+CSRF token management. See PREFECT_CLIENT_CSRF_SUPPORT_ENABLED for more.
+"""
 
 PREFECT_SERVER_CSRF_TOKEN_EXPIRATION = Setting(timedelta, default=timedelta(hours=1))
-"""How long a CSRF token is valid for. Defaults to 1 hour."""
+"""
+Specifies the duration for which a CSRF token remains valid after being issued
+by the server.
+
+The default expiration time is set to 1 hour, which offers a reasonable
+compromise. Adjust this setting based on your specific security requirements
+and usage patterns.
+"""
 
 PREFECT_UI_ENABLED = Setting(
     bool,
     default=True,
 )
diff --git a/tests/cli/test_worker.py b/tests/cli/test_worker.py
index 44ba752d0..da41f171b 100644
--- a/tests/cli/test_worker.py
+++ b/tests/cli/test_worker.py
@@ -65,10 +65,11 @@ async def kubernetes_work_pool(prefect_client: PrefectClient):
     )
 
     with respx.mock(
         assert_all_mocked=False, base_url=PREFECT_API_URL.value()
     ) as respx_mock:
+        respx_mock.get("/csrf-token", params={"client": ANY}).pass_through()
         respx_mock.route(path__startswith="/work_pools/").pass_through()
         respx_mock.route(path__startswith="/flow_runs/").pass_through()
         respx_mock.get("/collections/views/aggregate-worker-metadata").mock(
             return_value=httpx.Response(
                 200,
diff --git a/tests/client/test_base_client.py b/tests/client/test_base_client.py
index f0fce14a6..1c64cdb81 100644
--- a/tests/client/test_base_client.py
+++ b/tests/client/test_base_client.py
@@ -1,22 +1,28 @@
-from unittest.mock import call
+from contextlib import asynccontextmanager
+from datetime import datetime, timedelta, timezone
+from typing import AsyncGenerator, List, Tuple
+from unittest import mock
 
 import httpx
 import pytest
 from httpx import AsyncClient, Request, Response
 from prefect._vendor.starlette import status
 
 from prefect.client.base import PrefectHttpxClient, PrefectResponse
+from prefect.client.schemas.objects import CsrfToken
 from prefect.exceptions import PrefectHTTPStatusError
 from prefect.settings import (
     PREFECT_CLIENT_MAX_RETRIES,
     PREFECT_CLIENT_RETRY_EXTRA_CODES,
     PREFECT_CLIENT_RETRY_JITTER_FACTOR,
     temporary_settings,
 )
 from prefect.testing.utilities import AsyncMock
 
+now = datetime.now(timezone.utc)
+
 RESPONSE_429_RETRY_AFTER_0 = Response(
     status.HTTP_429_TOO_MANY_REQUESTS,
     headers={"Retry-After": "0"},
     request=Request("a test request", "fake.url/fake/route"),
 )
@@ -30,10 +36,35 @@ RESPONSE_429_RETRY_AFTER_MISSING = Response(
 RESPONSE_200 = Response(
     status.HTTP_200_OK,
     request=Request("a test request", "fake.url/fake/route"),
 )
 
+RESPONSE_CSRF = Response(
+    status.HTTP_200_OK,
+    json=CsrfToken(
+        client="test_client", token="test_token", expiration=now + timedelta(days=1)
+    ).dict(json_compatible=True, exclude_unset=True),
+    request=Request("a test request", "fake.url/fake/route"),
+)
+
+RESPONSE_400 = Response(
+    status.HTTP_400_BAD_REQUEST,
+    json={"detail": "You done bad things"},
+    request=Request("a test request", "fake.url/fake/route"),
+)
+
+RESPONSE_404 = Response(
+    status.HTTP_404_NOT_FOUND,
+    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."},
+    request=Request("a test request", "fake.url/fake/route"),
+)
+
 
 @pytest.fixture
 def disable_jitter():
     with temporary_settings({PREFECT_CLIENT_RETRY_JITTER_FACTOR: 0}):
         yield
@@ -341,11 +372,11 @@ class TestPrefectHttpxClient:
             async with client:
                 response = await client.post(
                     url="fake.url/fake/route", data={"evenmorefake": "data"}
                 )
         assert response.status_code == status.HTTP_200_OK
-        mock_anyio_sleep.assert_has_awaits([call(2), call(4), call(8)])
+        mock_anyio_sleep.assert_has_awaits([mock.call(2), mock.call(4), mock.call(8)])
 
     @pytest.mark.usefixtures("disable_jitter")
     async def test_prefect_httpx_client_respects_retry_header_per_response(
         self, mock_anyio_sleep, monkeypatch
     ):
@@ -368,11 +399,13 @@ class TestPrefectHttpxClient:
             async with client:
                 response = await client.post(
                     url="fake.url/fake/route", data={"evenmorefake": "data"}
                 )
         assert response.status_code == status.HTTP_200_OK
-        mock_anyio_sleep.assert_has_awaits([call(5), call(0), call(10), call(2.0)])
+        mock_anyio_sleep.assert_has_awaits(
+            [mock.call(5), mock.call(0), mock.call(10), mock.call(2.0)]
+        )
 
     async def test_prefect_httpx_client_adds_jitter_with_retry_header(
         self, monkeypatch, mock_anyio_sleep
     ):
         base_client_send = AsyncMock()
@@ -431,11 +464,11 @@ class TestPrefectHttpxClient:
                 response = await client.post(
                     url="fake.url/fake/route", data={"evenmorefake": "data"}
                 )
         assert response.status_code == status.HTTP_200_OK
         mock_anyio_sleep.assert_has_awaits(
-            [call(pytest.approx(n, rel=0.2)) for n in [2, 4, 8]]
+            [mock.call(pytest.approx(n, rel=0.2)) for n in [2, 4, 8]]
         )
 
     async def test_prefect_httpx_client_does_not_retry_other_exceptions(
         self, mock_anyio_sleep, monkeypatch
     ):
@@ -488,5 +521,163 @@ class TestPrefectHttpxClient:
                 await client.post(
                     url="fake.url/fake/route", data={"evenmorefake": "data"}
                 )
         expected = "Response: {'extra_info': [{'message': 'a test error message'}]}"
         assert expected in str(exc.exconly())
+
+
+@asynccontextmanager
+async def mocked_client(
+    responses: List[Response],
+) -> AsyncGenerator[Tuple[PrefectHttpxClient, mock.AsyncMock], None]:
+    with mock.patch("httpx.AsyncClient.send", autospec=True) as send:
+        send.side_effect = responses
+        client = PrefectHttpxClient(enable_csrf_support=True)
+        async with client:
+            try:
+                yield client, send
+            finally:
+                pass
+
+
+class TestCsrfSupport:
+    async def test_no_csrf_headers_not_change_request(self):
+        async with mocked_client(responses=[RESPONSE_200]) as (client, send):
+            await client.get(url="fake.url/fake/route")
+
+        request = send.call_args[0][1]
+        assert isinstance(request, httpx.Request)
+
+        assert "Prefect-Csrf-Token" not in request.headers
+        assert "Prefect-Csrf-Client" not in request.headers
+
+    @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
+    async def test_csrf_headers_on_change_request(self, method: str):
+        async with mocked_client(responses=[RESPONSE_CSRF, RESPONSE_200]) as (
+            client,
+            send,
+        ):
+            await getattr(client, method)(url="fake.url/fake/route")
+
+        assert send.await_count == 2
+
+        # The first call should be for the CSRF token
+        request = send.call_args_list[0][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.method == "GET"
+        assert request.url == httpx.URL(
+            f"/csrf-token?client={str(client.csrf_client_id)}"
+        )
+
+        # The second call should be for the actual request
+        request = send.call_args_list[1][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.method == method.upper()
+        assert request.url == httpx.URL("/fake.url/fake/route")
+        assert request.headers["Prefect-Csrf-Token"] == "test_token"
+        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)
+
+    async def test_refreshes_token_on_csrf_403(self):
+        async with mocked_client(
+            responses=[
+                RESPONSE_CSRF,
+                RESPONSE_INVALID_TOKEN,
+                RESPONSE_CSRF,
+                RESPONSE_200,
+            ]
+        ) as (
+            client,
+            send,
+        ):
+            await client.post(url="fake.url/fake/route")
+
+        assert send.await_count == 4
+
+        # The first call should be for the CSRF token
+        request = send.call_args_list[0][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.method == "GET"
+        assert request.url == httpx.URL(
+            f"/csrf-token?client={str(client.csrf_client_id)}"
+        )
+
+        # The second call should be for the actual request
+        request = send.call_args_list[1][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.url == httpx.URL("/fake.url/fake/route")
+        assert request.headers["Prefect-Csrf-Token"] == "test_token"
+        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)
+
+        # The third call should be a refresh of the CSRF token
+        request = send.call_args_list[0][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.method == "GET"
+        assert request.url == httpx.URL(
+            f"/csrf-token?client={str(client.csrf_client_id)}"
+        )
+
+        # The fourth call should be for the actual request
+        request = send.call_args_list[1][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.url == httpx.URL("/fake.url/fake/route")
+        assert request.headers["Prefect-Csrf-Token"] == "test_token"
+        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)
+
+    async def test_does_not_refresh_csrf_token_not_expired(self):
+        async with mocked_client(responses=[RESPONSE_200]) as (
+            client,
+            send,
+        ):
+            client.csrf_token = "fresh_token"
+            client.csrf_token_expiration = now + timedelta(days=1)
+            await client.post(url="fake.url/fake/route")
+
+        assert send.await_count == 1
+
+        request = send.call_args_list[0][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.url == httpx.URL("/fake.url/fake/route")
+        assert request.headers["Prefect-Csrf-Token"] == "fresh_token"
+        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)
+
+    async def test_does_refresh_csrf_token_when_expired(self):
+        async with mocked_client(responses=[RESPONSE_CSRF, RESPONSE_200]) as (
+            client,
+            send,
+        ):
+            client.csrf_token = "old_token"
+            client.csrf_token_expiration = now - timedelta(days=1)
+            await client.post(url="fake.url/fake/route")
+
+        assert send.await_count == 2
+
+        # The first call should be for the CSRF token
+        request = send.call_args_list[0][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.method == "GET"
+        assert request.url == httpx.URL(
+            f"/csrf-token?client={str(client.csrf_client_id)}"
+        )
+
+        # The second call should be for the actual request
+        request = send.call_args_list[1][0][1]
+        assert isinstance(request, httpx.Request)
+        assert request.url == httpx.URL("/fake.url/fake/route")
+        assert request.headers["Prefect-Csrf-Token"] == "test_token"
+        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)
+
+    async def test_raises_exception_bad_csrf_token_response(self):
+        async with mocked_client(responses=[RESPONSE_400]) as (
+            client,
+            _,
+        ):
+            with pytest.raises(PrefectHTTPStatusError):
+                await client.post(url="fake.url/fake/route")
+
+    async def test_disables_csrf_support_404_token_endpoint(self):
+        async with mocked_client(responses=[RESPONSE_404, 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/client/test_prefect_client.py b/tests/client/test_prefect_client.py
index b7a0a19eb..4683451bb 100644
--- a/tests/client/test_prefect_client.py
+++ b/tests/client/test_prefect_client.py
@@ -74,10 +74,11 @@ from prefect.server.api.server import create_app
 from prefect.settings import (
     PREFECT_API_DATABASE_MIGRATE_ON_START,
     PREFECT_API_KEY,
     PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
     PREFECT_API_URL,
+    PREFECT_CLIENT_CSRF_SUPPORT_ENABLED,
     PREFECT_CLOUD_API_URL,
     PREFECT_EXPERIMENTAL_ENABLE_FLOW_RUN_INFRA_OVERRIDES,
     PREFECT_UNIT_TEST_MODE,
     temporary_settings,
 )
@@ -1350,10 +1351,11 @@ async def test_prefect_api_tls_insecure_skip_verify_setting_set_to_true(monkeypa
         headers=ANY,
         verify=False,
         transport=ANY,
         base_url=ANY,
         timeout=ANY,
+        enable_csrf_support=ANY,
     )
 
 
 async def test_prefect_api_tls_insecure_skip_verify_setting_set_to_false(monkeypatch):
     with temporary_settings(updates={PREFECT_API_TLS_INSECURE_SKIP_VERIFY: False}):
@@ -1364,10 +1366,11 @@ async def test_prefect_api_tls_insecure_skip_verify_setting_set_to_false(monkeyp
     mock.assert_called_once_with(
         headers=ANY,
         transport=ANY,
         base_url=ANY,
         timeout=ANY,
+        enable_csrf_support=ANY,
     )
 
 
 async def test_prefect_api_tls_insecure_skip_verify_default_setting(monkeypatch):
     mock = Mock()
@@ -1376,10 +1379,11 @@ async def test_prefect_api_tls_insecure_skip_verify_default_setting(monkeypatch)
     mock.assert_called_once_with(
         headers=ANY,
         transport=ANY,
         base_url=ANY,
         timeout=ANY,
+        enable_csrf_support=ANY,
     )
 
 
 class TestResolveDataDoc:
     @pytest.fixture(autouse=True)
@@ -2240,5 +2244,27 @@ class TestPrefectClientDeploymentSchedules:
         nonexistent_schedule_id = str(uuid4())
         with pytest.raises(prefect.exceptions.ObjectNotFound):
             await prefect_client.delete_deployment_schedule(
                 deployment.id, nonexistent_schedule_id
             )
+
+
+class TestPrefectClientCsrfSupport:
+    def test_enabled_ephemeral(self, prefect_client: PrefectClient):
+        assert prefect_client.server_type == ServerType.EPHEMERAL
+        assert prefect_client._client.enable_csrf_support
+
+    async def test_enabled_server_type(self, hosted_api_server):
+        async with PrefectClient(hosted_api_server) as prefect_client:
+            assert prefect_client.server_type == ServerType.SERVER
+            assert prefect_client._client.enable_csrf_support
+
+    async def test_not_enabled_server_type_cloud(self):
+        async with PrefectClient(PREFECT_CLOUD_API_URL.value()) as prefect_client:
+            assert prefect_client.server_type == ServerType.CLOUD
+            assert not prefect_client._client.enable_csrf_support
+
+    async def test_disabled_setting_disabled(self, hosted_api_server):
+        with temporary_settings({PREFECT_CLIENT_CSRF_SUPPORT_ENABLED: False}):
+            async with PrefectClient(hosted_api_server) as prefect_client:
+                assert prefect_client.server_type == ServerType.SERVER
+                assert not prefect_client._client.enable_csrf_support
diff --git a/tests/test_deployments.py b/tests/test_deployments.py
index 4522ccd86..36b8721f0 100644
--- a/tests/test_deployments.py
+++ b/tests/test_deployments.py
@@ -1010,10 +1010,11 @@ class TestRunDeployment:
                     200,
                     json={**mock_flowrun_response, "state": {"type": terminal_state}},
                 ),
             ]
 
+            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
             router.get(f"/deployments/name/{d.flow_name}/{d.name}").pass_through()
             router.post(f"/deployments/{deployment_id}/create_flow_run").pass_through()
             flow_polls = router.get(re.compile("/flow_runs/.*")).mock(
                 side_effect=poll_responses
             )
@@ -1057,10 +1058,11 @@ class TestRunDeployment:
                     200,
                     json={**mock_flowrun_response, "state": {"type": terminal_state}},
                 ),
             ]
 
+            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
             router.get(f"/deployments/name/{d.flow_name}/{d.name}").pass_through()
             router.post(f"/deployments/{deployment_id}/create_flow_run").pass_through()
             flow_polls = router.get(re.compile("/flow_runs/.*")).mock(
                 side_effect=poll_responses
             )
@@ -1182,10 +1184,11 @@ class TestRunDeployment:
         }
 
         with respx.mock(
             base_url=PREFECT_API_URL.value(), assert_all_mocked=True
         ) as router:
+            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
             router.get(f"/deployments/name/{d.flow_name}/{d.name}").pass_through()
             router.post(f"/deployments/{deployment_id}/create_flow_run").pass_through()
             flow_polls = router.request(
                 "GET", re.compile(PREFECT_API_URL.value() + "/flow_runs/.*")
             ).mock(
@@ -1215,10 +1218,11 @@ class TestRunDeployment:
         with respx.mock(
             base_url=PREFECT_API_URL.value(),
             assert_all_mocked=True,
             assert_all_called=False,
         ) as router:
+            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
             router.get(f"/deployments/name/{d.flow_name}/{d.name}").pass_through()
             router.post(f"/deployments/{deployment_id}/create_flow_run").pass_through()
             flow_polls = router.request(
                 "GET", re.compile(PREFECT_API_URL.value() + "/flow_runs/.*")
             ).mock(
@@ -1259,10 +1263,11 @@ class TestRunDeployment:
         with respx.mock(
             base_url=PREFECT_API_URL.value(),
             assert_all_mocked=True,
             assert_all_called=False,
         ) as router:
+            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
             router.get(f"/deployments/name/{d.flow_name}/{d.name}").pass_through()
             router.post(f"/deployments/{deployment_id}/create_flow_run").pass_through()
             flow_polls = router.request(
                 "GET", re.compile(PREFECT_API_URL.value() + "/flow_runs/.*")
             ).mock(side_effect=side_effects)

Vulnerable code: from functools import partial
from httpx import HTTPStatusError, Response
        request: Callable,
                response = await request()
                # We got a response; return immediately if it is not retryable
                if response.status_code not in retry_codes:
    async def send(self, *args, **kwargs) -> Response:
        api_request = partial(super().send, *args, **kwargs)

            request=api_request,
        self._client = PrefectHttpxClient(**httpx_settings)
        self._client = PrefectHttpxClient(**httpx_settings)
"""Whether or not to protect the API from CSRF attacks. Experimental and
currently defaults to `False`."""
"""How long a CSRF token is valid for. Defaults to 1 hour."""
from unittest.mock import call
        mock_anyio_sleep.assert_has_awaits([call(2), call(4), call(8)])
        mock_anyio_sleep.assert_has_awaits([call(5), call(0), call(10), call(2.0)])
            [call(pytest.approx(n, rel=0.2)) for n in [2, 4, 8]]
Fixed code: import uuid
from datetime import datetime, timezone
    Optional,
from httpx import HTTPStatusError, Request, Response
from prefect.client.schemas.objects import CsrfToken
    def __init__(self, *args, enable_csrf_support: bool = False, **kwargs):
        self.enable_csrf_support: bool = enable_csrf_support
        self.csrf_token: Optional[str] = None
        self.csrf_token_expiration: Optional[datetime] = None
        self.csrf_client_id: uuid.UUID = uuid.uuid4()

        super().__init__(*args, **kwargs)

        request: Request,
        send: Callable[[Request], Awaitable[Response]],
        send_args: Tuple,
        send_kwargs: Dict,
        is_change_request = request.method.lower() in {"post", "put", "patch", "delete"}

        if self.enable_csrf_support and is_change_request:
            await self._add_csrf_headers(request=request)

                response = await send(request, *send_args, **send_kwargs)
                # We got a response; check if it's a CSRF error, otherwise
                # return immediately if it is not retryable
                if (
                    response.status_code == status.HTTP_403_FORBIDDEN
                    and "Invalid CSRF token" in response.text
                ):
                    # We got a CSRF error, clear the token and try again
                    self.csrf_token = None
                    await self._add_csrf_headers(request)
                elif response.status_code not in retry_codes:
    async def send(self, request: Request, *args, **kwargs) -> Response:
        - 403 Forbidden, if the request failed due to CSRF protection
        - 408 Request Timeout
        - Any additional status codes provided in `PREFECT_CLIENT_RETRY_EXTRA_CODES`
        super_send = super().send
            request=request,
            send=super_send,
            send_args=args,
            send_kwargs=kwargs,

    async def _add_csrf_headers(self, request: Request):
        now = datetime.now(timezone.utc)

        if not self.enable_csrf_support:
            return

        if not self.csrf_token or (
            self.csrf_token_expiration and now > self.csrf_token_expiration
        ):
            token_request = self.build_request(
                "GET", f"/csrf-token?client={self.csrf_client_id}"
            )

            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
                    self.enable_csrf_support = False
                    return

                raise

            token: CsrfToken = CsrfToken.parse_obj(token_response.json())
            self.csrf_token = token.token
            self.csrf_token_expiration = token.expiration

        request.headers["Prefect-Csrf-Token"] = self.csrf_token
        request.headers["Prefect-Csrf-Client"] = str(self.csrf_client_id)
        self._client = PrefectHttpxClient(**httpx_settings, enable_csrf_support=False)
    PREFECT_CLIENT_CSRF_SUPPORT_ENABLED,

        enable_csrf_support = (
            self.server_type != ServerType.CLOUD
            and PREFECT_CLIENT_CSRF_SUPPORT_ENABLED.value()
        )

        self._client = PrefectHttpxClient(
            **httpx_settings, enable_csrf_support=enable_csrf_support
        )


class CsrfToken(ObjectBaseModel):
    token: str = Field(
        default=...,
        description="The CSRF token",
    )
    client: str = Field(
        default=..., description="The client id associated with the CSRF token"
    )
    expiration: DateTimeTZ = Field(
        default=..., description="The expiration time of the CSRF token"
    )
PREFECT_CLIENT_CSRF_SUPPORT_ENABLED = Setting(bool, default=True)
"""
Determines if CSRF token handling is active in the Prefect client for API
requests.

When enabled (`True`), the client automatically manages CSRF tokens by
retrieving, storing, and including them in applicable state-changing requests
(POST, PUT, PATCH, DELETE) to the API.

Disabling this setting (`False`) means the client will not handle CSRF tokens,
which might be suitable for environments where CSRF protection is disabled.

Defaults to `True`, ensuring CSRF protection is enabled by default.
"""

"""
Controls the activation of CSRF protection for the Prefect server API.

When enabled (`True`), the server enforces CSRF validation checks on incoming
state-changing requests (POST, PUT, PATCH, DELETE), requiring a valid CSRF
token to be included in the request headers or body. This adds a layer of
security by preventing unauthorized or malicious sites from making requests on
behalf of authenticated users.

It is recommended to enable this setting in production environments where the
API is exposed to web clients to safeguard against CSRF attacks.

Note: Enabling this setting requires corresponding support in the client for
CSRF token management. See PREFECT_CLIENT_CSRF_SUPPORT_ENABLED for more.
"""
"""
Specifies the duration for which a CSRF token remains valid after being issued
by the server.

The default expiration time is set to 1 hour, which offers a reasonable
compromise. Adjust this setting based on your specific security requirements
and usage patterns.
"""
        respx_mock.get("/csrf-token", params={"client": ANY}).pass_through()
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from typing import AsyncGenerator, List, Tuple
from unittest import mock
from prefect.client.schemas.objects import CsrfToken
now = datetime.now(timezone.utc)

RESPONSE_CSRF = Response(
    status.HTTP_200_OK,
    json=CsrfToken(
        client="test_client", token="test_token", expiration=now + timedelta(days=1)
    ).dict(json_compatible=True, exclude_unset=True),
    request=Request("a test request", "fake.url/fake/route"),
)

RESPONSE_400 = Response(
    status.HTTP_400_BAD_REQUEST,
    json={"detail": "You done bad things"},
    request=Request("a test request", "fake.url/fake/route"),
)

RESPONSE_404 = Response(
    status.HTTP_404_NOT_FOUND,
    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."},
    request=Request("a test request", "fake.url/fake/route"),
)

        mock_anyio_sleep.assert_has_awaits([mock.call(2), mock.call(4), mock.call(8)])
        mock_anyio_sleep.assert_has_awaits(
            [mock.call(5), mock.call(0), mock.call(10), mock.call(2.0)]
        )
            [mock.call(pytest.approx(n, rel=0.2)) for n in [2, 4, 8]]


@asynccontextmanager
async def mocked_client(
    responses: List[Response],
) -> AsyncGenerator[Tuple[PrefectHttpxClient, mock.AsyncMock], None]:
    with mock.patch("httpx.AsyncClient.send", autospec=True) as send:
        send.side_effect = responses
        client = PrefectHttpxClient(enable_csrf_support=True)
        async with client:
            try:
                yield client, send
            finally:
                pass


class TestCsrfSupport:
    async def test_no_csrf_headers_not_change_request(self):
        async with mocked_client(responses=[RESPONSE_200]) as (client, send):
            await client.get(url="fake.url/fake/route")

        request = send.call_args[0][1]
        assert isinstance(request, httpx.Request)

        assert "Prefect-Csrf-Token" not in request.headers
        assert "Prefect-Csrf-Client" not in request.headers

    @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
    async def test_csrf_headers_on_change_request(self, method: str):
        async with mocked_client(responses=[RESPONSE_CSRF, RESPONSE_200]) as (
            client,
            send,
        ):
            await getattr(client, method)(url="fake.url/fake/route")

        assert send.await_count == 2

        # The first call should be for the CSRF token
        request = send.call_args_list[0][0][1]
        assert isinstance(request, httpx.Request)
        assert request.method == "GET"
        assert request.url == httpx.URL(
            f"/csrf-token?client={str(client.csrf_client_id)}"
        )

        # The second call should be for the actual request
        request = send.call_args_list[1][0][1]
        assert isinstance(request, httpx.Request)
        assert request.method == method.upper()
        assert request.url == httpx.URL("/fake.url/fake/route")
        assert request.headers["Prefect-Csrf-Token"] == "test_token"
        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)

    async def test_refreshes_token_on_csrf_403(self):
        async with mocked_client(
            responses=[
                RESPONSE_CSRF,
                RESPONSE_INVALID_TOKEN,
                RESPONSE_CSRF,
                RESPONSE_200,
            ]
        ) as (
            client,
            send,
        ):
            await client.post(url="fake.url/fake/route")

        assert send.await_count == 4

        # The first call should be for the CSRF token
        request = send.call_args_list[0][0][1]
        assert isinstance(request, httpx.Request)
        assert request.method == "GET"
        assert request.url == httpx.URL(
            f"/csrf-token?client={str(client.csrf_client_id)}"
        )

        # The second call should be for the actual request
        request = send.call_args_list[1][0][1]
        assert isinstance(request, httpx.Request)
        assert request.url == httpx.URL("/fake.url/fake/route")
        assert request.headers["Prefect-Csrf-Token"] == "test_token"
        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)

        # The third call should be a refresh of the CSRF token
        request = send.call_args_list[0][0][1]
        assert isinstance(request, httpx.Request)
        assert request.method == "GET"
        assert request.url == httpx.URL(
            f"/csrf-token?client={str(client.csrf_client_id)}"
        )

        # The fourth call should be for the actual request
        request = send.call_args_list[1][0][1]
        assert isinstance(request, httpx.Request)
        assert request.url == httpx.URL("/fake.url/fake/route")
        assert request.headers["Prefect-Csrf-Token"] == "test_token"
        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)

    async def test_does_not_refresh_csrf_token_not_expired(self):
        async with mocked_client(responses=[RESPONSE_200]) as (
            client,
            send,
        ):
            client.csrf_token = "fresh_token"
            client.csrf_token_expiration = now + timedelta(days=1)
            await client.post(url="fake.url/fake/route")

        assert send.await_count == 1

        request = send.call_args_list[0][0][1]
        assert isinstance(request, httpx.Request)
        assert request.url == httpx.URL("/fake.url/fake/route")
        assert request.headers["Prefect-Csrf-Token"] == "fresh_token"
        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)

    async def test_does_refresh_csrf_token_when_expired(self):
        async with mocked_client(responses=[RESPONSE_CSRF, RESPONSE_200]) as (
            client,
            send,
        ):
            client.csrf_token = "old_token"
            client.csrf_token_expiration = now - timedelta(days=1)
            await client.post(url="fake.url/fake/route")

        assert send.await_count == 2

        # The first call should be for the CSRF token
        request = send.call_args_list[0][0][1]
        assert isinstance(request, httpx.Request)
        assert request.method == "GET"
        assert request.url == httpx.URL(
            f"/csrf-token?client={str(client.csrf_client_id)}"
        )

        # The second call should be for the actual request
        request = send.call_args_list[1][0][1]
        assert isinstance(request, httpx.Request)
        assert request.url == httpx.URL("/fake.url/fake/route")
        assert request.headers["Prefect-Csrf-Token"] == "test_token"
        assert request.headers["Prefect-Csrf-Client"] == str(client.csrf_client_id)

    async def test_raises_exception_bad_csrf_token_response(self):
        async with mocked_client(responses=[RESPONSE_400]) as (
            client,
            _,
        ):
            with pytest.raises(PrefectHTTPStatusError):
                await client.post(url="fake.url/fake/route")

    async def test_disables_csrf_support_404_token_endpoint(self):
        async with mocked_client(responses=[RESPONSE_404, 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
    PREFECT_CLIENT_CSRF_SUPPORT_ENABLED,
        enable_csrf_support=ANY,
        enable_csrf_support=ANY,
        enable_csrf_support=ANY,


class TestPrefectClientCsrfSupport:
    def test_enabled_ephemeral(self, prefect_client: PrefectClient):
        assert prefect_client.server_type == ServerType.EPHEMERAL
        assert prefect_client._client.enable_csrf_support

    async def test_enabled_server_type(self, hosted_api_server):
        async with PrefectClient(hosted_api_server) as prefect_client:
            assert prefect_client.server_type == ServerType.SERVER
            assert prefect_client._client.enable_csrf_support

    async def test_not_enabled_server_type_cloud(self):
        async with PrefectClient(PREFECT_CLOUD_API_URL.value()) as prefect_client:
            assert prefect_client.server_type == ServerType.CLOUD
            assert not prefect_client._client.enable_csrf_support

    async def test_disabled_setting_disabled(self, hosted_api_server):
        with temporary_settings({PREFECT_CLIENT_CSRF_SUPPORT_ENABLED: False}):
            async with PrefectClient(hosted_api_server) as prefect_client:
                assert prefect_client.server_type == ServerType.SERVER
                assert not prefect_client._client.enable_csrf_support
            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()
            router.get("/csrf-token", params={"client": mock.ANY}).pass_through()