cookie attribute options

* cookie attribute options, closes #25 * Make test more stable in face of code changes * Configuring the cookie docs, refs #25 * Tests for cookie options, closes #25 --------- Co-authored-by: Kevin Abraham <[email protected]> Co-authored-by: Simon Willison <[email protected]>
simonw · Aug 15, 2024 · a18f7d0 · a18f7d0
1 parent 4db11ae
commit a18f7d0
Show file tree

Hide file tree

Showing 3 changed files with 143 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -64,6 +64,28 @@ If you would like the middleware to set that cookie for any incoming request tha
 app = asgi_csrf(app, signing_secret="secret-goes-here", always_set_cookie=True)
 ```
 
+## Configuring the cookie
+
+The middleware can be configured with several options to control how the CSRF cookie is set:
+
+```python
+app = asgi_csrf(
+    app,
+    signing_secret="secret-goes-here",
+    cookie_name="csrftoken",
+    cookie_path="/",
+    cookie_domain=None,
+    cookie_secure=False,
+    cookie_samesite="Lax"
+)
+```
+
+- `cookie_name`: The name of the cookie to set. Defaults to `"csrftoken"`.
+- `cookie_path`: The path for which the cookie is valid. Defaults to `"/"`, meaning the cookie is valid for the entire domain.
+- `cookie_domain`: The domain for which the cookie is valid. Defaults to `None`, which means the cookie will only be valid for the current domain.
+- `cookie_secure`: If set to `True`, the cookie will only be sent over HTTPS connections. Defaults to `False`.
+- `cookie_samesite`: Controls how the cookie is sent with cross-site requests. Can be set to `"Strict"`, `"Lax"`, or `"None"`. Defaults to `"Lax"`.
+
 ## Other cases that skip CSRF protection
 
 If the request includes an `Authorization: Bearer ...` header, commonly used by OAuth and JWT authentication, the request will not be required to include a CSRF token. This is because browsers cannot send those headers in a context that can be abused.
@@ -101,7 +123,7 @@ app = asgi_csrf(
 )
 ```
 
-### send_csrf_failed
+## Custom errors with 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.
 

diff --git a/asgi_csrf.py b/asgi_csrf.py
@@ -10,6 +10,10 @@
 import secrets
 
 DEFAULT_COOKIE_NAME = "csrftoken"
+DEFAULT_COOKIE_PATH = "/"
+DEFAULT_COOKIE_DOMAIN = None
+DEFAULT_COOKIE_SECURE = False
+DEFAULT_COOKIE_SAMESITE = "Lax"
 DEFAULT_FORM_INPUT = "csrftoken"
 DEFAULT_HTTP_HEADER = "x-csrftoken"
 DEFAULT_SIGNING_NAMESPACE = "csrftoken"
@@ -41,6 +45,10 @@ def asgi_csrf_decorator(
     always_protect=None,
     always_set_cookie=False,
     skip_if_scope=None,
+    cookie_path=DEFAULT_COOKIE_PATH,
+    cookie_domain=DEFAULT_COOKIE_DOMAIN,
+    cookie_secure=DEFAULT_COOKIE_SECURE,
+    cookie_samesite=DEFAULT_COOKIE_SAMESITE,
     send_csrf_failed=None,
 ):
     send_csrf_failed = send_csrf_failed or default_send_csrf_failed
@@ -106,12 +114,22 @@ async def wrapped_send(event):
                     else:
                         new_headers = original_headers
                     if should_set_cookie:
+                        cookie_attrs = [
+                            "{}={}".format(cookie_name, csrftoken),
+                            "Path={}".format(cookie_path),
+                            "SameSite={}".format(cookie_samesite),
+                        ]
+
+                        if cookie_domain is not None:
+                            cookie_attrs.append("Domain={}".format(cookie_domain))
+
+                        if cookie_secure:
+                            cookie_attrs.append("Secure")
+
                         new_headers.append(
                             (
                                 b"set-cookie",
-                                "{}={}; Path=/".format(cookie_name, csrftoken).encode(
-                                    "utf-8"
-                                ),
+                                "; ".join(cookie_attrs).encode("utf-8"),
                             )
                         )
                     event = {
@@ -309,6 +327,10 @@ def asgi_csrf(
     always_protect=None,
     always_set_cookie=False,
     skip_if_scope=None,
+    cookie_path=DEFAULT_COOKIE_PATH,
+    cookie_domain=DEFAULT_COOKIE_DOMAIN,
+    cookie_secure=DEFAULT_COOKIE_SECURE,
+    cookie_samesite=DEFAULT_COOKIE_SAMESITE,
     send_csrf_failed=None,
 ):
     return asgi_csrf_decorator(
@@ -319,6 +341,10 @@ def asgi_csrf(
         always_protect=always_protect,
         always_set_cookie=always_set_cookie,
         skip_if_scope=skip_if_scope,
+        cookie_path=cookie_path,
+        cookie_domain=cookie_domain,
+        cookie_secure=cookie_secure,
+        cookie_samesite=cookie_samesite,
         send_csrf_failed=send_csrf_failed,
     )(app)
 

diff --git a/test_asgi_csrf.py b/test_asgi_csrf.py
@@ -92,12 +92,22 @@ 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__[7].cell_contents.secret_key, bytes)
+    def _get_secret_key(app):
+        found = [
+            cell.cell_contents
+            for cell in app.__closure__
+            if "URLSafeSerializer" in repr(cell)
+        ]
+        assert found
+        return found[0].secret_key
+
+    assert isinstance(_get_secret_key(app), 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__[7].cell_contents.secret_key == b"secret-from-environment"
+    assert _get_secret_key(app2) == b"secret-from-environment"
 
 
 @pytest.mark.asyncio
@@ -106,7 +116,7 @@ async def test_asgi_csrf_sets_cookie(app_csrf):
         response = await client.get("http://localhost/")
     assert b'{"hello":"world"}' == response.content
     assert "csrftoken" in response.cookies
-    assert response.headers["set-cookie"].endswith("; Path=/")
+    assert response.headers["set-cookie"].endswith("; Path=/; SameSite=Lax")
     assert "Cookie" == response.headers["vary"]
 
 
@@ -116,7 +126,7 @@ async def test_asgi_csrf_modifies_existing_vary_header(app_csrf):
         response = await client.get("http://localhost/?_vary=User-Agent")
     assert b'{"hello":"world"}' == response.content
     assert "csrftoken" in response.cookies
-    assert response.headers["set-cookie"].endswith("; Path=/")
+    assert response.headers["set-cookie"].endswith("; Path=/; SameSite=Lax")
     assert "User-Agent, Cookie" == response.headers["vary"]
 
 
@@ -430,3 +440,80 @@ async def test_asgi_lifespan():
                 cookies={"foo": "bar"},
             )
             assert 200 == response.status_code
+
+
+# Tests for different cookie options
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cookie_name", ["csrftoken", "custom_csrf"])
+async def test_cookie_name(cookie_name):
+    wrapped_app = asgi_csrf(
+        hello_world_app, signing_secret="secret", cookie_name=cookie_name
+    )
+    async with httpx.AsyncClient(app=wrapped_app) as client:
+        response = await client.get("http://testserver/")
+    assert cookie_name in response.cookies
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cookie_path", ["/", "/custom"])
+async def test_cookie_path(cookie_path):
+    wrapped_app = asgi_csrf(
+        hello_world_app, signing_secret="secret", cookie_path=cookie_path
+    )
+    async with httpx.AsyncClient(app=wrapped_app) as client:
+        response = await client.get("http://testserver/")
+    assert f"Path={cookie_path}" in response.headers["set-cookie"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cookie_domain", [None, "example.com"])
+async def test_cookie_domain(cookie_domain):
+    wrapped_app = asgi_csrf(
+        hello_world_app, signing_secret="secret", cookie_domain=cookie_domain
+    )
+    async with httpx.AsyncClient(app=wrapped_app) as client:
+        response = await client.get("http://testserver/")
+    if cookie_domain:
+        assert f"Domain={cookie_domain}" in response.headers["set-cookie"]
+    else:
+        assert "Domain" not in response.headers["set-cookie"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cookie_secure", [True, False])
+async def test_cookie_secure(cookie_secure):
+    wrapped_app = asgi_csrf(
+        hello_world_app, signing_secret="secret", cookie_secure=cookie_secure
+    )
+    async with httpx.AsyncClient(app=wrapped_app) as client:
+        response = await client.get("http://testserver/")
+    if cookie_secure:
+        assert "Secure" in response.headers["set-cookie"]
+    else:
+        assert "Secure" not in response.headers["set-cookie"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cookie_samesite", ["Strict", "Lax", "None"])
+async def test_cookie_samesite(cookie_samesite):
+    wrapped_app = asgi_csrf(
+        hello_world_app, signing_secret="secret", cookie_samesite=cookie_samesite
+    )
+    async with httpx.AsyncClient(app=wrapped_app) as client:
+        response = await client.get("http://testserver/")
+    assert f"SameSite={cookie_samesite}" in response.headers["set-cookie"]
+
+
+@pytest.mark.asyncio
+async def test_default_cookie_options():
+    wrapped_app = asgi_csrf(hello_world_app, signing_secret="secret")
+    async with httpx.AsyncClient(app=wrapped_app) as client:
+        response = await client.get("http://testserver/")
+    set_cookie = response.headers["set-cookie"]
+    assert "csrftoken" in set_cookie
+    assert "Path=/" in set_cookie
+    assert "Domain" not in set_cookie
+    assert "Secure" not in set_cookie
+    assert "SameSite=Lax" in set_cookie