Merge pull request #6 from CodeSignal/store-sessions

Storing session IDs to avoid accessing protected routes with old cookies

Author: matheusgalvao1

app/middleware/auth_middleware.py

@@ -2,65 +2,63 @@
 from flask import request, jsonify, session, Blueprint
 import jwt
 from config.auth_config import AuthMethod, AuthConfig
-from services.auth_service import blacklisted_tokens
+from services.auth_service import blacklisted_tokens, invalidated_sessions
 
 class AuthMiddleware:
-    def __init__(self, auth_config: AuthConfig):
-        self.auth_config = auth_config
+    def __init__(self, config: AuthConfig):
+        self.config = config
 
-    def protect_blueprint(self, blueprint: Blueprint):
+    def protect_blueprint(self, blueprint):
+        """Add authentication middleware to all routes in a blueprint"""
         @blueprint.before_request
-        def verify_request():
-            if self.auth_config.auth_method == AuthMethod.NONE:
+        @wraps(blueprint)
+        def authenticate():
+            if self.config.auth_method == AuthMethod.NONE:
                 return None
 
-            auth_handlers = {
-                AuthMethod.API_KEY: self._handle_api_key,
-                AuthMethod.JWT: self._handle_jwt,
-                AuthMethod.SESSION: self._handle_session
-            }
+            if self.config.auth_method == AuthMethod.API_KEY:
+                return self._validate_api_key()
+            elif self.config.auth_method == AuthMethod.JWT:
+                return self._validate_jwt()
+            elif self.config.auth_method == AuthMethod.SESSION:
+                return self._validate_session()
 
-            handler = auth_handlers.get(self.auth_config.auth_method)
-            if not handler:
-                return jsonify({"error": "Invalid authentication method"}), 500
-
-            result = handler()
-            if result is not True:
-                return result
-
-            return None
-
-    def _handle_api_key(self):
+    def _validate_api_key(self):
+        """Validate API key from request header"""
         api_key = request.headers.get('X-API-Key')
         if not api_key:
             return jsonify({"error": "API key is required"}), 401
-        if api_key != self.auth_config.api_key:
+        if api_key != self.config.api_key:
             return jsonify({"error": "Invalid API key"}), 401
-        return True
+        return None
 
-    def _handle_jwt(self):
+    def _validate_jwt(self):
+        """Validate JWT from Authorization header"""
         auth_header = request.headers.get('Authorization')
         if not auth_header or not auth_header.startswith('Bearer '):
             return jsonify({"error": "JWT token is required"}), 401
-        
+
         token = auth_header.split(' ')[1]
         if token in blacklisted_tokens:
             return jsonify({"error": "You have been logged out. Please log in again."}), 401
-        
+
         try:
-            jwt.decode(
-                token, 
-                self.auth_config.jwt_secret, 
-                algorithms=["HS256"],
-                options={"verify_exp": True}
-            )
-            return True
+            jwt.decode(token, self.config.jwt_secret, algorithms=["HS256"])
+            return None
         except jwt.ExpiredSignatureError:
             return jsonify({"error": "Token has expired"}), 401
         except jwt.InvalidTokenError as e:
             return jsonify({"error": f"Invalid JWT token: {str(e)}"}), 401
 
-    def _handle_session(self):
-        if not session.get('authenticated'):
+    def _validate_session(self):
+        """Validate session authentication"""
+        if not session.get("authenticated"):
             return jsonify({"error": "Valid session required"}), 401
-        return True 
+            
+        # Check if session has been invalidated
+        current_session = request.cookies.get('session')
+        if current_session and current_session in invalidated_sessions:
+            session.clear()
+            return jsonify({"error": "Session has been invalidated"}), 401
+            
+        return None 

app/services/auth_service.py

@@ -3,7 +3,7 @@
 import secrets
 from config.auth_config import AuthConfig, AuthMethod
 from models.user import User
-from flask import session, jsonify
+from flask import session, jsonify, request
 
 # --- Configuration ---
 auth_config = None  # Global configuration object set during initialization
@@ -16,6 +16,7 @@ def init_auth_service(config: AuthConfig):
 users = []  # In-memory storage for user objects
 refresh_tokens = {}  # Maps refresh tokens to usernames
 blacklisted_tokens = set()  # Set of invalidated access tokens
+invalidated_sessions = set()  # Set of invalidated session IDs
 
 # --- User Management ---
 def is_username_taken(username):
@@ -112,9 +113,17 @@ def logout_jwt(access_token, refresh_token):
     return jsonify({"message": "Logout successful"})
 
 def logout_session():
-    """Clear user session if authenticated"""
+    """Clear user session if authenticated and invalidate the session cookie"""
     if not session.get("authenticated"):
         return jsonify({"error": "Not authenticated"}), 401
 
+    response = jsonify({"message": "Logout successful"})
+    
+    # Add current session ID to invalidated sessions set
+    if request.cookies.get('session'):
+        invalidated_sessions.add(request.cookies.get('session'))
+    
     session.clear()
-    return jsonify({"message": "Logout successful"}) 
+    # Set the session cookie to expire immediately
+    response.set_cookie('session', '', expires=0)
+    return response 

test/auth_api_collection.json

@@ -933,7 +933,40 @@
 					}
 				},
 				{
-					"name": "6. Logout with Session",
+					"name": "6. Save Session Cookie Before Logout",
+					"event": [
+						{
+							"listen": "test",
+							"script": {
+								"exec": [
+									"// Save the current session cookie for later testing",
+									"const sessionCookie = pm.cookies.get('session');",
+									"if (sessionCookie) {",
+									"    pm.environment.set('old_session', sessionCookie);",
+									"    console.log('Saved session cookie:', sessionCookie);",
+									"}",
+									"",
+									"pm.test('Status code is 200', function() {",
+									"    pm.response.to.have.status(200);",
+									"});"
+								],
+								"type": "text/javascript"
+							}
+						}
+					],
+					"request": {
+						"method": "GET",
+						"url": {
+							"raw": "http://localhost:8000/todos",
+							"protocol": "http",
+							"host": ["localhost"],
+							"port": "8000",
+							"path": ["todos"]
+						}
+					}
+				},
+				{
+					"name": "7. Logout with Session",
 					"event": [
 						{
 							"listen": "test",
@@ -959,7 +992,7 @@
 					}
 				},
 				{
-					"name": "7. Try to Get Todos after Logout (Should Fail)",
+					"name": "8. Try to Get Todos with Normal Request after Logout (Should Fail)",
 					"event": [
 						{
 							"listen": "test",
@@ -990,7 +1023,45 @@
 					}
 				},
 				{
-					"name": "8. Get API Documentation with Session",
+					"name": "9. Try to Get Todos with Old Session Cookie (Should Fail)",
+					"event": [
+						{
+							"listen": "test",
+							"script": {
+								"exec": [
+									"pm.test('Status code is 401', function() {",
+									"    pm.response.to.have.status(401);",
+									"});",
+									"",
+									"pm.test('Error message is correct', function() {",
+									"    var jsonData = pm.response.json();",
+									"    pm.expect(jsonData.error).to.equal('Session has been invalidated');",
+									"});"
+								],
+								"type": "text/javascript"
+							}
+						}
+					],
+					"request": {
+						"method": "GET",
+						"header": [
+							{
+								"key": "Cookie",
+								"value": "session={{old_session}}",
+								"type": "text"
+							}
+						],
+						"url": {
+							"raw": "http://localhost:8000/todos",
+							"protocol": "http",
+							"host": ["localhost"],
+							"port": "8000",
+							"path": ["todos"]
+						}
+					}
+				},
+				{
+					"name": "10. Get API Documentation with Session",
 					"event": [
 						{
 							"listen": "test",