feat(contrib-jwt): Add parameters to set the JWT extras field (#2313)

* feat(contrib-jwt): Add parameters to set the JWT `extras` field

- Add `token_extras` to both `login()` and `create_token()` methods.
- Update tests and docs example to use the new parameter.

* fix(contrib-jwt): Use `typing.Dict` for 3.8 compatibility

docs/examples/contrib/jwt/using_jwt_auth.py

@@ -48,7 +48,7 @@ async def login_handler(data: User) -> Response[User]:
     MOCK_DB[str(data.id)] = data
     # you can do whatever you want to update the response instance here
     # e.g. response.set_cookie(...)
-    return jwt_auth.login(identifier=str(data.id), response_body=data)
+    return jwt_auth.login(identifier=str(data.id), token_extras={"email": data.email}, response_body=data)
 
 
 # We also have some other routes, for example:

litestar/contrib/jwt/jwt_auth.py

@@ -127,6 +127,7 @@ def login(
         token_issuer: str | None = None,
         token_audience: str | None = None,
         token_unique_jwt_id: str | None = None,
+        token_extras: dict[str, Any] | None = None,
         send_token_as_response_body: bool = False,
     ) -> Response[Any]:
         """Create a response with a JWT header.
@@ -140,6 +141,7 @@ def login(
             token_issuer: An optional value of the token ``iss`` field.
             token_audience: An optional value for the token ``aud`` field.
             token_unique_jwt_id: An optional value for the token ``jti`` field.
+            token_extras: An optional dictionary to include in the token ``extras`` field.
             send_token_as_response_body: If ``True`` the response will be a dict including the token: ``{ "token": <token> }``
                 will be returned as the response body. Note: if a response body is passed this setting will be ignored.
 
@@ -152,6 +154,7 @@ def login(
             token_issuer=token_issuer,
             token_audience=token_audience,
             token_unique_jwt_id=token_unique_jwt_id,
+            token_extras=token_extras,
         )
 
         if response_body is not Empty:
@@ -175,6 +178,7 @@ def create_token(
         token_issuer: str | None = None,
         token_audience: str | None = None,
         token_unique_jwt_id: str | None = None,
+        token_extras: dict | None = None,
     ) -> str:
         """Create a Token instance from the passed in parameters, persists and returns it.
 
@@ -184,6 +188,7 @@ def create_token(
             token_issuer: An optional value of the token ``iss`` field.
             token_audience: An optional value for the token ``aud`` field.
             token_unique_jwt_id: An optional value for the token ``jti`` field.
+            token_extras: An optional dictionary to include in the token ``extras`` field.
 
         Returns:
             The created token.
@@ -194,6 +199,7 @@ def create_token(
             iss=token_issuer,
             aud=token_audience,
             jti=token_unique_jwt_id,
+            extras=token_extras or {},
         )
         return token.encode(secret=self.token_secret, algorithm=self.algorithm)
 
@@ -404,6 +410,7 @@ def login(
         token_issuer: str | None = None,
         token_audience: str | None = None,
         token_unique_jwt_id: str | None = None,
+        token_extras: dict[str, Any] | None = None,
         send_token_as_response_body: bool = False,
     ) -> Response[Any]:
         """Create a response with a JWT header.
@@ -417,6 +424,7 @@ def login(
             token_issuer: An optional value of the token ``iss`` field.
             token_audience: An optional value for the token ``aud`` field.
             token_unique_jwt_id: An optional value for the token ``jti`` field.
+            token_extras: An optional dictionary to include in the token ``extras`` field.
             send_token_as_response_body: If ``True`` the response will be a dict including the token: ``{ "token": <token> }``
                 will be returned as the response body. Note: if a response body is passed this setting will be ignored.
 
@@ -430,6 +438,7 @@ def login(
             token_issuer=token_issuer,
             token_audience=token_audience,
             token_unique_jwt_id=token_unique_jwt_id,
+            token_extras=token_extras,
         )
         cookie = Cookie(
             key=self.key,
@@ -620,6 +629,7 @@ def login(
         token_issuer: str | None = None,
         token_audience: str | None = None,
         token_unique_jwt_id: str | None = None,
+        token_extras: dict[str, Any] | None = None,
         send_token_as_response_body: bool = True,
     ) -> Response[Any]:
         """Create a response with a JWT header.
@@ -633,6 +643,7 @@ def login(
             token_issuer: An optional value of the token ``iss`` field.
             token_audience: An optional value for the token ``aud`` field.
             token_unique_jwt_id: An optional value for the token ``jti`` field.
+            token_extras: An optional dictionary to include in the token ``extras`` field.
             send_token_as_response_body: If ``True`` the response will be an oAuth2 token response dict.
                 Note: if a response body is passed this setting will be ignored.
 
@@ -645,6 +656,7 @@ def login(
             token_issuer=token_issuer,
             token_audience=token_audience,
             token_unique_jwt_id=token_unique_jwt_id,
+            token_extras=token_extras,
         )
         expires_in = int((token_expiration or self.default_token_expiration).total_seconds())
         cookie = Cookie(

tests/unit/test_contrib/test_jwt/test_auth.py

@@ -5,7 +5,7 @@
 
 import pytest
 from hypothesis import given, settings
-from hypothesis.strategies import integers, none, one_of, sampled_from, text, timedeltas
+from hypothesis.strategies import dictionaries, integers, none, one_of, sampled_from, text, timedeltas
 from pydantic import BaseModel, Field
 
 from litestar import Litestar, Request, Response, get
@@ -41,6 +41,7 @@ def mock_db() -> MemoryStore:
     token_issuer=one_of(none(), text(max_size=256)),
     token_audience=one_of(none(), text(max_size=256, alphabet=string.ascii_letters)),
     token_unique_jwt_id=one_of(none(), text(max_size=256)),
+    token_extras=one_of(none(), dictionaries(text(max_size=256), text(max_size=256))),
 )
 @settings(deadline=None)
 async def test_jwt_auth(
@@ -54,6 +55,7 @@ async def test_jwt_auth(
     token_issuer: Optional[str],
     token_audience: Optional[str],
     token_unique_jwt_id: Optional[str],
+    token_extras: Optional[Dict[str, Any]],
 ) -> None:
     user = UserFactory.build()
 
@@ -86,6 +88,7 @@ def login_handler() -> Response["User"]:
             token_issuer=token_issuer,
             token_audience=token_audience,
             token_unique_jwt_id=token_unique_jwt_id,
+            token_extras=token_extras,
         )
 
     with create_test_client(route_handlers=[my_handler, login_handler]) as client:
@@ -98,6 +101,9 @@ def login_handler() -> Response["User"]:
         assert decoded_token.iss == token_issuer
         assert decoded_token.aud == token_audience
         assert decoded_token.jti == token_unique_jwt_id
+        if token_extras is not None:
+            for key, value in token_extras.items():
+                assert decoded_token.extras[key] == value
 
         response = client.get("/my-endpoint")
         assert response.status_code == HTTP_401_UNAUTHORIZED
@@ -140,6 +146,7 @@ def login_handler() -> Response["User"]:
     token_issuer=one_of(none(), text(max_size=256)),
     token_audience=one_of(none(), text(max_size=256, alphabet=string.ascii_letters)),
     token_unique_jwt_id=one_of(none(), text(max_size=256)),
+    token_extras=one_of(none(), dictionaries(text(max_size=256), text(max_size=256))),
 )
 @settings(deadline=None)
 async def test_jwt_cookie_auth(
@@ -154,6 +161,7 @@ async def test_jwt_cookie_auth(
     token_issuer: Optional[str],
     token_audience: Optional[str],
     token_unique_jwt_id: Optional[str],
+    token_extras: Optional[Dict[str, Any]],
 ) -> None:
     user = UserFactory.build()
 
@@ -188,6 +196,7 @@ def login_handler() -> Response["User"]:
             token_issuer=token_issuer,
             token_audience=token_audience,
             token_unique_jwt_id=token_unique_jwt_id,
+            token_extras=token_extras,
         )
 
     with create_test_client(route_handlers=[my_handler, login_handler]) as client:
@@ -200,6 +209,9 @@ def login_handler() -> Response["User"]:
         assert decoded_token.iss == token_issuer
         assert decoded_token.aud == token_audience
         assert decoded_token.jti == token_unique_jwt_id
+        if token_extras is not None:
+            for key, value in token_extras.items():
+                assert decoded_token.extras[key] == value
 
         client.cookies.clear()
         response = client.get("/my-endpoint")

tests/unit/test_contrib/test_jwt/test_token.py

@@ -2,7 +2,7 @@
 import sys
 from dataclasses import asdict
 from datetime import datetime, timedelta, timezone
-from typing import Optional
+from typing import Any, Dict, Optional
 from uuid import uuid4
 
 import pytest
@@ -20,11 +20,13 @@
 @pytest.mark.parametrize(
     "token_unique_jwt_id", [None, "10f5c6967783ddd6bb0c4e8262d7097caeae64705e45f83275e3c32eee5d30f2"]
 )
+@pytest.mark.parametrize("token_extras", [None, {"email": "[email protected]"}])
 def test_token(
     algorithm: str,
     token_issuer: Optional[str],
     token_audience: Optional[str],
     token_unique_jwt_id: Optional[str],
+    token_extras: Optional[Dict[str, Any]],
 ) -> None:
     token_secret = secrets.token_hex()
     token = Token(
@@ -33,6 +35,7 @@ def test_token(
         aud=token_audience,
         iss=token_issuer,
         jti=token_unique_jwt_id,
+        extras=token_extras or {},
     )
     encoded_token = token.encode(secret=token_secret, algorithm=algorithm)
     decoded_token = token.decode(encoded_token=encoded_token, secret=token_secret, algorithm=algorithm)