crackerjack 0.33.0__py3-none-any.whl → 0.33.2__py3-none-any.whl

crackerjack/services/websocket_resource_limiter.py

@@ -1,11 +1,3 @@
-"""
-WebSocket Resource Limiter to prevent resource exhaustion attacks.
-
-Provides comprehensive resource monitoring and limiting for WebSocket
-connections, including connection limits, message limits, memory usage
-tracking, and DoS prevention.
-"""
-
 import asyncio
 import time
 import typing as t
@@ -18,8 +10,6 @@ from .security_logger import SecurityEventLevel, SecurityEventType, get_security
 
 @dataclass
 class ConnectionMetrics:
-    """Metrics for individual WebSocket connections."""
-
     client_id: str
     connect_time: float = field(default_factory=time.time)
     message_count: int = 0
@@ -29,70 +19,44 @@ class ConnectionMetrics:
 
     @property
     def connection_duration(self) -> float:
-        """Get connection duration in seconds."""
         return time.time() - self.connect_time
 
     @property
     def idle_time(self) -> float:
-        """Get idle time since last activity."""
         return time.time() - self.last_activity
 
 
 @dataclass
 class ResourceLimits:
-    """Resource limits configuration."""
-
     max_connections: int = 50
     max_connections_per_ip: int = 10
-    max_message_size: int = 64 * 1024  # 64KB
+    max_message_size: int = 64 * 1024
     max_messages_per_minute: int = 100
     max_messages_per_connection: int = 10000
-    max_connection_duration: float = 3600.0  # 1 hour
-    max_idle_time: float = 300.0  # 5 minutes
+    max_connection_duration: float = 3600.0
+    max_idle_time: float = 300.0
     max_memory_usage_mb: int = 100
     connection_timeout: float = 30.0
     message_timeout: float = 10.0
 
 
 class ResourceExhaustedError(Exception):
-    """Raised when resource limits are exceeded."""
-
     pass
 
 
 class WebSocketResourceLimiter:
-    """
-    Resource limiter for WebSocket connections.
-
-    Features:
-    - Connection count and per-IP limits
-    - Message size and rate limiting
-    - Memory usage monitoring
-    - Connection duration limits
-    - Idle connection cleanup
-    - DoS attack prevention
-    """
-
     def __init__(self, limits: ResourceLimits | None = None):
-        """
-        Initialize WebSocket resource limiter.
-
-        Args:
-            limits: Resource limits configuration
-        """
         self.limits = limits or ResourceLimits()
         self.security_logger = get_security_logger()
 
         self._setup_limiter_components()
 
     def _setup_limiter_components(self) -> None:
-        """Set up all limiter components in initialization."""
         self._initialize_connection_tracking()
         self._initialize_metrics_tracking()
         self._initialize_cleanup_system()
 
     def _initialize_connection_tracking(self) -> None:
-        """Initialize thread-safe connection tracking structures."""
         self._lock = RLock()
         self._connections: dict[str, ConnectionMetrics] = {}
         self._ip_connections: dict[str, set[str]] = defaultdict(set)
@@ -102,20 +66,16 @@ class WebSocketResourceLimiter:
         self._blocked_ips: dict[str, float] = {}
 
     def _initialize_metrics_tracking(self) -> None:
-        """Initialize metrics tracking with proper typing."""
         self._connection_metrics: dict[str, ConnectionMetrics] = {}
         self._message_queues: dict[str, deque[bytes]] = defaultdict(deque)
         self._resource_usage: dict[str, dict[str, t.Any]] = {}
         self._memory_usage: int = 0
 
     def _initialize_cleanup_system(self) -> None:
-        """Initialize the background cleanup system."""
         self._cleanup_task: asyncio.Task[t.Any] | None = None
         self._shutdown_event = asyncio.Event()
 
     async def start(self) -> None:
-        """Start the resource limiter with background cleanup."""
-
         if self._cleanup_task is None:
             self._cleanup_task = asyncio.create_task(self._cleanup_loop())
 
@@ -127,8 +87,6 @@ class WebSocketResourceLimiter:
             )
 
     async def stop(self) -> None:
-        """Stop the resource limiter and cleanup resources."""
-
         self._shutdown_event.set()
 
         if self._cleanup_task:
@@ -143,7 +101,6 @@ class WebSocketResourceLimiter:
 
             self._cleanup_task = None
 
-        # Force cleanup all connections
         with self._lock:
             self._connections.clear()
             self._ip_connections.clear()
@@ -157,17 +114,6 @@ class WebSocketResourceLimiter:
         )
 
     def validate_new_connection(self, client_id: str, client_ip: str) -> None:
-        """
-        Validate if a new connection can be accepted.
-
-        Args:
-            client_id: Unique client identifier
-            client_ip: Client IP address
-
-        Raises:
-            ResourceExhaustedError: If connection limits exceeded
-        """
-
         with self._lock:
             current_time = time.time()
 
@@ -176,7 +122,6 @@ class WebSocketResourceLimiter:
             self._check_per_ip_connection_limit(client_id, client_ip, current_time)
 
     def _check_ip_blocking_status(self, client_ip: str, current_time: float) -> None:
-        """Check if IP is currently blocked."""
         if client_ip in self._blocked_ips:
             if current_time < self._blocked_ips[client_ip]:
                 raise ResourceExhaustedError(f"IP {client_ip} is temporarily blocked")
@@ -184,7 +129,6 @@ class WebSocketResourceLimiter:
                 del self._blocked_ips[client_ip]
 
     def _check_total_connection_limit(self, client_id: str, client_ip: str) -> None:
-        """Check if total connection limit is exceeded."""
         if len(self._connections) >= self.limits.max_connections:
             self.security_logger.log_security_event(
                 event_type=SecurityEventType.RATE_LIMIT_EXCEEDED,
@@ -201,7 +145,6 @@ class WebSocketResourceLimiter:
     def _check_per_ip_connection_limit(
         self, client_id: str, client_ip: str, current_time: float
     ) -> None:
-        """Check if per-IP connection limit is exceeded."""
         ip_connection_count = len(self._ip_connections[client_ip])
         if ip_connection_count >= self.limits.max_connections_per_ip:
             self.security_logger.log_security_event(
@@ -213,22 +156,13 @@ class WebSocketResourceLimiter:
                 additional_data={"client_ip": client_ip},
             )
 
-            # Block IP temporarily for repeated violations
-            self._blocked_ips[client_ip] = current_time + 300.0  # 5 minute block
+            self._blocked_ips[client_ip] = current_time + 300.0
 
             raise ResourceExhaustedError(
                 f"Maximum connections per IP exceeded: {ip_connection_count}"
             )
 
     def register_connection(self, client_id: str, client_ip: str) -> None:
-        """
-        Register a new WebSocket connection.
-
-        Args:
-            client_id: Unique client identifier
-            client_ip: Client IP address
-        """
-
         with self._lock:
             metrics = ConnectionMetrics(client_id=client_id)
             self._connections[client_id] = metrics
@@ -244,14 +178,6 @@ class WebSocketResourceLimiter:
         )
 
     def unregister_connection(self, client_id: str, client_ip: str) -> None:
-        """
-        Unregister a WebSocket connection.
-
-        Args:
-            client_id: Unique client identifier
-            client_ip: Client IP address
-        """
-
         with self._lock:
             if client_id in self._connections:
                 metrics = self._connections[client_id]
@@ -272,28 +198,15 @@ class WebSocketResourceLimiter:
                     },
                 )
 
-            # Remove from IP tracking
             if client_ip in self._ip_connections:
                 self._ip_connections[client_ip].discard(client_id)
                 if not self._ip_connections[client_ip]:
                     del self._ip_connections[client_ip]
 
-            # Clear message history
             if client_id in self._message_history:
                 del self._message_history[client_id]
 
     def validate_message(self, client_id: str, message_size: int) -> None:
-        """
-        Validate if a message can be processed.
-
-        Args:
-            client_id: Client identifier
-            message_size: Size of the message in bytes
-
-        Raises:
-            ResourceExhaustedError: If message limits exceeded
-        """
-
         with self._lock:
             self._check_message_size(client_id, message_size)
             metrics = self._get_connection_metrics(client_id)
@@ -301,7 +214,6 @@ class WebSocketResourceLimiter:
             self._check_message_rate(client_id)
 
     def _check_message_size(self, client_id: str, message_size: int) -> None:
-        """Check if message size exceeds limits."""
         if message_size > self.limits.max_message_size:
             self.security_logger.log_security_event(
                 event_type=SecurityEventType.RATE_LIMIT_EXCEEDED,
@@ -313,13 +225,11 @@ class WebSocketResourceLimiter:
             raise ResourceExhaustedError(f"Message too large: {message_size} bytes")
 
     def _get_connection_metrics(self, client_id: str) -> ConnectionMetrics:
-        """Get connection metrics, raising error if connection doesn't exist."""
         if client_id not in self._connections:
             raise ResourceExhaustedError(f"Connection not registered: {client_id}")
         return self._connections[client_id]
 
     def _check_message_count(self, client_id: str, metrics: ConnectionMetrics) -> None:
-        """Check if total message count per connection exceeds limits."""
         if metrics.message_count >= self.limits.max_messages_per_connection:
             self.security_logger.log_security_event(
                 event_type=SecurityEventType.RATE_LIMIT_EXCEEDED,
@@ -333,11 +243,9 @@ class WebSocketResourceLimiter:
             )
 
     def _check_message_rate(self, client_id: str) -> None:
-        """Check if message rate exceeds per-minute limits."""
         current_time = time.time()
         message_times = self._message_history[client_id]
 
-        # Remove old messages (older than 1 minute)
         while message_times and current_time - message_times[0] > 60.0:
             message_times.popleft()
 
@@ -356,15 +264,6 @@ class WebSocketResourceLimiter:
     def track_message(
         self, client_id: str, message_size: int, is_sent: bool = True
     ) -> None:
-        """
-        Track a processed message.
-
-        Args:
-            client_id: Client identifier
-            message_size: Size of the message in bytes
-            is_sent: True if message was sent, False if received
-        """
-
         with self._lock:
             current_time = time.time()
 
@@ -378,20 +277,9 @@ class WebSocketResourceLimiter:
                 else:
                     metrics.bytes_received += message_size
 
-            # Track message timing for rate limiting
             self._message_history[client_id].append(current_time)
 
     async def check_connection_limits(self, client_id: str) -> None:
-        """
-        Check if connection has exceeded limits.
-
-        Args:
-            client_id: Client identifier
-
-        Raises:
-            ResourceExhaustedError: If connection limits exceeded
-        """
-
         with self._lock:
             if client_id not in self._connections:
                 return
@@ -399,45 +287,40 @@ class WebSocketResourceLimiter:
             metrics = self._connections[client_id]
             time.time()
 
-            # Check connection duration
             if metrics.connection_duration > self.limits.max_connection_duration:
                 self.security_logger.log_security_event(
                     event_type=SecurityEventType.CONNECTION_TIMEOUT,
                     level=SecurityEventLevel.WARNING,
-                    message=f"Connection duration exceeded: {metrics.connection_duration:.1f}s",
+                    message=f"Connection duration exceeded: {metrics.connection_duration: .1f}s",
                     client_id=client_id,
                     operation="connection_limit_check",
                 )
                 raise ResourceExhaustedError(
-                    f"Connection duration exceeded: {metrics.connection_duration:.1f}s"
+                    f"Connection duration exceeded: {metrics.connection_duration: .1f}s"
                 )
 
-            # Check idle time
             if metrics.idle_time > self.limits.max_idle_time:
                 self.security_logger.log_security_event(
                     event_type=SecurityEventType.CONNECTION_IDLE,
                     level=SecurityEventLevel.INFO,
-                    message=f"Connection idle timeout: {metrics.idle_time:.1f}s",
+                    message=f"Connection idle timeout: {metrics.idle_time: .1f}s",
                     client_id=client_id,
                     operation="connection_limit_check",
                 )
                 raise ResourceExhaustedError(
-                    f"Connection idle timeout: {metrics.idle_time:.1f}s"
+                    f"Connection idle timeout: {metrics.idle_time: .1f}s"
                 )
 
     async def _cleanup_loop(self) -> None:
-        """Background cleanup loop for expired connections and resources."""
-
         while not self._shutdown_event.is_set():
             try:
                 await asyncio.wait_for(
                     self._shutdown_event.wait(),
-                    timeout=30.0,  # Cleanup every 30 seconds
+                    timeout=30.0,
                 )
-                break  # Shutdown requested
+                break
 
             except TimeoutError:
-                # Perform cleanup
                 await self._perform_cleanup()
 
         self.security_logger.log_security_event(
@@ -448,7 +331,6 @@ class WebSocketResourceLimiter:
         )
 
     async def _perform_cleanup(self) -> None:
-        """Perform resource cleanup."""
         current_time = time.time()
 
         with self._lock:
@@ -459,7 +341,6 @@ class WebSocketResourceLimiter:
         self._log_cleanup_results(cleanup_count)
 
     def _cleanup_expired_connections(self, current_time: float) -> int:
-        """Clean up expired connections and return count of cleaned connections."""
         expired_connections = self._find_expired_connections()
         cleanup_count = 0
 
@@ -470,7 +351,6 @@ class WebSocketResourceLimiter:
         return cleanup_count
 
     def _find_expired_connections(self) -> list[str]:
-        """Find connections that have exceeded duration or idle time limits."""
         return [
             client_id
             for client_id, metrics in self._connections.items()
@@ -481,31 +361,25 @@ class WebSocketResourceLimiter:
         ]
 
     def _remove_expired_connection(self, client_id: str) -> bool:
-        """Remove a specific expired connection and its associated data."""
         if client_id not in self._connections:
             return False
 
-        # Remove connection metrics
         del self._connections[client_id]
 
-        # Clean up IP tracking
         for client_set in self._ip_connections.values():
             client_set.discard(client_id)
 
-        # Clean up message history
         if client_id in self._message_history:
             del self._message_history[client_id]
 
         return True
 
     def _cleanup_empty_ip_entries(self) -> None:
-        """Remove IP entries that no longer have any connections."""
         empty_ips = [ip for ip, clients in self._ip_connections.items() if not clients]
         for ip in empty_ips:
             del self._ip_connections[ip]
 
     def _cleanup_expired_ip_blocks(self, current_time: float) -> None:
-        """Remove IP blocks that have expired."""
         expired_blocks = [
             ip
             for ip, block_until in self._blocked_ips.items()
@@ -515,7 +389,6 @@ class WebSocketResourceLimiter:
             del self._blocked_ips[ip]
 
     def _log_cleanup_results(self, cleanup_count: int) -> None:
-        """Log cleanup results if any connections were cleaned up."""
         if cleanup_count > 0:
             self.security_logger.log_security_event(
                 event_type=SecurityEventType.RESOURCE_CLEANUP,
@@ -525,8 +398,6 @@ class WebSocketResourceLimiter:
             )
 
     def get_resource_status(self) -> dict[str, t.Any]:
-        """Get current resource usage status."""
-
         with self._lock:
             return {
                 "connections": {
@@ -551,21 +422,16 @@ class WebSocketResourceLimiter:
             }
 
     def get_connection_metrics(self, client_id: str) -> ConnectionMetrics | None:
-        """Get metrics for a specific connection."""
-
         with self._lock:
             return self._connections.get(client_id)
 
 
-# Global singleton instance
 _resource_limiter: WebSocketResourceLimiter | None = None
 
 
 def get_websocket_resource_limiter(
     limits: ResourceLimits | None = None,
 ) -> WebSocketResourceLimiter:
-    """Get the global WebSocket resource limiter instance."""
-
     global _resource_limiter
     if _resource_limiter is None:
         _resource_limiter = WebSocketResourceLimiter(limits)