Use IDToken Subclasses (HarryMWinters#18)

- Add token_type argument to get_config.
- Adds tests for subclassing functionality.
- Update docs for subclassing IDToken

Author: HarryMWinters

README.md

@@ -18,9 +18,11 @@
 
 ---
 
-:warning: **See [this issue](https://github.com/HarryMWinters/fastapi-oidc/issues/1) for simple role-your-own example of checking OIDC tokens.**
+:warning: **See [this issue](https://github.com/HarryMWinters/fastapi-oidc/issues/1) for
+simple role-your-own example of checking OIDC tokens.**
 
-Verify and decrypt 3rd party OIDC ID tokens to protect your [fastapi](https://github.com/tiangolo/fastapi) endpoints.
+Verify and decrypt 3rd party OIDC ID tokens to protect your
+[fastapi](https://github.com/tiangolo/fastapi) endpoints.
 
 **Documentation:** [ReadTheDocs](https://fastapi-oidc.readthedocs.io/en/latest/)
 
@@ -30,10 +32,13 @@ Verify and decrypt 3rd party OIDC ID tokens to protect your [fastapi](https://gi
 
 `pip install fastapi-oidc`
 
+## Usage
+
 ### Verify ID Tokens Issued by Third Party
 
 This is great if you just want to use something like Okta or google to handle
-your auth. All you need to do is verify the token and then you can extract user ID info from it.
+your auth. All you need to do is verify the token and then you can extract user ID info
+from it.
 
 ```python3
 from fastapi import Depends
@@ -60,3 +65,24 @@ app = FastAPI()
 def protected(id_token: IDToken = Depends(authenticate_user)):
     return {"Hello": "World", "user_email": id_token.email}
 ```
+
+#### Using your own tokens
+
+The IDToken class will accept any number of extra field but if you want to craft your
+own token class and validation that's accounted for too.
+
+```python3
+class CustomIDToken(fastapi_oidc.IDToken):
+    custom_field: str
+    custom_default: float = 3.14
+
+
+authenticate_user: Callable = get_auth(**OIDC_config, token_type=CustomIDToken)
+
+app = FastAPI()
+
+
+@app.get("/protected")
+def protected(id_token: CustomIDToken = Depends(authenticate_user)):
+    return {"Hello": "World", "user_email": id_token.custom_default}
+```

fastapi_oidc/auth.py

@@ -18,6 +18,7 @@ def test_auth(authenticated_user: AuthenticatedUser = Depends(authenticate_user)
 
 from typing import Callable
 from typing import Optional
+from typing import Type
 
 from fastapi import Depends
 from fastapi import HTTPException
@@ -28,6 +29,7 @@ def test_auth(authenticated_user: AuthenticatedUser = Depends(authenticate_user)
 from jose.exceptions import JWTClaimsError
 
 from fastapi_oidc import discovery
+from fastapi_oidc.exceptions import TokenSpecificationError
 from fastapi_oidc.types import IDToken
 
 
@@ -38,31 +40,43 @@ def get_auth(
     base_authorization_server_uri: str,
     issuer: str,
     signature_cache_ttl: int,
+    token_type: Type[IDToken] = IDToken,
 ) -> Callable[[str], IDToken]:
     """Take configurations and return the authenticate_user function.
 
     This function should only be invoked once at the beggining of your
     server code. The function it returns should be used to check user credentials.
 
     Args:
-        client_id (str): This string is provided when you register with your resource server.
+        client_id (str): This string is provided when you register with your resource
+            server.
+        base_authorization_server_uri(URL): Everything before /.wellknow in your auth
+            server URL. I.E. https://dev-123456.okta.com
+        issuer (URL): Same as base_authorization. This is used to generating OpenAPI3.0
+            docs which is broken (in OpenAPI/FastAPI) right now.
+        signature_cache_ttl (int): How many seconds your app should cache the
+            authorization server's public signatures.
         audience (str): (Optional) The audience string configured by your auth server.
             If not set defaults to client_id
-        base_authorization_server_uri(URL): Everything before /.wellknow in your auth server URL.
-            I.E. https://dev-123456.okta.com
-        issuer (URL): Same as base_authorization. This is used to generating OpenAPI3.0 docs which
-            is broken (in OpenAPI/FastAPI) right now.
-        signature_cache_ttl (int): How many seconds your app should cache the authorization
-            server's public signatures.
+        token_type (IDToken or subclass): (Optional) An optional class to be returned by
+            the authenticate_user function.
 
 
     Returns:
-        func: authenticate_user(auth_header: str)
+        func: authenticate_user(auth_header: str) -> IDToken (or token_type)
 
     Raises:
         Nothing intentional
     """
-    # As far as I can tell this does two things.
+
+    if not issubclass(token_type, IDToken):
+        raise TokenSpecificationError(
+            "Invalid argument for token_type. "
+            "Token type must be a subclass of fastapi_oidc.type.IDToken. "
+            f"Received {token_type=}"
+        )
+
+    # As far as I can tell the oauth2_scheme does two things.
     # 1. Extracts and returns the Authorization header.
     # 2. Integrates with the OpenAPI3.0 doc generation in FastAPI.
     #    This integration doesn't matter much now since OpenAPI
@@ -79,8 +93,8 @@ def authenticate_user(auth_header: str = Depends(oauth2_scheme)) -> IDToken:
         for signature_cache_ttl seconds.
 
         Args:
-            auth_header (str): Base64 encoded OIDC Token. This is invoked behind the scenes
-                by Depends.
+            auth_header (str): Base64 encoded OIDC Token. This is invoked behind the
+                scenes by Depends.
 
         Return:
             IDToken (types.IDToken):
@@ -103,27 +117,9 @@ def authenticate_user(auth_header: str = Depends(oauth2_scheme)) -> IDToken:
                 # Disabled at_hash check since we aren't using the access token
                 options={"verify_at_hash": False},
             )
-            return IDToken.parse_obj(token)
+            return token_type.parse_obj(token)
 
         except (ExpiredSignatureError, JWTError, JWTClaimsError) as err:
             raise HTTPException(status_code=401, detail=f"Unauthorized: {err}")
 
     return authenticate_user
-
-
-# This is a dummy method for sphinx docs. DO NOT User.
-# TODO Find a way to doc higher order functions w/ sphinx.
-def authenticate_user(auth_header: str) -> IDToken:  # type: ignore
-    """
-    Validate and parse OIDC ID token against issuer in config.
-    Note this function caches the signatures and algorithms of the issuing server
-    for signature_cache_ttl seconds.
-
-    Args:
-        auth_header (str): Base64 encoded OIDC Token. This is invoked behind the scenes
-            by Depends.
-
-    Return:
-        IDToken (types.IDToken):
-    """
-    pass

fastapi_oidc/exceptions.py

@@ -0,0 +1,2 @@
+class TokenSpecificationError(BaseException):
+    pass

tests/test_auth.py

@@ -1,4 +1,7 @@
+import pytest
+
 from fastapi_oidc import auth
+from fastapi_oidc.exceptions import TokenSpecificationError
 from fastapi_oidc.types import IDToken
 
 
@@ -40,3 +43,28 @@ def test__authenticate_user_no_aud(
 
     assert id_token.email == test_email  # nosec
     assert id_token.aud == no_audience_config["client_id"]
+
+
+def test__get_auth_raises_if_token_type_is_not_subclass_of_IDToken(no_audience_config):
+    class BadToken:
+        pass
+
+    with pytest.raises(TokenSpecificationError):
+        auth.get_auth(**no_audience_config, token_type=BadToken)
+
+
+def test__authenticate_user_returns_custom_tokens(
+    monkeypatch, mock_discovery, token_without_audience, no_audience_config
+):
+    class CustomToken(IDToken):
+        custom_field: str = "OnlySlightlyBent"
+
+    monkeypatch.setattr(auth.discovery, "configure", mock_discovery)
+
+    token = token_without_audience
+
+    authenticate_user = auth.get_auth(**no_audience_config, token_type=CustomToken)
+
+    custom_token: CustomToken = authenticate_user(auth_header=f"Bearer {token}")
+
+    assert custom_token.custom_field == "OnlySlightlyBent"