GH-9: Add callback argument to middleware

Author: ArtyomVancyan

src/fastapi_oauth2/middleware.py

@@ -1,5 +1,7 @@
 from datetime import datetime
 from datetime import timedelta
+from typing import Awaitable
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -87,13 +89,20 @@ def identity(self) -> str:
 
 
 class OAuth2Backend(AuthenticationBackend):
-    def __init__(self, config: OAuth2Config) -> None:
+    """Authentication backend for AuthenticationMiddleware."""
+
+    def __init__(
+            self,
+            config: OAuth2Config,
+            callback: Callable[[User], Union[Awaitable[None], None]] = None,
+    ) -> None:
         Auth.set_http(config.allow_http)
         Auth.set_secret(config.jwt_secret)
         Auth.set_expires(config.jwt_expires)
         Auth.set_algorithm(config.jwt_algorithm)
         for client in config.clients:
             Auth.register_client(client)
+        self.callback = callback
 
     async def authenticate(self, request: Request) -> Optional[Tuple["Auth", "User"]]:
         authorization = request.headers.get(
@@ -106,18 +115,39 @@ async def authenticate(self, request: Request) -> Optional[Tuple["Auth", "User"]
             return Auth(), User()
 
         user = Auth.jwt_decode(param)
-        return Auth(user.pop("scope", [])), User(user)
+        auth, user = Auth(user.pop("scope", [])), User(user)
+
+        # Call the callback function on authentication
+        if callable(self.callback):
+            coroutine = self.callback(user)
+            if issubclass(type(coroutine), Awaitable):
+                await coroutine
+        return auth, user
 
 
 class OAuth2Middleware:
+    """Wrapper for the Starlette AuthenticationMiddleware."""
+
     auth_middleware: AuthenticationMiddleware = None
 
-    def __init__(self, app: ASGIApp, config: Union[OAuth2Config, dict]) -> None:
+    def __init__(
+            self,
+            app: ASGIApp,
+            config: Union[OAuth2Config, dict],
+            callback: Callable[[User], Union[Awaitable[None], None]] = None,
+            **kwargs,  # AuthenticationMiddleware kwargs
+    ) -> None:
+        """Initiates the middleware with the given configuration.
+
+        :param app: FastAPI application instance
+        :param config: middleware configuration
+        :param callback: callback function to be called after authentication
+        """
         if isinstance(config, dict):
             config = OAuth2Config(**config)
         elif not isinstance(config, OAuth2Config):
             raise TypeError("config is not a valid type")
-        self.auth_middleware = AuthenticationMiddleware(app, OAuth2Backend(config))
+        self.auth_middleware = AuthenticationMiddleware(app, backend=OAuth2Backend(config, callback), **kwargs)
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         await self.auth_middleware(scope, receive, send)