Skip to content

Commit 3756c9c

Browse files
matheusgalvao1matheusgalvao1
authored
Merge pull request #6 from CodeSignal/store-sessions
Storing session IDs to avoid accessing protected routes with old cookies
2 parents 00a4627 + 47aaefe commit 3756c9c

3 files changed

Lines changed: 121 additions & 43 deletions

File tree

‎app/middleware/auth_middleware.py‎

Lines changed: 35 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -2,65 +2,63 @@
22
from flask import request, jsonify, session, Blueprint
33
import jwt
44
from config.auth_config import AuthMethod, AuthConfig
5-
from services.auth_service import blacklisted_tokens
5+
from services.auth_service import blacklisted_tokens, invalidated_sessions
66

77
class AuthMiddleware:
8-
def __init__(self, auth_config: AuthConfig):
9-
self.auth_config = auth_config
8+
def __init__(self, config: AuthConfig):
9+
self.config = config
1010

11-
def protect_blueprint(self, blueprint: Blueprint):
11+
def protect_blueprint(self, blueprint):
12+
"""Add authentication middleware to all routes in a blueprint"""
1213
@blueprint.before_request
13-
def verify_request():
14-
if self.auth_config.auth_method == AuthMethod.NONE:
14+
@wraps(blueprint)
15+
def authenticate():
16+
if self.config.auth_method == AuthMethod.NONE:
1517
return None
1618

17-
auth_handlers = {
18-
AuthMethod.API_KEY: self._handle_api_key,
19-
AuthMethod.JWT: self._handle_jwt,
20-
AuthMethod.SESSION: self._handle_session
21-
}
19+
if self.config.auth_method == AuthMethod.API_KEY:
20+
return self._validate_api_key()
21+
elif self.config.auth_method == AuthMethod.JWT:
22+
return self._validate_jwt()
23+
elif self.config.auth_method == AuthMethod.SESSION:
24+
return self._validate_session()
2225

23-
handler = auth_handlers.get(self.auth_config.auth_method)
24-
if not handler:
25-
return jsonify({"error": "Invalid authentication method"}), 500
26-
27-
result = handler()
28-
if result is not True:
29-
return result
30-
31-
return None
32-
33-
def _handle_api_key(self):
26+
def _validate_api_key(self):
27+
"""Validate API key from request header"""
3428
api_key = request.headers.get('X-API-Key')
3529
if not api_key:
3630
return jsonify({"error": "API key is required"}), 401
37-
if api_key != self.auth_config.api_key:
31+
if api_key != self.config.api_key:
3832
return jsonify({"error": "Invalid API key"}), 401
39-
return True
33+
return None
4034

41-
def _handle_jwt(self):
35+
def _validate_jwt(self):
36+
"""Validate JWT from Authorization header"""
4237
auth_header = request.headers.get('Authorization')
4338
if not auth_header or not auth_header.startswith('Bearer '):
4439
return jsonify({"error": "JWT token is required"}), 401
45-
40+
4641
token = auth_header.split(' ')[1]
4742
if token in blacklisted_tokens:
4843
return jsonify({"error": "You have been logged out. Please log in again."}), 401
49-
44+
5045
try:
51-
jwt.decode(
52-
token,
53-
self.auth_config.jwt_secret,
54-
algorithms=["HS256"],
55-
options={"verify_exp": True}
56-
)
57-
return True
46+
jwt.decode(token, self.config.jwt_secret, algorithms=["HS256"])
47+
return None
5848
except jwt.ExpiredSignatureError:
5949
return jsonify({"error": "Token has expired"}), 401
6050
except jwt.InvalidTokenError as e:
6151
return jsonify({"error": f"Invalid JWT token: {str(e)}"}), 401
6252

63-
def _handle_session(self):
64-
if not session.get('authenticated'):
53+
def _validate_session(self):
54+
"""Validate session authentication"""
55+
if not session.get("authenticated"):
6556
return jsonify({"error": "Valid session required"}), 401
66-
return True
57+
58+
# Check if session has been invalidated
59+
current_session = request.cookies.get('session')
60+
if current_session and current_session in invalidated_sessions:
61+
session.clear()
62+
return jsonify({"error": "Session has been invalidated"}), 401
63+
64+
return None

‎app/services/auth_service.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import secrets
44
from config.auth_config import AuthConfig, AuthMethod
55
from models.user import User
6-
from flask import session, jsonify
6+
from flask import session, jsonify, request
77

88
# --- Configuration ---
99
auth_config = None # Global configuration object set during initialization
@@ -16,6 +16,7 @@ def init_auth_service(config: AuthConfig):
1616
users = [] # In-memory storage for user objects
1717
refresh_tokens = {} # Maps refresh tokens to usernames
1818
blacklisted_tokens = set() # Set of invalidated access tokens
19+
invalidated_sessions = set() # Set of invalidated session IDs
1920

2021
# --- User Management ---
2122
def is_username_taken(username):
@@ -112,9 +113,17 @@ def logout_jwt(access_token, refresh_token):
112113
return jsonify({"message": "Logout successful"})
113114

114115
def logout_session():
115-
"""Clear user session if authenticated"""
116+
"""Clear user session if authenticated and invalidate the session cookie"""
116117
if not session.get("authenticated"):
117118
return jsonify({"error": "Not authenticated"}), 401
118119

120+
response = jsonify({"message": "Logout successful"})
121+
122+
# Add current session ID to invalidated sessions set
123+
if request.cookies.get('session'):
124+
invalidated_sessions.add(request.cookies.get('session'))
125+
119126
session.clear()
120-
return jsonify({"message": "Logout successful"})
127+
# Set the session cookie to expire immediately
128+
response.set_cookie('session', '', expires=0)
129+
return response

‎test/auth_api_collection.json‎

Lines changed: 74 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -933,7 +933,40 @@
933933
}
934934
},
935935
{
936-
"name": "6. Logout with Session",
936+
"name": "6. Save Session Cookie Before Logout",
937+
"event": [
938+
{
939+
"listen": "test",
940+
"script": {
941+
"exec": [
942+
"// Save the current session cookie for later testing",
943+
"const sessionCookie = pm.cookies.get('session');",
944+
"if (sessionCookie) {",
945+
" pm.environment.set('old_session', sessionCookie);",
946+
" console.log('Saved session cookie:', sessionCookie);",
947+
"}",
948+
"",
949+
"pm.test('Status code is 200', function() {",
950+
" pm.response.to.have.status(200);",
951+
"});"
952+
],
953+
"type": "text/javascript"
954+
}
955+
}
956+
],
957+
"request": {
958+
"method": "GET",
959+
"url": {
960+
"raw": "http://localhost:8000/todos",
961+
"protocol": "http",
962+
"host": ["localhost"],
963+
"port": "8000",
964+
"path": ["todos"]
965+
}
966+
}
967+
},
968+
{
969+
"name": "7. Logout with Session",
937970
"event": [
938971
{
939972
"listen": "test",
@@ -959,7 +992,7 @@
959992
}
960993
},
961994
{
962-
"name": "7. Try to Get Todos after Logout (Should Fail)",
995+
"name": "8. Try to Get Todos with Normal Request after Logout (Should Fail)",
963996
"event": [
964997
{
965998
"listen": "test",
@@ -990,7 +1023,45 @@
9901023
}
9911024
},
9921025
{
993-
"name": "8. Get API Documentation with Session",
1026+
"name": "9. Try to Get Todos with Old Session Cookie (Should Fail)",
1027+
"event": [
1028+
{
1029+
"listen": "test",
1030+
"script": {
1031+
"exec": [
1032+
"pm.test('Status code is 401', function() {",
1033+
" pm.response.to.have.status(401);",
1034+
"});",
1035+
"",
1036+
"pm.test('Error message is correct', function() {",
1037+
" var jsonData = pm.response.json();",
1038+
" pm.expect(jsonData.error).to.equal('Session has been invalidated');",
1039+
"});"
1040+
],
1041+
"type": "text/javascript"
1042+
}
1043+
}
1044+
],
1045+
"request": {
1046+
"method": "GET",
1047+
"header": [
1048+
{
1049+
"key": "Cookie",
1050+
"value": "session={{old_session}}",
1051+
"type": "text"
1052+
}
1053+
],
1054+
"url": {
1055+
"raw": "http://localhost:8000/todos",
1056+
"protocol": "http",
1057+
"host": ["localhost"],
1058+
"port": "8000",
1059+
"path": ["todos"]
1060+
}
1061+
}
1062+
},
1063+
{
1064+
"name": "10. Get API Documentation with Session",
9941065
"event": [
9951066
{
9961067
"listen": "test",

0 commit comments

Comments
 (0)