rate limit

Author: pedro-cf

docker-compose.yml

@@ -22,6 +22,7 @@ services:
       - ES_USE_SSL=false
       - ES_VERIFY_CERTS=false
       - BACKEND=elasticsearch
+      - STAC_FASTAPI_RATE_LIMIT=2/5second
     ports:
       - "8080:8080"
     volumes:
@@ -54,6 +55,7 @@ services:
       - ES_USE_SSL=false
       - ES_VERIFY_CERTS=false
       - BACKEND=opensearch
+      - STAC_FASTAPI_RATE_LIMIT=200/minute
     ports:
       - "8082:8082"
     volumes:

stac_fastapi/core/setup.py

@@ -19,6 +19,7 @@
     "pygeofilter==0.2.1",
     "typing_extensions==4.8.0",
     "jsonschema",
+    "slowapi==0.1.9",
 ]
 
 setup(

stac_fastapi/core/stac_fastapi/core/rate_limit.py

@@ -0,0 +1,27 @@
+import os
+from fastapi import FastAPI, Request
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from slowapi.middleware import SlowAPIMiddleware
+import logging
+
+logger = logging.getLogger(__name__)
+
+limiter = Limiter(key_func=get_remote_address)
+
+def setup_rate_limit(app: FastAPI):
+    RATE_LIMIT = os.getenv("STAC_FASTAPI_RATE_LIMIT")
+    logger.info(f"Setting up rate limit with RATE_LIMIT={RATE_LIMIT}")
+    if RATE_LIMIT:
+        app.state.limiter = limiter
+        app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+        app.add_middleware(SlowAPIMiddleware)
+
+        @app.middleware("http")
+        @limiter.limit(RATE_LIMIT)
+        async def rate_limit_middleware(request: Request, call_next):
+            response = await call_next(request)
+            return response
+
+    logger.info("Rate limit setup complete")

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/app.py

@@ -1,6 +1,7 @@
 """FastAPI application."""
 
 import os
+from stac_fastapi.core.rate_limit import setup_rate_limit
 
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
@@ -97,6 +98,8 @@
 app = api.app
 app.root_path = os.getenv("STAC_FASTAPI_ROOT_PATH", "")
 
+# Add rate limit
+setup_rate_limit(app)
 
 @app.on_event("startup")
 async def _startup_event() -> None:

stac_fastapi/opensearch/stac_fastapi/opensearch/app.py

@@ -1,6 +1,7 @@
 """FastAPI application."""
 
 import os
+from stac_fastapi.core.rate_limit import setup_rate_limit
 
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
@@ -97,6 +98,8 @@
 app = api.app
 app.root_path = os.getenv("STAC_FASTAPI_ROOT_PATH", "")
 
+# Add rate limit
+setup_rate_limit(app)
 
 @app.on_event("startup")
 async def _startup_event() -> None:

stac_fastapi/tests/basic_auth/test_basic_auth.py

@@ -18,7 +18,7 @@ async def test_get_search_not_authenticated(app_client_basic_auth, ctx):
 
 @pytest.mark.asyncio
 async def test_post_search_authenticated(app_client_basic_auth, ctx):
-    """Test protected endpoint [POST /search] with reader auhtentication"""
+    """Test protected endpoint [POST /search] with reader authentication"""
     if not os.getenv("BASIC_AUTH"):
         pytest.skip()
     params = {"id": ctx.item["id"]}
@@ -34,7 +34,7 @@ async def test_post_search_authenticated(app_client_basic_auth, ctx):
 async def test_delete_resource_anonymous(
     app_client_basic_auth,
 ):
-    """Test protected endpoint [DELETE /collections/{collection_id}] without auhtentication"""
+    """Test protected endpoint [DELETE /collections/{collection_id}] without authentication"""
     if not os.getenv("BASIC_AUTH"):
         pytest.skip()
 

stac_fastapi/tests/conftest.py

@@ -25,6 +25,7 @@
     EsAsyncAggregationClient,
 )
 from stac_fastapi.core.route_dependencies import get_route_dependencies
+from stac_fastapi.core.rate_limit import setup_rate_limit
 
 if os.getenv("BACKEND", "elasticsearch").lower() == "opensearch":
     from stac_fastapi.opensearch.config import AsyncOpensearchSettings as AsyncSettings
@@ -223,7 +224,56 @@ async def app():
 
     post_request_model = create_post_request_model(search_extensions)
 
-    return StacApi(
+    app = StacApi(
+        settings=settings,
+        client=CoreClient(
+            database=database,
+            session=None,
+            extensions=extensions,
+            post_request_model=post_request_model,
+        ),
+        extensions=extensions,
+        search_get_request_model=create_get_request_model(search_extensions),
+        search_post_request_model=post_request_model,
+    ).app
+
+    return app
+
+
+@pytest_asyncio.fixture(scope="function")
+async def app_rate_limit(monkeypatch):
+    monkeypatch.setenv("STAC_FASTAPI_RATE_LIMIT", "2/minute")
+    
+    settings = AsyncSettings()
+    
+    aggregation_extension = AggregationExtension(
+        client=EsAsyncAggregationClient(
+            database=database, session=None, settings=settings
+        )
+    )
+    aggregation_extension.POST = EsAggregationExtensionPostRequest
+    aggregation_extension.GET = EsAggregationExtensionGetRequest
+
+    search_extensions = [
+        TransactionExtension(
+            client=TransactionsClient(
+                database=database, session=None, settings=settings
+            ),
+            settings=settings,
+        ),
+        SortExtension(),
+        FieldsExtension(),
+        QueryExtension(),
+        TokenPaginationExtension(),
+        FilterExtension(),
+        FreeTextExtension(),
+    ]
+
+    extensions = [aggregation_extension] + search_extensions
+
+    post_request_model = create_post_request_model(search_extensions)
+
+    app = StacApi(
         settings=settings,
         client=CoreClient(
             database=database,
@@ -236,6 +286,12 @@ async def app():
         search_post_request_model=post_request_model,
     ).app
 
+    # Set up rate limit
+    setup_rate_limit(app)
+
+    return app
+
+
 
 @pytest_asyncio.fixture(scope="session")
 async def app_client(app):
@@ -246,6 +302,15 @@ async def app_client(app):
         yield c
 
 
+@pytest_asyncio.fixture(scope="function")
+async def app_client_rate_limit(app_rate_limit):
+    await create_index_templates()
+    await create_collection_index()
+
+    async with AsyncClient(app=app_rate_limit, base_url="http://test-server") as c:
+        yield c
+
+
 @pytest_asyncio.fixture(scope="session")
 async def app_basic_auth():
 

stac_fastapi/tests/rate_limit/test_rate_limit.py

@@ -0,0 +1,20 @@
+import pytest
+from httpx import AsyncClient
+import logging
+from slowapi.errors import RateLimitExceeded
+
+logger = logging.getLogger(__name__)
+
+@pytest.mark.asyncio
+async def test_rate_limit(app_client_rate_limit: AsyncClient):
+    expected_status_codes = [200, 200, 429, 429, 429]
+    
+    for i, expected_status_code in enumerate(expected_status_codes):
+        try:
+            response = await app_client_rate_limit.get("/collections")
+            status_code = response.status_code
+        except RateLimitExceeded:
+            status_code = 429
+        
+        logger.info(f"Request {i+1}: Status code {status_code}")
+        assert status_code == expected_status_code, f"Expected status code {expected_status_code}, but got {status_code}"