Users with large JWT tokens are unauthorized (401) to access chailit due to having a cookie larger than 4kb

[jsprpl]

Some users with large JWT tokens get a 401 error when trying to log on to Chainlit. The Set-Cookie operation fails.

For these users, the login process fails for all devices/browsers.

[anthonym1991]

I can confirm that some authentication providers return larger tokens. I’m using Azure AD authentication, and the token size is about 4,400 bytes (excluding the key). The maximum cookie size is approximately 4,096 bytes.

I resolved this issue by chunking the token and storing the parts in separate cookies. After some tweaking in socket.py, I got everything working.

However, while this works as a workaround, wouldn’t it be better to store the token locally in some kind of datastore and reference it instead?

[willydouhard]

If you are up for it you could try to contribute your cookie split approach? I think that would be the easiest for everyone (rather than storing and referencing)

[anthonym1991]

Sure, here is my approach.

I added a reference in the file /auth/__init__.py:

from .cookie import OAuth2PasswordBearerWithCookie, reconstruct_token_from_cookies

In the file /auth/cookie.py i changed the code to this:

class OAuth2PasswordBearerWithCookie(SecurityBase):
    """
    OAuth2 password flow with cookie support with fallback to bearer token.
    """

    def __init__(
        self,
        tokenUrl: str,
        scheme_name: Optional[str] = None,
        auto_error: bool = True,
    ):
        self.tokenUrl = tokenUrl
        self.scheme_name = scheme_name or self.__class__.__name__
        self.auto_error = auto_error

    async def __call__(self, request: Request) -> Optional[str]:
        # First try to get the token from the cookie
        token = reconstruct_token_from_cookies(request.cookies)

        # If no cookie, try the Authorization header as fallback
        if not token:
            # TODO: Only bother to check if cookie auth is explicitly disabled.
            authorization = request.headers.get("Authorization")
            if authorization:
                scheme, token = get_authorization_scheme_param(authorization)
                if scheme.lower() != "bearer":
                    if self.auto_error:
                        raise HTTPException(
                            status_code=HTTP_401_UNAUTHORIZED,
                            detail="Invalid authentication credentials",
                            headers={"WWW-Authenticate": "Bearer"},
                        )
                    else:
                        return None
            else:
                if self.auto_error:
                    raise HTTPException(
                        status_code=HTTP_401_UNAUTHORIZED,
                        detail="Not authenticated",
                        headers={"WWW-Authenticate": "Bearer"},
                    )
                else:
                    return None

        return token

def reconstruct_token_from_cookies(request_cookies: dict) -> str:
    """
    Read all chunk cookies and reconstruct the token
    """

    print("INFO: Reconstructing token from cookies")

    # Gather all auth_chunk_i cookies, sorted by their index
    chunk_parts = []
    i = 0
    while True:
        cookie_key = f"{_auth_cookie_name}_{i}"
        if cookie_key not in request_cookies:
            break
        chunk_parts.append(request_cookies[cookie_key])
        i += 1
    
    joined = "".join(chunk_parts)

    return joined if joined != '' else None

def set_auth_cookie(response: Response, token: str):
    """
    Helper function to set the authentication cookie with secure parameters
    """

    _chunk_size = 3000
    # Split the token into multiple chunks of up to _chunk_size
    chunks = [token[i : i + _chunk_size] for i in range(0, len(token), _chunk_size)]

    # For each chunk, set a separate cookie: auth_chunk_0, auth_chunk_1, etc.
    for i, chunk in enumerate(chunks):
        cookie_key = f"{_auth_cookie_name}_{i}"
        response.set_cookie(
            key=cookie_key,
            value=chunk,
            httponly=True,
            secure=_cookie_secure,
            samesite=_cookie_samesite,
            max_age=_auth_cookie_lifetime,
        )


def clear_auth_cookie(request: Request, response: Response):
    """
    Helper function to clear the authentication cookie
    """
    
    i = 0
    while True:
        cookie_key = f"{_auth_cookie_name}_{i}"
        if cookie_key not in request.cookies:
            break
        response.delete_cookie(key=cookie_key, path="/")
        i += 1

In the file /socket.py i changed the code at line 82:

def _get_token_from_cookie(environ: WSGIEnvironment) -> Optional[str]:
    if cookie_header := environ.get("HTTP_COOKIE", None):
        cookies = cookie_parser(cookie_header)
        return reconstruct_token_from_cookies(cookies)

    return None

[willydouhard]

@dokterbob

[dokterbob]

Some users with large JWT tokens get a 401 error when trying to log on to Chainlit. The Set-Cookie operation fails.

For these users, the login process fails for all devices/browsers.

It's always M$ who think they need to do stuff in their way, isn't it? What on earth do they need such huge tokens for! ;)

Anyways, besides the fun and games, the real solution here is to store the token in the session on the backend side. #1513 In addition, this would allow essentially an unbounded session context size.

Pending aforementioned, perhaps we could roll the proof of concept above into a patch, see whether it has any unintended side effects by running CI on it?

Another workaround, for now, would be: not using such huge JWT tokens!
What is stored there that's so huge? Can it be stored in some other way or place? Do you really need all the 4kb?

[jsprpl]

@dokterbob Unfortunately the JWT token size is not something we can change. We don't really know if there will be more users than will have a token that's larger than 4kb either. Probably yes.

I tried making a pull request for this, but it mysteriously failed. I'll look into this again somewhere this week.