Coverage for src / local_deep_research / security / rate_limiter.py: 82%

1""" 

2Rate limiting for API endpoints using Flask-Limiter. 

3 

4Provides decorators and configuration for rate limiting HTTP requests 

5to prevent abuse and resource exhaustion. 

6 

7Configuration: 

8 RATE_LIMIT_FAIL_CLOSED: Set to "true" in production to fail closed 

9 when rate limiter is not initialized. Default is fail-open for 

10 easier development setup. 

11 

12 For production, also configure: 

13 storage_uri="redis://localhost:6379" for multi-worker support 

14""" 

15 

16import os 

17 

18from flask import g 

19from flask_limiter import Limiter 

20from flask_limiter.util import get_remote_address 

21from loguru import logger 

22 

23 

24# Global limiter instance (initialized by app factory) 

25_limiter = None 

26 

27# Configuration: fail-closed in production, fail-open in development 

28RATE_LIMIT_FAIL_CLOSED = ( 

29 os.environ.get("RATE_LIMIT_FAIL_CLOSED", "false").lower() == "true" 

30) 

31 

32 

33def get_rate_limiter() -> Limiter: 

34 """ 

35 Get the global rate limiter instance. 

36 

37 Returns: 

38 Flask-Limiter instance 

39 

40 Raises: 

41 RuntimeError: If limiter hasn't been initialized 

42 """ 

43 global _limiter 

44 if _limiter is None: 

45 raise RuntimeError( 

46 "Rate limiter not initialized. Call init_rate_limiter(app) first." 

47 ) 

48 return _limiter 

49 

50 

51def init_rate_limiter(app): 

52 """ 

53 Initialize the rate limiter with the Flask app. 

54 

55 This should be called once during app initialization. 

56 

57 Args: 

58 app: Flask application instance 

59 

60 Returns: 

61 Configured Limiter instance 

62 """ 

63 global _limiter 

64 

65 # Use authenticated username if available, otherwise fall back to IP 

66 def get_user_identifier(): 

67 # Check if user is authenticated 

68 username = g.get("current_user") 

69 if username: 

70 return f"user:{username}" 

71 # Fall back to IP address for unauthenticated requests 

72 return f"ip:{get_remote_address()}" 

73 

74 _limiter = Limiter( 

75 app=app, 

76 key_func=get_user_identifier, 

77 default_limits=[], # No default limits - apply via decorators only 

78 storage_uri="memory://", # Use in-memory storage (for development) 

79 # For production, consider: storage_uri="redis://localhost:6379" 

80 strategy="fixed-window", 

81 headers_enabled=True, # Add rate limit headers to responses 

82 ) 

83 

84 logger.info("Rate limiter initialized successfully") 

85 return _limiter 

86 

87 

88def upload_rate_limit(f): 

89 """ 

90 Decorator for rate limiting file upload endpoints. 

91 

92 Limits: 

93 - 10 uploads per minute per user 

94 - 100 uploads per hour per user 

95 

96 Usage: 

97 @research_bp.route("/api/upload/pdf", methods=["POST"]) 

98 @login_required 

99 @upload_rate_limit 

100 def upload_pdf(): 

101 ... 

102 

103 Returns: 

104 Decorated function with rate limiting 

105 

106 Note: 

107 If the rate limiter is not initialized (e.g., flask_limiter not installed), 

108 behavior depends on RATE_LIMIT_FAIL_CLOSED environment variable: 

109 - "false" (default): Pass through without rate limiting (fail-open) 

110 - "true": Raise RuntimeError, crashing the application (fail-closed) 

111 

112 Set RATE_LIMIT_FAIL_CLOSED=true in production to ensure rate limiting 

113 is always active. 

114 """ 

115 try: 

116 limiter = get_rate_limiter() 

117 # Use Flask-Limiter's limit decorator 

118 return limiter.limit("10 per minute;100 per hour")(f) 

119 except RuntimeError: 

120 # Rate limiter not initialized - this is expected if flask_limiter 

121 # is not installed or init_rate_limiter was not called 

122 if RATE_LIMIT_FAIL_CLOSED: 

123 logger.exception( 

124 f"Rate limiter not initialized for {f.__name__} and " 

125 "RATE_LIMIT_FAIL_CLOSED is enabled. Application will fail." 

126 ) 

127 raise 

128 logger.warning( 

129 f"Rate limiting disabled for {f.__name__}: " 

130 "limiter not initialized. Install flask_limiter for rate limiting. " 

131 "Set RATE_LIMIT_FAIL_CLOSED=true to enforce rate limiting in production." 

132 ) 

133 return f 

134 except Exception: 

135 # Unexpected error - log it but don't break the application 

136 if RATE_LIMIT_FAIL_CLOSED: 

137 logger.exception( 

138 f"Unexpected error applying rate limit to {f.__name__} and " 

139 "RATE_LIMIT_FAIL_CLOSED is enabled. Application will fail." 

140 ) 

141 raise 

142 logger.exception( 

143 f"Unexpected error applying rate limit to {f.__name__}. " 

144 "Rate limiting disabled for this endpoint. " 

145 "Set RATE_LIMIT_FAIL_CLOSED=true to enforce rate limiting in production." 

146 ) 

147 return f 

148 

149 

150# Export convenience decorator 

151__all__ = ["init_rate_limiter", "get_rate_limiter", "upload_rate_limit"]