feat[bckend-middleware]:Implemented custom middleware to check for blacklisted access tokens in Redis.

Author: shikharpa

services/api/kalvi/api/views/auth.py

@@ -5,7 +5,7 @@
 from rest_framework_simplejwt.views import TokenRefreshView as SimpleJWTTokenRefreshView
 from django.contrib.auth import authenticate
 from db.renderer import UserRenderer
-from rest_framework_simplejwt.tokens import RefreshToken
+from rest_framework_simplejwt.tokens import RefreshToken, AccessToken
 from rest_framework_simplejwt.exceptions import TokenError
 from rest_framework.permissions import AllowAny
 from django.utils.encoding import smart_str
@@ -155,19 +155,24 @@ def post(self, request, uid, token, format=None):
         except ValidationError as e:
             return Response({'error': e.detail}, status=status.HTTP_400_BAD_REQUEST)
 
-# Viewset class for blocking refresh tokens after logging out.       
+# Viewset class for blocking tokens after logging out.       
 class SignOutEndpoint(APIView):
     def post(self, request):
         refresh_token = request.data.get('refresh_token')
-        if refresh_token:
+        access_token = request.headers.get('Authorization').split(' ')[1] #We are catching access tokens from authorization header.
+        if refresh_token:  
             try:
                 # Connect to Redis
                 redis_conn = get_redis_connection()
-                token = str(RefreshToken(refresh_token))
-                # Blacklist the token in Redis
-                expiration_time = int(settings.SIMPLE_JWT['REFRESH_TOKEN_LIFETIME'].total_seconds())
-                redis_conn.set(token, 'blacklisted')
-                redis_conn.expire(token, expiration_time)
+                refresh_token_str = str(RefreshToken(refresh_token))
+                access_token_str = str(AccessToken(access_token))
+                # Blacklist tokens in Redis
+                refresh_exp_time = int(settings.SIMPLE_JWT['REFRESH_TOKEN_LIFETIME'].total_seconds())
+                access_exp_time = int(settings.SIMPLE_JWT['ACCESS_TOKEN_LIFETIME'].total_seconds())
+                with redis_conn.pipeline() as pipe:
+                    pipe.set(refresh_token_str, 'blacklisted', ex=refresh_exp_time)
+                    pipe.set(access_token_str, 'blacklisted', ex=access_exp_time)
+                    pipe.execute()
                 return Response({'message': 'Logged out successfully'}, status=status.HTTP_200_OK)
             except Exception:
                 return Response({'error': 'Please try after some time'}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

services/api/kalvi/middlewares/__init__.py

@@ -0,0 +1,27 @@
+from rest_framework_simplejwt.tokens import AccessToken
+from rest_framework import status
+from kalvi.api.views.auth import get_redis_connection
+from django.http import JsonResponse
+
+class CustomOutstandingTokenMiddleware:
+    """
+    This middleware checks for blacklisted tokens in Redis. It does not perform any token validation.
+    """
+    def __init__(self, get_response=None):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        try:
+            auth_header = request.headers.get('Authorization')
+            if not auth_header:
+                return self.get_response(request)
+            token = auth_header.split()[1]
+            token_obj = str(AccessToken(token))
+            redis_conn = get_redis_connection()
+            # Check if the token is blacklisted in Redis
+            if redis_conn.exists(token_obj):
+                return JsonResponse({'error': "Access token is blacklisted"}, status=status.HTTP_401_UNAUTHORIZED)
+            else:
+                return self.get_response(request)
+        except Exception:
+            return self.get_response(request)

services/api/kalvi/middlewares/tokenMiddleware.py

@@ -57,6 +57,7 @@
     "django.middleware.common.CommonMiddleware",
     "django.middleware.csrf.CsrfViewMiddleware",
     "django.contrib.auth.middleware.AuthenticationMiddleware",
+    "kalvi.middlewares.tokenMiddleware.CustomOutstandingTokenMiddleware",
     "django.contrib.messages.middleware.MessageMiddleware",
     "django.middleware.clickjacking.XFrameOptionsMiddleware",
 ]