Setup token validators at app initialisation.

2 files changed, 8 insertions, 2 deletions

diff --git a/gn_auth/auth/authentication/oauth2/resource_server.py b/gn_auth/auth/authentication/oauth2/resource_server.py
index c062b28..2405ee2 100644
--- a/gn_auth/auth/authentication/oauth2/resource_server.py
+++ b/gn_auth/auth/authentication/oauth2/resource_server.py
@@ -15,5 +15,3 @@ class BearerTokenValidator(_BearerTokenValidator):
                 None, lambda tok: tok)
 
 require_oauth = ResourceProtector()
-
-require_oauth.register_token_validator(BearerTokenValidator())
diff --git a/gn_auth/auth/authentication/oauth2/server.py b/gn_auth/auth/authentication/oauth2/server.py
index 09a3449..db2a0d5 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.oauth2.rfc7523 import JWTBearerTokenValidator
 from authlib.oauth2.rfc6749.errors import InvalidClientError
 from authlib.integrations.flask_oauth2 import AuthorizationServer
 
@@ -19,6 +20,8 @@ from .grants.jwt_bearer_grant import JWTBearerGrant, JWTBearerTokenGenerator
 from .endpoints.revocation import RevocationEndpoint
 from .endpoints.introspection import IntrospectionEndpoint
 
+from .resource_server import require_oauth, BearerTokenValidator
+
 def create_query_client_func() -> Callable:
     """Create the function that loads the client."""
     def __query_client__(client_id: uuid.UUID):
@@ -75,3 +78,8 @@ def setup_oauth2_server(app: Flask) -> None:
         query_client=create_query_client_func(),
         save_token=create_save_token_func(OAuth2Token))
     app.config["OAUTH2_SERVER"] = server
+
+    ## Set up the token validators
+    require_oauth.register_token_validator(BearerTokenValidator())
+    require_oauth.register_token_validator(
+        JWTBearerTokenValidator(app.config["JWT_PUBLIC_KEY"]))