add util.add_samesite_fields(response, value)

This function adds f"SameSite={value}" to each cookie's header
stored in response.

This is needed because starlette does not currently support
the `samesite` argument in Response.set_cookie. It is merged,
however, and waiting for next release.

Signed-off-by: Kevin Morris <kevr@0cost.org>

aurweb/routers/auth.py

@@ -6,6 +6,7 @@ from fastapi.responses import HTMLResponse, RedirectResponse
 
 import aurweb.config
 
+from aurweb import util
 from aurweb.auth import auth_required
 from aurweb.models.user import User
 from aurweb.templates import make_context, render_template
@@ -63,7 +64,7 @@ async def login_post(request: Request,
     secure_cookies = aurweb.config.getboolean("options", "disable_http_login")
     response.set_cookie("AURSID", sid, expires=expires_at,
                         secure=secure_cookies, httponly=True)
-    return response
+    return util.add_samesite_fields(response, "strict")
 
 
 @router.get("/logout")

aurweb/routers/html.py

@@ -8,6 +8,7 @@ from fastapi.responses import HTMLResponse, RedirectResponse
 
 import aurweb.config
 
+from aurweb import util
 from aurweb.templates import make_context, render_template
 
 router = APIRouter()
@@ -50,7 +51,7 @@ async def language(request: Request,
     secure_cookies = aurweb.config.getboolean("options", "disable_http_login")
     response.set_cookie("AURLANG", set_lang,
                         secure=secure_cookies, httponly=True)
-    return response
+    return util.add_samesite_fields(response, "strict")
 
 
 @router.get("/", response_class=HTMLResponse)

aurweb/routers/sso.py

@@ -14,6 +14,7 @@ from starlette.requests import Request
 import aurweb.config
 import aurweb.db
 
+from aurweb import util
 from aurweb.l10n import get_translator_for_request
 from aurweb.schema import Bans, Sessions, Users
 
@@ -140,7 +141,7 @@ async def authenticate(request: Request, redirect: str = None, conn=Depends(aurw
             response.set_cookie(key="SSO_ID_TOKEN", value=token["id_token"],
                                 path="/sso/", httponly=True,
                                 secure=secure_cookies)
-        return response
+        return util.add_samesite_fields(response, "strict")
     else:
         # We’ve got a severe integrity violation.
         raise Exception("Multiple accounts found for SSO account " + sub)

aurweb/templates.py

@@ -94,4 +94,4 @@ def render_template(request: Request,
                         secure=secure_cookies, httponly=True)
     response.set_cookie("AURTZ", context.get("timezone"),
                         secure=secure_cookies, httponly=True)
-    return response
+    return util.add_samesite_fields(response, "strict")

aurweb/util.py

@@ -9,6 +9,7 @@ from urllib.parse import quote_plus, urlparse
 from zoneinfo import ZoneInfo
 
 from email_validator import EmailNotValidError, EmailUndeliverableError, validate_email
+from fastapi.responses import Response
 from jinja2 import pass_context
 
 import aurweb.config
@@ -88,7 +89,7 @@ def migrate_cookies(request, response):
     secure_cookies = aurweb.config.getboolean("options", "disable_http_login")
     for k, v in request.cookies.items():
         response.set_cookie(k, v, secure=secure_cookies, httponly=True)
-    return response
+    return add_samesite_fields(response, "strict")
 
 
 @pass_context
@@ -136,3 +137,15 @@ def jsonify(obj):
     if isinstance(obj, datetime):
         obj = int(obj.timestamp())
     return obj
+
+
+def add_samesite_fields(response: Response, value: str):
+    """ Set the SameSite field on all cookie headers found.
+    Taken from https://github.com/tiangolo/fastapi/issues/1099. """
+    for idx, header in enumerate(response.raw_headers):
+        if header[0].decode() == "set-cookie":
+            cookie = header[1].decode()
+            if f"SameSite={value}" not in cookie:
+                cookie += f"; SameSite={value}"
+                response.raw_headers[idx] = (header[0], cookie.encode())
+    return response

test/test_auth_routes.py

@@ -111,6 +111,8 @@ def test_secure_login(mock):
     cookie = next(c for c in response.cookies if c.name == "AURSID")
     assert cookie.secure is True
     assert cookie.has_nonstandard_attr("HttpOnly") is True
+    assert cookie.has_nonstandard_attr("SameSite") is True
+    assert cookie.get_nonstandard_attr("SameSite") == "strict"
     assert cookie.value is not None and len(cookie.value) > 0
 
     # Let's make sure we actually have a session relationship