diff --git a/backend/chainlit/auth/cookie.py b/backend/chainlit/auth/cookie.py index 2fff180606..3edd64856b 100644 --- a/backend/chainlit/auth/cookie.py +++ b/backend/chainlit/auth/cookie.py @@ -25,6 +25,8 @@ _cookie_secure = _cookie_samesite == "none" _state_cookie_lifetime = 3 * 60 # 3m +# TODO: Prefix with __Host- to use in conjunction with the partitioned cookie as described here +# [CHIPS](https://github.com/privacycg/CHIPS/blob/main/README.md) _auth_cookie_name = "access_token" _state_cookie_name = "oauth_state" @@ -89,6 +91,14 @@ def set_auth_cookie(response: Response, token: str): samesite=_cookie_samesite, max_age=config.project.user_session_timeout, ) + if _cookie_secure: + # Set partitioned cookie to avoid cookie rejection errors + # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#partitioned + # TODO: Use FastAPI native methods once it supports partitioned cookie + # https://github.com/fastapi/fastapi/discussions/11285 + response.headers["Set-Cookie"] = ( + response.headers["Set-Cookie"] + "; Partitioned" + ) def clear_auth_cookie(response: Response): @@ -107,6 +117,14 @@ def set_oauth_state_cookie(response: Response, token: str): secure=_cookie_secure, max_age=_state_cookie_lifetime, ) + if _cookie_secure: + # Set partitioned cookie to avoid cookie rejection errors + # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#partitioned + # TODO: Use FastAPI native methods once it supports partitioned cookie + # https://github.com/fastapi/fastapi/discussions/11285 + response.headers["Set-Cookie"] = ( + response.headers["Set-Cookie"] + "; Partitioned" + ) def validate_oauth_state_cookie(request: Request, state: str): diff --git a/backend/chainlit/cli/__init__.py b/backend/chainlit/cli/__init__.py index f4671221ce..20d5fbec10 100644 --- a/backend/chainlit/cli/__init__.py +++ b/backend/chainlit/cli/__init__.py @@ -200,6 +200,7 @@ def chainlit_run( os.environ["OPENAI_API_KEY"] = "sk-FAKE-OPENAI-API-KEY" # This is required for authentication tests os.environ["CHAINLIT_AUTH_SECRET"] = "SUPER_SECRET" # nosec B105 + os.environ["CHAINLIT_COOKIE_SAMESITE"] = "none" else: trace_event("chainlit run") diff --git a/backend/chainlit/element.py b/backend/chainlit/element.py index 47de8102f9..75eb0d0681 100644 --- a/backend/chainlit/element.py +++ b/backend/chainlit/element.py @@ -95,7 +95,7 @@ class Element: for_id: Optional[str] = None # The language, if relevant language: Optional[str] = None - # Mime type, infered based on content if not provided + # Mime type, inferred based on content if not provided mime: Optional[str] = None def __post_init__(self) -> None: diff --git a/backend/chainlit/socket.py b/backend/chainlit/socket.py index 598e540bc2..5df243b49f 100644 --- a/backend/chainlit/socket.py +++ b/backend/chainlit/socket.py @@ -7,6 +7,7 @@ from typing_extensions import TypeAlias from chainlit.auth import get_current_user, require_login +from chainlit.auth.cookie import _auth_cookie_name from chainlit.chat_context import chat_context from chainlit.config import config from chainlit.context import init_ws_context @@ -83,7 +84,7 @@ def load_user_env(user_env): def _get_token_from_cookie(environ: WSGIEnvironment) -> Optional[str]: if cookie_header := environ.get("HTTP_COOKIE", None): cookies = cookie_parser(cookie_header) - return cookies.get("access_token", None) + return cookies.get(_auth_cookie_name, None) return None