Custom errors (#30)

* send_csrf_failed for customizing errors, closes #28

Author: simonw

README.md

@@ -100,3 +100,45 @@ app = asgi_csrf(
     skip_if_scope=skip_api_paths
 )
 ```
+
+### send_csrf_failed
+
+By default, when a CSRF token is missing or invalid, the middleware will return a 403 Forbidden response page with a short error message.
+
+You can customize this behavior by passing a `send_csrf_failed` function to the middleware. This function should accept the ASGI `scope` and `send` functions, and the `message_id` of the error that occurred.
+
+The `message_id` will be an integer representing an item from the `asgi_csrf.Errors` enum.
+
+This example shows how you could customize the error message based on that `message_id`:
+
+```python
+async def custom_csrf_failed(scope, send, message_id):
+    assert scope["type"] == "http"
+    await send(
+        {
+            "type": "http.response.start",
+            "status": 403,
+            "headers": [[b"content-type", b"text/html; charset=utf-8"]],
+        }
+    )
+    await send(
+        {
+            "type": "http.response.body",
+            "body": {
+                Errors.FORM_URLENCODED_MISMATCH: "custom form-urlencoded error",
+                Errors.MULTIPART_MISMATCH: "custom multipart error",
+                Errors.FILE_BEFORE_TOKEN: "custom file before token error",
+                Errors.UNKNOWN_CONTENT_TYPE: "custom unknown content type error",
+            }
+            .get(message_id, "")
+            .encode("utf-8"),
+        }
+    )
+
+
+app = asgi_csrf(
+    app,
+    signing_secret="secret-goes-here",
+    send_csrf_failed=custom_csrf_failed
+)
+```

asgi_csrf.py

@@ -1,4 +1,5 @@
 from http.cookies import SimpleCookie
+from enum import Enum
 import fnmatch
 from functools import wraps
 from multipart import FormParser
@@ -16,6 +17,21 @@
 ENV_SECRET = "ASGI_CSRF_SECRET"
 
 
+class Errors(Enum):
+    FORM_URLENCODED_MISMATCH = 1
+    MULTIPART_MISMATCH = 2
+    FILE_BEFORE_TOKEN = 3
+    UNKNOWN_CONTENT_TYPE = 4
+
+
+error_messages = {
+    Errors.FORM_URLENCODED_MISMATCH: "form-urlencoded POST field did not match cookie",
+    Errors.MULTIPART_MISMATCH: "multipart/form-data POST field did not match cookie",
+    Errors.FILE_BEFORE_TOKEN: "File encountered before csrftoken - make sure csrftoken is first in the HTML",
+    Errors.UNKNOWN_CONTENT_TYPE: "Unknown content-type",
+}
+
+
 def asgi_csrf_decorator(
     cookie_name=DEFAULT_COOKIE_NAME,
     http_header=DEFAULT_HTTP_HEADER,
@@ -25,7 +41,9 @@ def asgi_csrf_decorator(
     always_protect=None,
     always_set_cookie=False,
     skip_if_scope=None,
+    send_csrf_failed=None,
 ):
+    send_csrf_failed = send_csrf_failed or default_send_csrf_failed
     if signing_secret is None:
         signing_secret = os.environ.get(ENV_SECRET, None)
     if signing_secret is None:
@@ -144,9 +162,7 @@ async def wrapped_send(event):
                         return
                     else:
                         await send_csrf_failed(
-                            scope,
-                            wrapped_send,
-                            "form-urlencoded POST field did not match cookie",
+                            scope, wrapped_send, Errors.FORM_URLENCODED_MISMATCH
                         )
                         return
                 elif content_type == b"multipart/form-data":
@@ -168,22 +184,22 @@ async def wrapped_send(event):
                             await send_csrf_failed(
                                 scope,
                                 wrapped_send,
-                                "multipart/form-data POST field did not match cookie",
+                                Errors.MULTIPART_MISMATCH,
                             )
                             return
                     except FileBeforeToken:
                         await send_csrf_failed(
                             scope,
                             wrapped_send,
-                            "File encountered before csrftoken - make sure csrftoken is first in the HTML",
+                            Errors.FILE_BEFORE_TOKEN,
                         )
                         return
                     # Now replay the body
                     await app(scope, replay_receive, wrapped_send)
                     return
                 else:
                     await send_csrf_failed(
-                        scope, wrapped_send, message="Unknown content-type"
+                        scope, wrapped_send, Errors.UNKNOWN_CONTENT_TYPE
                     )
                     return
 
@@ -271,7 +287,7 @@ async def replay_receive():
     return None, replay_receive
 
 
-async def send_csrf_failed(scope, send, message="CSRF check failed"):
+async def default_send_csrf_failed(scope, send, message_id):
     assert scope["type"] == "http"
     await send(
         {
@@ -280,6 +296,7 @@ async def send_csrf_failed(scope, send, message="CSRF check failed"):
             "headers": [[b"content-type", b"text/html; charset=utf-8"]],
         }
     )
+    message = error_messages.get(message_id) or "CSRF validation failed"
     await send({"type": "http.response.body", "body": message.encode("utf-8")})
 
 
@@ -292,6 +309,7 @@ def asgi_csrf(
     always_protect=None,
     always_set_cookie=False,
     skip_if_scope=None,
+    send_csrf_failed=None,
 ):
     return asgi_csrf_decorator(
         cookie_name,
@@ -301,6 +319,7 @@ def asgi_csrf(
         always_protect=always_protect,
         always_set_cookie=always_set_cookie,
         skip_if_scope=skip_if_scope,
+        send_csrf_failed=send_csrf_failed,
     )(app)
 
 

test_asgi_csrf.py

@@ -2,7 +2,7 @@
 from starlette.applications import Starlette
 from starlette.responses import JSONResponse
 from starlette.routing import Route
-from asgi_csrf import asgi_csrf
+from asgi_csrf import asgi_csrf, Errors
 from itsdangerous.url_safe import URLSafeSerializer
 import httpx
 import json
@@ -45,6 +45,39 @@ def app_csrf():
     return asgi_csrf(hello_world_app, signing_secret=SECRET)
 
 
+async def custom_csrf_failed(scope, send, message_id):
+    assert scope["type"] == "http"
+    await send(
+        {
+            "type": "http.response.start",
+            "status": 403,
+            "headers": [[b"content-type", b"text/html; charset=utf-8"]],
+        }
+    )
+    await send(
+        {
+            "type": "http.response.body",
+            "body": {
+                Errors.FORM_URLENCODED_MISMATCH: "custom form-urlencoded error",
+                Errors.MULTIPART_MISMATCH: "custom multipart error",
+                Errors.FILE_BEFORE_TOKEN: "custom file before token error",
+                Errors.UNKNOWN_CONTENT_TYPE: "custom unknown content type error",
+            }
+            .get(message_id, "")
+            .encode("utf-8"),
+        }
+    )
+
+
+@pytest.fixture
+def app_csrf_custom_errors():
+    return asgi_csrf(
+        hello_world_app,
+        signing_secret=SECRET,
+        send_csrf_failed=custom_csrf_failed,
+    )
+
+
 @pytest.fixture
 def csrftoken():
     return URLSafeSerializer(SECRET).dumps("token", "csrftoken")
@@ -60,11 +93,11 @@ async def test_hello_world_app():
 def test_signing_secret_if_none_provided(monkeypatch):
     app = asgi_csrf(hello_world_app)
     # Should be randomly generated
-    assert isinstance(app.__closure__[6].cell_contents.secret_key, bytes)
+    assert isinstance(app.__closure__[7].cell_contents.secret_key, bytes)
     # Should pick up `ASGI_CSRF_SECRET` if available
     monkeypatch.setenv("ASGI_CSRF_SECRET", "secret-from-environment")
     app2 = asgi_csrf(hello_world_app)
-    assert app2.__closure__[6].cell_contents.secret_key == b"secret-from-environment"
+    assert app2.__closure__[7].cell_contents.secret_key == b"secret-from-environment"
 
 
 @pytest.mark.asyncio
@@ -142,15 +175,24 @@ async def test_prevents_post_if_cookie_not_sent_in_post(app_csrf, csrftoken):
 
 
 @pytest.mark.asyncio
-async def test_prevents_post_if_cookie_not_sent_in_post(app_csrf, csrftoken):
-    async with httpx.AsyncClient(app=app_csrf) as client:
+@pytest.mark.parametrize("custom_errors", (False, True))
+async def test_prevents_post_if_cookie_not_sent_in_post(
+    custom_errors, app_csrf, app_csrf_custom_errors, csrftoken
+):
+    async with httpx.AsyncClient(
+        app=app_csrf_custom_errors if custom_errors else app_csrf
+    ) as client:
         response = await client.post(
             "http://localhost/",
             cookies={"csrftoken": csrftoken},
             data={"csrftoken": csrftoken[-1]},
         )
     assert 403 == response.status_code
-    assert response.text == "form-urlencoded POST field did not match cookie"
+    assert (
+        response.text == "custom form-urlencoded error"
+        if custom_errors
+        else "form-urlencoded POST field did not match cookie"
+    )
 
 
 @pytest.mark.asyncio
@@ -194,9 +236,14 @@ async def test_multipart(csrftoken):
 
 
 @pytest.mark.asyncio
-async def test_multipart_failure_wrong_token(csrftoken):
+@pytest.mark.parametrize("custom_errors", (False, True))
+async def test_multipart_failure_wrong_token(csrftoken, custom_errors):
     async with httpx.AsyncClient(
-        app=asgi_csrf(hello_world_app, signing_secret=SECRET)
+        app=asgi_csrf(
+            hello_world_app,
+            signing_secret=SECRET,
+            send_csrf_failed=custom_csrf_failed if custom_errors else None,
+        )
     ) as client:
         response = await client.post(
             "http://localhost/",
@@ -205,7 +252,11 @@ async def test_multipart_failure_wrong_token(csrftoken):
             cookies={"csrftoken": csrftoken[:-1]},
         )
         assert response.status_code == 403
-        assert response.text == "multipart/form-data POST field did not match cookie"
+        assert (
+            response.text == "custom multipart error"
+            if custom_errors
+            else "multipart/form-data POST field did not match cookie"
+        )
 
 
 class TrickEmptyDictionary(dict):
@@ -215,9 +266,14 @@ def __bool__(self):
 
 
 @pytest.mark.asyncio
-async def test_multipart_failure_missing_token(csrftoken):
+@pytest.mark.parametrize("custom_errors", (False, True))
+async def test_multipart_failure_missing_token(csrftoken, custom_errors):
     async with httpx.AsyncClient(
-        app=asgi_csrf(hello_world_app, signing_secret=SECRET)
+        app=asgi_csrf(
+            hello_world_app,
+            signing_secret=SECRET,
+            send_csrf_failed=custom_csrf_failed if custom_errors else None,
+        )
     ) as client:
         response = await client.post(
             "http://localhost/",
@@ -226,18 +282,27 @@ async def test_multipart_failure_missing_token(csrftoken):
             cookies={"csrftoken": csrftoken},
         )
         assert response.status_code == 403
-        assert response.text == "multipart/form-data POST field did not match cookie"
+        assert response.text == (
+            "custom multipart error"
+            if custom_errors
+            else "multipart/form-data POST field did not match cookie"
+        )
 
 
 @pytest.mark.asyncio
-async def test_multipart_failure_file_comes_before_token(csrftoken):
+@pytest.mark.parametrize("custom_errors", (False, True))
+async def test_multipart_failure_file_comes_before_token(csrftoken, custom_errors):
     async with httpx.AsyncClient(
-        app=asgi_csrf(hello_world_app, signing_secret=SECRET)
+        app=asgi_csrf(
+            hello_world_app,
+            signing_secret=SECRET,
+            send_csrf_failed=custom_csrf_failed if custom_errors else None,
+        )
     ) as client:
         request = httpx.Request(
             url="http://localhost/",
             method="POST",
-            data=(
+            content=(
                 b"--boo\r\n"
                 b'Content-Disposition: form-data; name="csv"; filename="data.csv"'
                 b"\r\nContent-Type: text/csv\r\n\r\n"
@@ -255,8 +320,9 @@ async def test_multipart_failure_file_comes_before_token(csrftoken):
         response = await client.send(request)
         assert response.status_code == 403
         assert (
-            response.text
-            == "File encountered before csrftoken - make sure csrftoken is first in the HTML"
+            response.text == "custom file before token error"
+            if custom_errors
+            else "File encountered before csrftoken - make sure csrftoken is first in the HTML"
         )