Repository: litestar-org/litestar

Issue:
Bug: Some session data got lost when trying to set session with big payload
### Description

I'm trying to store acess and refresh tokens in request.session with some other user data, and got an error. In middleware `AuthRequiredMiddleware` i have no user data, but it should be. If I remove two long fields with tokens, it works normaly.

Additionaly, i got a second issue - `SerializationException` when it trying to decode session with big payload, but can't reproduce it

### URL to code causing the issue

_No response_

### MCVE

```python
import secrets
import time
from typing import Any

from litestar import Litestar, Request, Router, get
from litestar.config.cors import CORSConfig
from litestar.middleware.base import DefineMiddleware, MiddlewareProtocol
from litestar.middleware.session.client_side import CookieBackendConfig
from litestar.openapi import OpenAPIConfig
from litestar.openapi.plugins import SwaggerRenderPlugin
from litestar.response import Redirect
from litestar.response.redirect import ASGIRedirectResponse
from litestar.types import ASGIApp, Receive, Scope, Send
from pydantic import BaseModel, Field

USER_SESSION_KEY = "user"
AUTH_REDIRECT_KEY = "auth_redirect_from"


class UserSchema(BaseModel):
    username: str
    fullname: str | None = Field(default=None)
    firstname: str | None = Field(default=None)
    lastname: str | None = Field(default=None)
    email: str | None = Field(default=None)
    position: str | None = Field(default=None)


class UserSessionSchema(BaseModel):
    user_info: UserSchema | None = Field(default=None)
    access_token: str | None = Field(default=None)
    refresh_token: str | None = Field(default=None)
    expire_at: int | None = Field(default=None)


class AuthRequiredMiddleware(MiddlewareProtocol):
    def __init__(
        self,
        app: ASGIApp,
        api_path: str,
        auth_controller_path: str,
        login_endpoint_path: str = "/login",
        **_: Any,
    ) -> None:
        self.app = app

        self.api_path = api_path
        self.auth_controller_path = auth_controller_path
        self.login_endpoint_path = login_endpoint_path
        self.login_url = self.api_path + self.auth_controller_path + self.login_endpoint_path

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        request = Request(scope)
        print("In middleware session:", request.session)

        user_session = UserSessionSchema.model_validate(request.session.get(USER_SESSION_KEY, {}))

        def _redirect_login():
            response = ASGIRedirectResponse(path=self.login_url)

            request.session[USER_SESSION_KEY] = user_session.model_dump()
            request.session[AUTH_REDIRECT_KEY] = str(request.url)

            return response(scope, receive, send)

        if self.auth_controller_path not in request.url.path:
            if not user_session.user_info:
                return await _redirect_login()

            if AUTH_REDIRECT_KEY in request.session:
                del request.session[AUTH_REDIRECT_KEY]

        await self.app(scope, receive, send)


@get("/protected_resource")
async def test(request: Request) -> str:
    user_session = UserSessionSchema.model_validate(request.session[USER_SESSION_KEY])

    assert user_session.user_info

    return f"Hello, {user_session.user_info.username}!"


@get("/auth/login")
async def login(request: Request) -> Redirect:
    request.session[USER_SESSION_KEY] = UserSessionSchema(
        user_info=UserSchema(
            username="some_username",
            fullname="Иванов Ivan Jhon",
            firstname="Ivan",
            lastname="Иванов",
            email="some-email-ivan@notmail.com",
            position="Sernior999lvl",
        ),
        expire_at=int(time.time()) + 900,
        # XXX: Uncoment to reproduce
        # access_token="a" * 2078,
        # refresh_token="b" * 1074,
    )

    print("In login session:", request.session)

    if AUTH_REDIRECT_KEY in request.session:
        return Redirect(request.session.pop(AUTH_REDIRECT_KEY))

    return Redirect("/")


router = Router("/", route_handlers=[test, login])

app = Litestar(
    route_handlers=[router],
    openapi_config=OpenAPIConfig(
        title="Insourcing API",
        version="0.1.0",
        render_plugins=[SwaggerRenderPlugin()],
        path="/api/v1/docs/",
    ),
    cors_config=CORSConfig(
        allow_origins=["http://localhost:8000"],
        allow_credentials=True,
    ),
    middleware=[
        CookieBackendConfig(secret=secrets.token_urlsafe(24).encode()).middleware,
        DefineMiddleware(
            middleware=AuthRequiredMiddleware,
            api_path="",
            auth_controller_path="/auth",
        ),
    ],
    debug=True,
)
```


### Steps to reproduce

```bash
1. Install litestar + pydantic2 + uvicorn + uvloop
2. Do not uncomment lines and run with uvicorn FILE_NAME:app
3. open localhost:8000/protected_resource. it should work without errors and you'll see - Hello, some_username!
4. Uncomment lines and run again. You will fall into redirect cycle
```


### Screenshots

_No response_

### Logs

_No response_

### Litestar Version

2.8.2

### Platform
WSL 2.1.5.0 + Ubuntu 22.04

- [ ] Linux
- [ ] Mac
- [ ] Windows
- [X] Other (Please specify in the description above)

<!-- POLAR PLEDGE BADGE START -->
---
> [!NOTE]  
> While we are open for sponsoring on [GitHub Sponsors](https://github.com/sponsors/litestar-org/) and 
> [OpenCollective](https://opencollective.com/litestar), we also utilize [Polar.sh](https://polar.sh/) to engage in pledge-based sponsorship.
>
> Check out all issues funded or available for funding [on our Polar.sh dashboard](https://polar.sh/litestar-org)
> * If you would like to see an issue prioritized, make a pledge towards it!
> * We receive the pledge once the issue is completed & verified
> * This, along with engagement in the community, helps us know which features are a priority to our users.

<a href="https://polar.sh/litestar-org/litestar/issues/3441">
<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://polar.sh/api/github/litestar-org/litestar/issues/3441/pledge.svg?darkmode=1">
  <img alt="Fund with Polar" src="https://polar.sh/api/github/litestar-org/litestar/issues/3441/pledge.svg">
</picture>
</a>
<!-- POLAR PLEDGE BADGE END -->



Discussion:
Thanks @wallseat - good find!

Once cookies get above a certain size, we chunk them and store across multiple cookies. When this happens, cookies get stored with an enumeration, e.g., `session-0`, `session-1` etc.

In this case, the session is persisted in the headers for the first time when the request to `/protected_route` is redirected to the auth route. At this time, the session cookie is not greater than the chunk size and so it gets stored in the cookie under the name `session`.

After authentication, when the size of the session is much larger due to the presence of the tokens, the serialized session is greater than the chunk size, so the session cookie gets chunked and stored under `session-0`, `session-1`.

There is an issue with the algorithm that detects cookies that should be cleared under the condition where the cookie grows in size greater than a single chunk, and that is what we're hitting here. The original `session` cookie was not being cleared when it is superseded by a cookie called `session-0`.

Resolving patch:
diff --git a/litestar/middleware/session/client_side.py b/litestar/middleware/session/client_side.py
--- a/litestar/middleware/session/client_side.py
+++ b/litestar/middleware/session/client_side.py
@@ -5,9 +5,9 @@
 import re
 import time
 from base64 import b64decode, b64encode
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
 from os import urandom
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Final, Literal, Mapping
 
 from litestar.datastructures import MutableScopeHeaders
 from litestar.datastructures.cookie import Cookie
@@ -38,12 +38,19 @@
 NONCE_SIZE = 12
 CHUNK_SIZE = 4096 - 64
 AAD = b"additional_authenticated_data="
+SET_COOKIE_INCLUDE = {f.name for f in fields(Cookie) if f.name not in {"key", "secret"}}
+CLEAR_COOKIE_INCLUDE = {f.name for f in fields(Cookie) if f.name not in {"key", "secret", "max_age"}}
 
 
 class ClientSideSessionBackend(BaseSessionBackend["CookieBackendConfig"]):
     """Cookie backend for SessionMiddleware."""
 
-    __slots__ = ("aesgcm", "cookie_re")
+    __slots__ = (
+        "_clear_cookie_params",
+        "_set_cookie_params",
+        "aesgcm",
+        "cookie_re",
+    )
 
     def __init__(self, config: CookieBackendConfig) -> None:
         """Initialize ``ClientSideSessionBackend``.
@@ -54,6 +61,12 @@ def __init__(self, config: CookieBackendConfig) -> None:
         super().__init__(config)
         self.aesgcm = AESGCM(config.secret)
         self.cookie_re = re.compile(rf"{self.config.key}(?:-\d+)?")
+        self._set_cookie_params: Final[Mapping[str, Any]] = dict(
+            extract_dataclass_items(config, exclude_none=True, include=SET_COOKIE_INCLUDE)
+        )
+        self._clear_cookie_params: Final[Mapping[str, Any]] = dict(
+            extract_dataclass_items(config, exclude_none=True, include=CLEAR_COOKIE_INCLUDE)
+        )
 
     def dump_data(self, data: Any, scope: Scope | None = None) -> list[bytes]:
         """Given serializable data, including pydantic models and numpy types, dump it into a bytes string, encrypt,
@@ -107,19 +120,25 @@ def get_cookie_keys(self, connection: ASGIConnection) -> list[str]:
         """
         return sorted(key for key in connection.cookies if self.cookie_re.fullmatch(key))
 
-    def _create_session_cookies(self, data: list[bytes], cookie_params: dict[str, Any] | None = None) -> list[Cookie]:
+    def get_cookie_key_set(self, connection: ASGIConnection) -> set[str]:
+        """Return a set of cookie-keys from the connection if they match the session-cookie pattern.
+
+        .. versionadded:: 2.8.3
+
+        Args:
+            connection: An ASGIConnection instance
+
+        Returns:
+            A set of session-cookie keys
+        """
+        return {key for key in connection.cookies if self.cookie_re.fullmatch(key)}
+
+    def _create_session_cookies(self, data: list[bytes]) -> list[Cookie]:
         """Create a list of cookies containing the session data.
         If the data is split into multiple cookies, the key will be of the format ``session-{segment number}``,
         however if only one cookie is needed, the key will be ``session``.
         """
-        if cookie_params is None:
-            cookie_params = dict(
-                extract_dataclass_items(
-                    self.config,
-                    exclude_none=True,
-                    include={f for f in Cookie.__dict__ if f not in ("key", "secret")},
-                )
-            )
+        cookie_params = self._set_cookie_params
 
         if len(data) == 1:
             return [
@@ -156,38 +175,23 @@ async def store_in_message(self, scope_session: ScopeSession, message: Message,
 
         scope = connection.scope
         headers = MutableScopeHeaders.from_message(message)
-        cookie_keys = self.get_cookie_keys(connection)
+        connection_cookies = self.get_cookie_key_set(connection)
+        response_cookies: set[str] = set()
 
         if scope_session and scope_session is not Empty:
             data = self.dump_data(scope_session, scope=scope)
-            cookie_params = dict(
-                extract_dataclass_items(
-                    self.config,
-                    exclude_none=True,
-                    include={f for f in Cookie.__dict__ if f not in ("key", "secret")},
-                )
-            )
-            for cookie in self._create_session_cookies(data, cookie_params):
+            for cookie in self._create_session_cookies(data):
                 headers.add("Set-Cookie", cookie.to_header(header=""))
-            # Cookies with the same key overwrite the earlier cookie with that key. To expire earlier session
-            # cookies, first check how many session cookies will not be overwritten in this upcoming response.
-            # If leftover cookies are greater than or equal to 1, that means older session cookies have to be
-            # expired and their names are in cookie_keys.
-            cookies_to_clear = cookie_keys[len(data) :] if len(cookie_keys) - len(data) > 0 else []
+                response_cookies.add(cookie.key)
+
+            cookies_to_clear = connection_cookies - response_cookies
         else:
-            cookies_to_clear = cookie_keys
+            cookies_to_clear = connection_cookies
 
         for cookie_key in cookies_to_clear:
-            cookie_params = dict(
-                extract_dataclass_items(
-                    self.config,
-                    exclude_none=True,
-                    include={f for f in Cookie.__dict__ if f not in ("key", "secret", "max_age")},
-                )
-            )
             headers.add(
                 "Set-Cookie",
-                Cookie(value="null", key=cookie_key, expires=0, **cookie_params).to_header(header=""),
+                Cookie(value="null", key=cookie_key, expires=0, **self._clear_cookie_params).to_header(header=""),
             )
 
     async def load_from_connection(self, connection: ASGIConnection) -> dict[str, Any]:


Test patch:
diff --git a/tests/unit/test_middleware/test_session/test_client_side_backend.py b/tests/unit/test_middleware/test_session/test_client_side_backend.py
--- a/tests/unit/test_middleware/test_session/test_client_side_backend.py
+++ b/tests/unit/test_middleware/test_session/test_client_side_backend.py
@@ -9,6 +9,7 @@
 from cryptography.exceptions import InvalidTag
 
 from litestar import Request, get, post
+from litestar.datastructures.headers import MutableScopeHeaders
 from litestar.exceptions import ImproperlyConfiguredException
 from litestar.middleware.session import SessionMiddleware
 from litestar.middleware.session.client_side import (
@@ -18,7 +19,8 @@
     CookieBackendConfig,
 )
 from litestar.serialization import encode_json
-from litestar.testing import create_test_client
+from litestar.testing import RequestFactory, create_test_client
+from litestar.types.asgi_types import HTTPResponseStartEvent
 from tests.helpers import randbytes
 
 
@@ -220,3 +222,25 @@ def test_load_data_should_raise_invalid_tag_if_tampered_aad(cookie_session_backe
 
     with pytest.raises(InvalidTag):
         cookie_session_backend.load_data(encoded)
+
+
+async def test_store_in_message_clears_cookies_when_session_grows_gt_chunk_size(
+    cookie_session_backend: ClientSideSessionBackend,
+) -> None:
+    """Should clear the cookies when the session grows larger than the chunk size."""
+    # we have a connection that already contains a cookie header with the "session" key in it
+    connection = RequestFactory().get("/", headers={"Cookie": "session=foo"})
+    # we want to persist a new session that is larger than the chunk size
+    # by the time the encrypted data, nonce and associated data are b64 encoded, the size of
+    # this session will be > 2x larger than the chunk size
+    session = create_session(size=CHUNK_SIZE)
+    message: HTTPResponseStartEvent = {"type": "http.response.start", "status": 200, "headers": []}
+    await cookie_session_backend.store_in_message(session, message, connection)
+    # due to the large session stored in multiple chunks, we now enumerate the name of the cookies
+    # e.g., session-0, session-1, session-2, etc. This means we need to have a cookie with the name
+    # "session" in the response headers that is set to null to clear the original cookie.
+    headers = MutableScopeHeaders.from_message(message)
+    assert len(headers.headers) > 1
+    header_name, header_content = headers.headers[-1]
+    assert header_name == b"set-cookie"
+    assert header_content.startswith(b"session=null;")