Fix #701: factorize JWT handling and support secrets as SecretStr

fastapi_users/authentication/cookie.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, List, Optional
 
 import jwt
 from fastapi import Response
@@ -7,8 +7,8 @@ from pydantic import UUID4
 
 from fastapi_users.authentication import BaseAuthentication
 from fastapi_users.db.base import BaseUserDatabase
+from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
 from fastapi_users.models import BaseUserDB
-from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 
 class CookieAuthentication(BaseAuthentication[str]):
@@ -25,11 +25,12 @@ class CookieAuthentication(BaseAuthentication[str]):
     :param cookie_secure: Whether to only send the cookie to the server via SSL request.
     :param cookie_httponly: Whether to prevent access to the cookie via JavaScript.
     :param name: Name of the backend. It will be used to name the login route.
+    :param token_audience: List of valid audiences for the JWT.
     """
 
     scheme: APIKeyCookie
-    token_audience: str = "fastapi-users:auth"
-    secret: str
+    token_audience: List[str]
+    secret: SecretType
     lifetime_seconds: Optional[int]
     cookie_name: str
     cookie_path: str
@@ -40,7 +41,7 @@ class CookieAuthentication(BaseAuthentication[str]):
 
     def __init__(
         self,
-        secret: str,
+        secret: SecretType,
         lifetime_seconds: Optional[int] = None,
         cookie_name: str = "fastapiusersauth",
         cookie_path: str = "/",
@@ -49,6 +50,7 @@ class CookieAuthentication(BaseAuthentication[str]):
         cookie_httponly: bool = True,
         cookie_samesite: str = "lax",
         name: str = "cookie",
+        token_audience: List[str] = ["fastapi-users:auth"],
     ):
         super().__init__(name, logout=True)
         self.secret = secret
@@ -59,6 +61,7 @@ class CookieAuthentication(BaseAuthentication[str]):
         self.cookie_secure = cookie_secure
         self.cookie_httponly = cookie_httponly
         self.cookie_samesite = cookie_samesite
+        self.token_audience = token_audience
         self.scheme = APIKeyCookie(name=self.cookie_name, auto_error=False)
 
     async def __call__(
@@ -70,12 +73,7 @@ class CookieAuthentication(BaseAuthentication[str]):
             return None
 
         try:
-            data = jwt.decode(
-                credentials,
-                self.secret,
-                audience=self.token_audience,
-                algorithms=[JWT_ALGORITHM],
-            )
+            data = decode_jwt(credentials, self.secret, self.token_audience)
             user_id = data.get("user_id")
             if user_id is None:
                 return None
@@ -112,4 +110,4 @@ class CookieAuthentication(BaseAuthentication[str]):
 
     async def _generate_token(self, user: BaseUserDB) -> str:
         data = {"user_id": str(user.id), "aud": self.token_audience}
-        return generate_jwt(data, self.secret, self.lifetime_seconds, JWT_ALGORITHM)
+        return generate_jwt(data, self.secret, self.lifetime_seconds)

fastapi_users/authentication/jwt.py

@@ -7,8 +7,8 @@ from pydantic import UUID4
 
 from fastapi_users.authentication.base import BaseAuthentication
 from fastapi_users.db.base import BaseUserDatabase
+from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
 from fastapi_users.models import BaseUserDB
-from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 
 class JWTAuthentication(BaseAuthentication[str]):
@@ -19,11 +19,12 @@ class JWTAuthentication(BaseAuthentication[str]):
     :param lifetime_seconds: Lifetime duration of the JWT in seconds.
     :param tokenUrl: Path where to get a token.
     :param name: Name of the backend. It will be used to name the login route.
+    :param token_audience: List of valid audiences for the JWT.
     """
 
     scheme: OAuth2PasswordBearer
-    token_audience: List[str] = ["fastapi-users:auth"]
-    secret: str
+    token_audience: List[str]
+    secret: SecretType
     lifetime_seconds: int
 
     def __init__(
@@ -49,12 +50,7 @@ class JWTAuthentication(BaseAuthentication[str]):
             return None
 
         try:
-            data = jwt.decode(
-                credentials,
-                self.secret,
-                audience=self.token_audience,
-                algorithms=[JWT_ALGORITHM],
-            )
+            data = decode_jwt(credentials, self.secret, self.token_audience)
             user_id = data.get("user_id")
             if user_id is None:
                 return None
@@ -73,4 +69,4 @@ class JWTAuthentication(BaseAuthentication[str]):
 
     async def _generate_token(self, user: BaseUserDB) -> str:
         data = {"user_id": str(user.id), "aud": self.token_audience}
-        return generate_jwt(data, self.secret, self.lifetime_seconds, JWT_ALGORITHM)
+        return generate_jwt(data, self.secret, self.lifetime_seconds)

fastapi_users/fastapi_users.py

@@ -5,6 +5,7 @@ from fastapi import APIRouter, Request
 from fastapi_users import models
 from fastapi_users.authentication import Authenticator, BaseAuthentication
 from fastapi_users.db import BaseUserDatabase
+from fastapi_users.jwt import SecretType
 from fastapi_users.router import (
     get_auth_router,
     get_register_router,
@@ -132,7 +133,7 @@ class FastAPIUsers:
 
     def get_verify_router(
         self,
-        verification_token_secret: str,
+        verification_token_secret: SecretType,
         verification_token_lifetime_seconds: int = 3600,
         after_verification_request: Optional[
             Callable[[models.UD, str, Request], None]
@@ -161,7 +162,7 @@ class FastAPIUsers:
 
     def get_reset_password_router(
         self,
-        reset_password_token_secret: str,
+        reset_password_token_secret: SecretType,
         reset_password_token_lifetime_seconds: int = 3600,
         after_forgot_password: Optional[
             Callable[[models.UD, str, Request], None]
@@ -207,7 +208,7 @@ class FastAPIUsers:
     def get_oauth_router(
         self,
         oauth_client: BaseOAuth2,
-        state_secret: str,
+        state_secret: SecretType,
         redirect_url: str = None,
         after_register: Optional[Callable[[models.UD, Request], None]] = None,
     ) -> APIRouter:

fastapi_users/jwt.py

@@ -0,0 +1,41 @@
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union
+
+import jwt
+from pydantic import SecretStr
+
+SecretType = Union[str, SecretStr]
+JWT_ALGORITHM = "HS256"
+
+
+def _get_secret_value(secret: SecretType) -> str:
+    if isinstance(secret, SecretStr):
+        return secret.get_secret_value()
+    return secret
+
+
+def generate_jwt(
+    data: dict,
+    secret: SecretType,
+    lifetime_seconds: Optional[int] = None,
+    algorithm: str = JWT_ALGORITHM,
+) -> str:
+    payload = data.copy()
+    if lifetime_seconds:
+        expire = datetime.utcnow() + timedelta(seconds=lifetime_seconds)
+        payload["exp"] = expire
+    return jwt.encode(payload, _get_secret_value(secret), algorithm=algorithm)
+
+
+def decode_jwt(
+    encoded_jwt: str,
+    secret: SecretType,
+    audience: List[str],
+    algorithms: List[str] = [JWT_ALGORITHM],
+) -> Dict[str, Any]:
+    return jwt.decode(
+        encoded_jwt,
+        _get_secret_value(secret),
+        audience=audience,
+        algorithms=algorithms,
+    )

fastapi_users/router/oauth.py

@@ -8,27 +8,18 @@ from httpx_oauth.oauth2 import BaseOAuth2
 from fastapi_users import models
 from fastapi_users.authentication import Authenticator
 from fastapi_users.db import BaseUserDatabase
+from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
 from fastapi_users.password import generate_password, get_password_hash
 from fastapi_users.router.common import ErrorCode, run_handler
-from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 STATE_TOKEN_AUDIENCE = "fastapi-users:oauth-state"
 
 
 def generate_state_token(
-    data: Dict[str, str], secret: str, lifetime_seconds: int = 3600
+    data: Dict[str, str], secret: SecretType, lifetime_seconds: int = 3600
 ) -> str:
     data["aud"] = STATE_TOKEN_AUDIENCE
-    return generate_jwt(data, secret, lifetime_seconds, JWT_ALGORITHM)
-
-
-def decode_state_token(token: str, secret: str) -> Dict[str, str]:
-    return jwt.decode(
-        token,
-        secret,
-        audience=STATE_TOKEN_AUDIENCE,
-        algorithms=[JWT_ALGORITHM],
-    )
+    return generate_jwt(data, secret, lifetime_seconds)
 
 
 def get_oauth_router(
@@ -36,7 +27,7 @@ def get_oauth_router(
     user_db: BaseUserDatabase[models.BaseUserDB],
     user_db_model: Type[models.BaseUserDB],
     authenticator: Authenticator,
-    state_secret: str,
+    state_secret: SecretType,
     redirect_url: str = None,
     after_register: Optional[Callable[[models.UD, Request], None]] = None,
 ) -> APIRouter:
@@ -99,7 +90,7 @@ def get_oauth_router(
         )
 
         try:
-            state_data = decode_state_token(state, state_secret)
+            state_data = decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
         except jwt.DecodeError:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
 

fastapi_users/router/reset.py

@@ -6,17 +6,17 @@ from pydantic import UUID4, EmailStr
 
 from fastapi_users import models
 from fastapi_users.db import BaseUserDatabase
+from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
 from fastapi_users.password import get_password_hash
 from fastapi_users.router.common import ErrorCode, run_handler
 from fastapi_users.user import InvalidPasswordException, ValidatePasswordProtocol
-from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 RESET_PASSWORD_TOKEN_AUDIENCE = "fastapi-users:reset"
 
 
 def get_reset_password_router(
     user_db: BaseUserDatabase[models.BaseUserDB],
-    reset_password_token_secret: str,
+    reset_password_token_secret: SecretType,
     reset_password_token_lifetime_seconds: int = 3600,
     after_forgot_password: Optional[Callable[[models.UD, str, Request], None]] = None,
     after_reset_password: Optional[Callable[[models.UD, Request], None]] = None,
@@ -48,11 +48,8 @@ def get_reset_password_router(
         request: Request, token: str = Body(...), password: str = Body(...)
     ):
         try:
-            data = jwt.decode(
-                token,
-                reset_password_token_secret,
-                audience=RESET_PASSWORD_TOKEN_AUDIENCE,
-                algorithms=[JWT_ALGORITHM],
+            data = decode_jwt(
+                token, reset_password_token_secret, [RESET_PASSWORD_TOKEN_AUDIENCE]
             )
             user_id = data.get("user_id")
             if user_id is None:

fastapi_users/router/verify.py

@@ -5,6 +5,7 @@ from fastapi import APIRouter, Body, HTTPException, Request, status
 from pydantic import UUID4, EmailStr
 
 from fastapi_users import models
+from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
 from fastapi_users.router.common import ErrorCode, run_handler
 from fastapi_users.user import (
     GetUserProtocol,
@@ -12,7 +13,6 @@ from fastapi_users.user import (
     UserNotExists,
     VerifyUserProtocol,
 )
-from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 VERIFY_USER_TOKEN_AUDIENCE = "fastapi-users:verify"
 
@@ -21,7 +21,7 @@ def get_verify_router(
     verify_user: VerifyUserProtocol,
     get_user: GetUserProtocol,
     user_model: Type[models.BaseUser],
-    verification_token_secret: str,
+    verification_token_secret: SecretType,
     verification_token_lifetime_seconds: int = 3600,
     after_verification_request: Optional[
         Callable[[models.UD, str, Request], None]
@@ -58,11 +58,8 @@ def get_verify_router(
     @router.post("/verify", response_model=user_model)
     async def verify(request: Request, token: str = Body(..., embed=True)):
         try:
-            data = jwt.decode(
-                token,
-                verification_token_secret,
-                audience=VERIFY_USER_TOKEN_AUDIENCE,
-                algorithms=[JWT_ALGORITHM],
+            data = decode_jwt(
+                token, verification_token_secret, [VERIFY_USER_TOKEN_AUDIENCE]
             )
         except jwt.exceptions.ExpiredSignatureError:
             raise HTTPException(

fastapi_users/utils.py

@@ -1,19 +0,0 @@
-from datetime import datetime, timedelta
-from typing import Optional
-
-import jwt
-
-JWT_ALGORITHM = "HS256"
-
-
-def generate_jwt(
-    data: dict,
-    secret: str,
-    lifetime_seconds: Optional[int] = None,
-    algorithm: str = JWT_ALGORITHM,
-) -> str:
-    payload = data.copy()
-    if lifetime_seconds:
-        expire = datetime.utcnow() + timedelta(seconds=lifetime_seconds)
-        payload["exp"] = expire
-    return jwt.encode(payload, secret, algorithm=algorithm)