Save token with same ID as JWT's "jti" value.

1 files changed, 16 insertions, 5 deletions

diff --git a/gn_auth/auth/authentication/oauth2/server.py b/gn_auth/auth/authentication/oauth2/server.py
index 0669139..5ed7785 100644
--- a/gn_auth/auth/authentication/oauth2/server.py
+++ b/gn_auth/auth/authentication/oauth2/server.py
@@ -4,6 +4,7 @@ import datetime
 from typing import Callable
 
 from flask import Flask, current_app
+from authlib.jose import jwk, jwt
 from authlib.oauth2.rfc7523 import JWTBearerTokenValidator
 from authlib.oauth2.rfc6749.errors import InvalidClientError
 from authlib.integrations.flask_oauth2 import AuthorizationServer
@@ -37,13 +38,22 @@ def create_query_client_func() -> Callable:
 
     return __query_client__
 
-def create_save_token_func(token_model: type) -> Callable:
+def create_save_token_func(token_model: type, jwtkey: jwk) -> Callable:
     """Create the function that saves the token."""
     def __save_token__(token, request):
+        _jwt = jwt.decode(token["access_token"], jwtkey)
+        _token = token_model(
+            token_id=uuid.UUID(_jwt["jti"]),
+            client=request.client,
+            user=request.user,
+            **{
+                "refresh_token": None,
+                "revoked": False,
+                "issued_at": datetime.datetime.now(),
+                **token
+            })
         with db.connection(current_app.config["AUTH_DB"]) as conn:
-            save_token(
-                conn, token_model(
-                    token_id=uuid.uuid4(), client=request.client,
+            save_token(conn, _token)
                     user=request.user,
                     **{
                         "refresh_token": None, "revoked": False,
@@ -76,7 +86,8 @@ def setup_oauth2_server(app: Flask) -> None:
     server.init_app(
         app,
         query_client=create_query_client_func(),
-        save_token=create_save_token_func(OAuth2Token))
+        save_token=create_save_token_func(
+            OAuth2Token, app.config["SSL_PRIVATE_KEY"]))
     app.config["OAUTH2_SERVER"] = server
 
     ## Set up the token validators