fix: better fastapi middleware (#505)

* fix: better fastapi middleware

* fix: changelog and version

* fix: comments

Author: sattvikc

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased]
 
+## [0.20.2] - 2024-05-17
+
+-   Improves FastAPI middleware performance using recommended ASGI middleware implementation.
+
 ## [0.20.1] - 2024-05-10
 
 -   Fixes parameter mismatch in generating fake email

setup.py

@@ -83,7 +83,7 @@
 
 setup(
     name="supertokens_python",
-    version="0.20.1",
+    version="0.20.2",
     author="SuperTokens",
     license="Apache 2.0",
     author_email="[email protected]",

supertokens_python/constants.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 SUPPORTED_CDI_VERSIONS = ["3.0"]
-VERSION = "0.20.1"
+VERSION = "0.20.2"
 TELEMETRY = "/telemetry"
 USER_COUNT = "/users/count"
 USER_DELETE = "/user/remove"

supertokens_python/framework/fastapi/fastapi_middleware.py

@@ -11,64 +11,100 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations
 # under the License.
-from __future__ import annotations
+from typing import Union
 
-from typing import TYPE_CHECKING, Union
 
-from supertokens_python.framework import BaseResponse
+def get_middleware():
+    from supertokens_python import Supertokens
+    from supertokens_python.utils import default_user_context
+    from supertokens_python.exceptions import SuperTokensError
+    from supertokens_python.framework import BaseResponse
+    from supertokens_python.recipe.session import SessionContainer
+    from supertokens_python.supertokens import manage_session_post_response
 
-if TYPE_CHECKING:
-    from fastapi import Request
+    from starlette.requests import Request
+    from starlette.responses import Response
+    from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
+    from supertokens_python.framework.fastapi.fastapi_request import (
+        FastApiRequest,
+    )
+    from supertokens_python.framework.fastapi.fastapi_response import (
+        FastApiResponse,
+    )
 
-def get_middleware():
-    from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
-    from supertokens_python.utils import default_user_context
+    class ASGIMiddleware:
+        def __init__(self, app: ASGIApp) -> None:
+            self.app = app
 
-    class Middleware(BaseHTTPMiddleware):
-        async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
-            from supertokens_python import Supertokens
-            from supertokens_python.exceptions import SuperTokensError
-            from supertokens_python.framework.fastapi.fastapi_request import (
-                FastApiRequest,
-            )
-            from supertokens_python.framework.fastapi.fastapi_response import (
-                FastApiResponse,
-            )
-            from supertokens_python.recipe.session import SessionContainer
-            from supertokens_python.supertokens import manage_session_post_response
+        async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+            if scope["type"] != "http":  # we pass through the non-http requests, if any
+                await self.app(scope, receive, send)
+                return
 
             st = Supertokens.get_instance()
-            from fastapi.responses import Response
 
+            request = Request(scope, receive=receive)
             custom_request = FastApiRequest(request)
-            response = FastApiResponse(Response())
             user_context = default_user_context(custom_request)
 
             try:
+                response = FastApiResponse(Response())
                 result: Union[BaseResponse, None] = await st.middleware(
                     custom_request, response, user_context
                 )
                 if result is None:
-                    response = await call_next(request)
-                    result = FastApiResponse(response)
+                    # This means that the supertokens middleware did not handle the request,
+                    # however, we may need to handle the header changes in the response,
+                    # based on response mutators used by the session.
+                    async def send_wrapper(message: Message):
+                        if message["type"] == "http.response.start":
+                            # Start message has the headers, so we update the headers here
+                            # by using `manage_session_post_response` function, which will
+                            # apply all the Response Mutators. In the end, we just replace
+                            # the updated headers in the message.
+                            if hasattr(request.state, "supertokens") and isinstance(
+                                request.state.supertokens, SessionContainer
+                            ):
+                                fapi_response = Response()
+                                fapi_response.raw_headers = message["headers"]
+                                response = FastApiResponse(fapi_response)
+                                manage_session_post_response(
+                                    request.state.supertokens, response, user_context
+                                )
+                                message["headers"] = fapi_response.raw_headers
 
+                        # For `http.response.start` message, we might have the headers updated,
+                        # otherwise, we just send all the messages as is
+                        await send(message)
+
+                    await self.app(scope, receive, send_wrapper)
+                    return
+
+                # This means that the request was handled by the supertokens middleware
+                # and hence we respond using the response object returned by the middleware.
                 if hasattr(request.state, "supertokens") and isinstance(
                     request.state.supertokens, SessionContainer
                 ):
                     manage_session_post_response(
                         request.state.supertokens, result, user_context
                     )
+
                 if isinstance(result, FastApiResponse):
-                    return result.response
+                    await result.response(scope, receive, send)
+                    return
+
+                return
+
             except SuperTokensError as e:
                 response = FastApiResponse(Response())
                 result: Union[BaseResponse, None] = await st.handle_supertokens_error(
                     FastApiRequest(request), e, response, user_context
                 )
                 if isinstance(result, FastApiResponse):
-                    return result.response
+                    await result.response(scope, receive, send)
+                    return
 
             raise Exception("Should never come here")
 
-    return Middleware
+    return ASGIMiddleware