Coverage for mcpgateway / utils / security_cookies.py: 100%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/utils/security_cookies.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Mihai Criveti 

6 

7Security Cookie Utilities for MCP Gateway. 

8 

9This module provides utilities for setting secure authentication cookies with proper 

10security attributes to prevent common cookie-based attacks. 

11""" 

12 

13# Standard 

14import logging 

15 

16# Third-Party 

17from fastapi import Response 

18 

19# First-Party 

20from mcpgateway.config import settings 

21 

22logger = logging.getLogger(__name__) 

23 

24# RFC 6265 §6.1: browsers SHOULD support cookies of at least 4096 bytes 

25_COOKIE_HARD_LIMIT = 4096 

26_COOKIE_WARN_THRESHOLD = 3800 

27 

28 

29class CookieTooLargeError(Exception): 

30 """Raised when a cookie value exceeds the browser's 4KB limit.""" 

31 

32 def __init__(self, cookie_size: int, limit: int = _COOKIE_HARD_LIMIT): 

33 """Initialize with the actual cookie size and the browser limit. 

34 

35 Args: 

36 cookie_size: Actual size of the cookie in bytes. 

37 limit: Maximum allowed cookie size in bytes. 

38 """ 

39 self.cookie_size = cookie_size 

40 self.limit = limit 

41 super().__init__(f"Cookie size {cookie_size} bytes exceeds browser limit of {limit} bytes") 

42 

43 

44def set_auth_cookie(response: Response, token: str, remember_me: bool = False) -> None: 

45 """ 

46 Set authentication cookie with security flags and size validation. 

47 

48 Configures the JWT token as a secure HTTP-only cookie with appropriate 

49 security attributes to prevent XSS and CSRF attacks. 

50 

51 Args: 

52 response: FastAPI response object to set the cookie on 

53 token: JWT token to store in the cookie 

54 remember_me: If True, sets longer expiration time (30 days vs 1 hour) 

55 

56 Raises: 

57 CookieTooLargeError: If the cookie would exceed 4096 bytes 

58 

59 Security attributes set: 

60 - httponly: Prevents JavaScript access to the cookie 

61 - secure: HTTPS only in production environments 

62 - samesite: CSRF protection (configurable, defaults to 'lax') 

63 - path: Cookie scope limitation 

64 - max_age: Automatic expiration 

65 

66 Examples: 

67 Basic cookie set with remember_me disabled: 

68 >>> from fastapi import Response 

69 >>> from mcpgateway.utils.security_cookies import set_auth_cookie 

70 >>> resp = Response() 

71 >>> set_auth_cookie(resp, 'tok123', remember_me=False) 

72 >>> header = resp.headers.get('set-cookie') 

73 >>> 'jwt_token=' in header and 'HttpOnly' in header and 'Path=/' in header 

74 True 

75 

76 Extended expiration when remember_me is True: 

77 >>> resp2 = Response() 

78 >>> set_auth_cookie(resp2, 'tok123', remember_me=True) 

79 >>> 'Max-Age=2592000' in resp2.headers.get('set-cookie') # 30 days 

80 True 

81 """ 

82 # Set expiration based on remember_me preference 

83 max_age = 30 * 24 * 3600 if remember_me else 3600 # 30 days or 1 hour 

84 

85 # Determine if we should use secure flag 

86 # In production or when explicitly configured, require HTTPS 

87 use_secure = (settings.environment == "production") or settings.secure_cookies 

88 samesite = settings.cookie_samesite 

89 path = settings.app_root_path or "/" 

90 

91 # Estimate cookie size in bytes (Set-Cookie header format) 

92 # The cookie name, value, and attributes all count toward the limit 

93 cookie_header = f"jwt_token={token}; HttpOnly; SameSite={samesite}; Path={path}; Max-Age={max_age}" 

94 if use_secure: 

95 cookie_header += "; Secure" 

96 cookie_size = len(cookie_header.encode("ascii", errors="replace")) 

97 

98 # Extract sub claim for log context (best-effort, don't build control flow on it) 

99 sub_for_log = "" 

100 try: 

101 # Standard 

102 import base64 

103 import json 

104 

105 parts = token.split(".") 

106 if len(parts) >= 2: 

107 padded = parts[1] + "=" * (-len(parts[1]) % 4) 

108 payload = json.loads(base64.urlsafe_b64decode(padded)) 

109 sub_for_log = payload.get("sub", "") 

110 except Exception: # nosec B110 - Best-effort sub extraction for logging 

111 pass 

112 

113 if cookie_size > _COOKIE_HARD_LIMIT: 

114 logger.error("JWT cookie size %d bytes exceeds %d byte browser limit (user: %s)", cookie_size, _COOKIE_HARD_LIMIT, sub_for_log) 

115 raise CookieTooLargeError(cookie_size) 

116 

117 if cookie_size > _COOKIE_WARN_THRESHOLD: 

118 logger.warning("JWT cookie size %d bytes approaching %d byte browser limit (user: %s)", cookie_size, _COOKIE_HARD_LIMIT, sub_for_log) 

119 

120 response.set_cookie( 

121 key="jwt_token", 

122 value=token, 

123 max_age=max_age, 

124 httponly=True, # Prevents JavaScript access 

125 secure=use_secure, # HTTPS only in production 

126 samesite=samesite, # CSRF protection 

127 path=path, # Cookie scope 

128 ) 

129 

130 

131def clear_auth_cookie(response: Response) -> None: 

132 """ 

133 Clear authentication cookie securely. 

134 

135 Removes the JWT token cookie by setting it to expire immediately 

136 with the same security attributes used when setting it. 

137 

138 Args: 

139 response: FastAPI response object to clear the cookie from 

140 

141 Examples: 

142 >>> from fastapi import Response 

143 >>> resp = Response() 

144 >>> set_auth_cookie(resp, 'tok123') 

145 >>> clear_auth_cookie(resp) 

146 >>> # Deletion sets another Set-Cookie for jwt_token; presence indicates cleared cookie header 

147 >>> 'jwt_token=' in resp.headers.get('set-cookie') 

148 True 

149 """ 

150 # Use same security settings as when setting the cookie 

151 use_secure = (settings.environment == "production") or settings.secure_cookies 

152 

153 response.delete_cookie( 

154 key="jwt_token", 

155 path=settings.app_root_path or "/", 

156 secure=use_secure, 

157 httponly=True, 

158 samesite=settings.cookie_samesite, 

159 ) 

160 

161 

162def set_session_cookie(response: Response, session_id: str, max_age: int = 3600) -> None: 

163 """ 

164 Set session cookie with security flags. 

165 

166 Configures a session ID cookie with appropriate security attributes. 

167 

168 Args: 

169 response: FastAPI response object to set the cookie on 

170 session_id: Session identifier to store in the cookie 

171 max_age: Cookie expiration time in seconds (default: 1 hour) 

172 

173 Examples: 

174 >>> from fastapi import Response 

175 >>> resp = Response() 

176 >>> set_session_cookie(resp, 'sess-1', max_age=3600) 

177 >>> header = resp.headers.get('set-cookie') 

178 >>> 'session_id=sess-1' in header and 'HttpOnly' in header 

179 True 

180 """ 

181 use_secure = (settings.environment == "production") or settings.secure_cookies 

182 

183 response.set_cookie( 

184 key="session_id", 

185 value=session_id, 

186 max_age=max_age, 

187 httponly=True, 

188 secure=use_secure, 

189 samesite=settings.cookie_samesite, 

190 path=settings.app_root_path or "/", 

191 ) 

192 

193 

194def clear_session_cookie(response: Response) -> None: 

195 """ 

196 Clear session cookie securely. 

197 

198 Args: 

199 response: FastAPI response object to clear the cookie from 

200 

201 Examples: 

202 >>> from fastapi import Response 

203 >>> resp = Response() 

204 >>> set_session_cookie(resp, 'sess-2', max_age=60) 

205 >>> clear_session_cookie(resp) 

206 >>> 'session_id=' in resp.headers.get('set-cookie') 

207 True 

208 """ 

209 use_secure = (settings.environment == "production") or settings.secure_cookies 

210 

211 response.delete_cookie( 

212 key="session_id", 

213 path=settings.app_root_path or "/", 

214 secure=use_secure, 

215 httponly=True, 

216 samesite=settings.cookie_samesite, 

217 )