Backup-bdg/OpenHands

1import asyncio2import os3from collections import defaultdict4from datetime import datetime, timedelta5from urllib.parse import urlparse6 7from fastapi import Request8from fastapi.middleware.cors import CORSMiddleware9from fastapi.responses import JSONResponse10from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint11from starlette.requests import Request as StarletteRequest12from starlette.responses import Response13from starlette.types import ASGIApp14 15 16class LocalhostCORSMiddleware(CORSMiddleware):17    """18    Custom CORS middleware that allows any request from localhost/127.0.0.1 domains,19    while using standard CORS rules for other origins.20    """21 22    def __init__(self, app: ASGIApp) -> None:23        allow_origins_str = os.getenv('PERMITTED_CORS_ORIGINS')24        if allow_origins_str:25            allow_origins = tuple(26                origin.strip() for origin in allow_origins_str.split(',')27            )28        else:29            allow_origins = ()30        super().__init__(31            app,32            allow_origins=allow_origins,33            allow_credentials=True,34            allow_methods=['*'],35            allow_headers=['*'],36        )37 38    def is_allowed_origin(self, origin: str) -> bool:39        if origin and not self.allow_origins and not self.allow_origin_regex:40            parsed = urlparse(origin)41            hostname = parsed.hostname or ''42 43            # Allow any localhost/127.0.0.1 origin regardless of port44            if hostname in ['localhost', '127.0.0.1']:45                return True46 47        # For missing origin or other origins, use the parent class's logic48        result: bool = super().is_allowed_origin(origin)49        return result50 51 52class CacheControlMiddleware(BaseHTTPMiddleware):53    """54    Middleware to disable caching for all routes by adding appropriate headers55    """56 57    async def dispatch(58        self, request: Request, call_next: RequestResponseEndpoint59    ) -> Response:60        response = await call_next(request)61        if request.url.path.startswith('/assets'):62            # The content of the assets directory has fingerprinted file names so we cache aggressively63            response.headers['Cache-Control'] = 'public, max-age=2592000, immutable'64        else:65            response.headers['Cache-Control'] = (66                'no-cache, no-store, must-revalidate, max-age=0'67            )68            response.headers['Pragma'] = 'no-cache'69            response.headers['Expires'] = '0'70        return response71 72 73class InMemoryRateLimiter:74    history: dict[str, list[datetime]]75    requests: int76    seconds: int77    sleep_seconds: int78 79    def __init__(self, requests: int = 2, seconds: int = 1, sleep_seconds: int = 1):80        self.requests = requests81        self.seconds = seconds82        self.sleep_seconds = sleep_seconds83        self.history = defaultdict(list)84        self.sleep_seconds = sleep_seconds85 86    def _clean_old_requests(self, key: str) -> None:87        now = datetime.now()88        cutoff = now - timedelta(seconds=self.seconds)89        self.history[key] = [ts for ts in self.history[key] if ts > cutoff]90 91    async def __call__(self, request: Request) -> bool:92        key = request.client.host93        now = datetime.now()94 95        self._clean_old_requests(key)96 97        self.history[key].append(now)98 99        if len(self.history[key]) > self.requests * 2:100            return False101        elif len(self.history[key]) > self.requests:102            if self.sleep_seconds > 0:103                await asyncio.sleep(self.sleep_seconds)104                return True105            else:106                return False107 108        return True109 110 111class RateLimitMiddleware(BaseHTTPMiddleware):112    def __init__(self, app: ASGIApp, rate_limiter: InMemoryRateLimiter):113        super().__init__(app)114        self.rate_limiter = rate_limiter115 116    async def dispatch(117        self, request: Request, call_next: RequestResponseEndpoint118    ) -> Response:119        if not self.is_rate_limited_request(request):120            return await call_next(request)121        ok = await self.rate_limiter(request)122        if not ok:123            return JSONResponse(124                status_code=429,125                content={'message': 'Too many requests'},126                headers={'Retry-After': '1'},127            )128        return await call_next(request)129 130    def is_rate_limited_request(self, request: StarletteRequest) -> bool:131        if request.url.path.startswith('/assets'):132            return False133        # Put Other non rate limited checks here134        return True135