Fix fastapi support for 0.118+

.github/workflows/end2end.yml

@@ -37,6 +37,8 @@ jobs:
           - { name: flask-postgres-xml, testfile: end2end/flask_postgres_xml_lxml_test.py }
           - { name: quart-postgres-uvicorn, testfile: end2end/quart_postgres_uvicorn_test.py }
           - { name: starlette-postgres-uvicorn, testfile: end2end/starlette_postgres_uvicorn_test.py }
+          - { name: fastapi-postgres-uvicorn, testfile: end2end/fastapi_postgres_uvicorn_test.py }
+          - { name: fastapi-postgres-uvicorn-latest, testfile: end2end/fastapi_postgres_uvicorn_latest_test.py }
         python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Install packages

aikido_zen/__init__.py

@@ -61,6 +61,7 @@ def protect(mode="daemon", token=""):
     import aikido_zen.sources.flask
     import aikido_zen.sources.quart
     import aikido_zen.sources.starlette
+    import aikido_zen.sources.fastapi
     import aikido_zen.sources.xml_sources.xml
     import aikido_zen.sources.xml_sources.lxml
 

aikido_zen/sources/fastapi/__init__.py

@@ -0,0 +1 @@
+import aikido_zen.sources.fastapi.fastapi_routing

aikido_zen/sources/fastapi/fastapi_routing.py

@@ -0,0 +1,18 @@
+from ...sinks import on_import, patch_function
+from ..starlette.starlette_routing import _request_response
+
+
+@on_import("fastapi.routing", "fastapi")
+def patch(m):
+    """
+    patching module fastapi.routing
+    - patches: request_response
+
+    Newer FastAPI defines its own request_response rather than importing from
+    starlette.routing, so the starlette patch alone doesn't cover APIRoute endpoints.
+    In older versions FastAPI imports from starlette, so we skip to avoid double-wrapping.
+    """
+    import starlette.routing
+
+    if m.request_response is not starlette.routing.request_response:
+        patch_function(m, "request_response", _request_response)

end2end/fastapi_postgres_uvicorn_latest_test.py

@@ -0,0 +1,107 @@
+import time
+import pytest
+import requests
+from .server.check_events_from_mock import (
+    fetch_events_from_mock,
+    validate_started_event,
+    validate_heartbeat,
+    filter_on_event_type,
+)
+
+post_url_fw = "http://localhost:8116/create"
+post_url_nofw = "http://localhost:8117/create"
+sync_route_fw = "http://localhost:8116/sync_route"
+sync_route_nofw = "http://localhost:8117/sync_route"
+
+
+def test_firewall_started_okay():
+    events = fetch_events_from_mock("http://localhost:5000")
+    started_events = filter_on_event_type(events, "started")
+    assert len(started_events) == 1
+    validate_started_event(started_events[0], None)
+
+
+def test_safe_response_with_firewall():
+    res = requests.post(post_url_fw, data={"dog_name": "Bobby Tables"})
+    assert res.status_code == 201
+
+
+def test_safe_response_without_firewall():
+    res = requests.post(post_url_nofw, data={"dog_name": "Bobby Tables"})
+    assert res.status_code == 201
+
+
+def test_dangerous_response_with_firewall():
+    dog_name = "Dangerous Bobby', TRUE); -- "
+    res = requests.post(post_url_fw, data={"dog_name": dog_name})
+    assert res.status_code == 500
+
+    time.sleep(5)  # Wait for attack to be reported
+    events = fetch_events_from_mock("http://localhost:5000")
+    attacks = filter_on_event_type(events, "detected_attack")
+
+    assert len(attacks) == 1
+    del attacks[0]["attack"]["stack"]
+    assert attacks[0]["attack"]["blocked"] == True
+    assert attacks[0]["attack"]["kind"] == "sql_injection"
+    assert attacks[0]["attack"]["metadata"]["sql"] == "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Dangerous Bobby', TRUE); -- ', FALSE)"
+    assert attacks[0]["attack"]["metadata"]["dialect"] == "postgres"
+    assert attacks[0]["attack"]["operation"] == "asyncpg.connection.Connection.execute"
+    assert attacks[0]["attack"]["pathToPayload"] == ".dog_name"
+    assert attacks[0]["attack"]["payload"] == "\"Dangerous Bobby', TRUE); -- \""
+    assert attacks[0]["attack"]["source"] == "body"
+    assert attacks[0]["attack"]["user"]["id"] == "user123"
+    assert attacks[0]["attack"]["user"]["name"] == "John Doe"
+
+    # These assertions verify the pre/post response hooks fired for FastAPI APIRoute endpoints.
+    # Without patching fastapi.routing.request_response, route will be None/empty and
+    # source will not reflect the fastapi framework context.
+    assert attacks[0]["request"]["route"] == "/create"
+    assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
+
+
+def test_dangerous_response_without_firewall():
+    dog_name = "Dangerous Bobby', TRUE); -- "
+    res = requests.post(post_url_nofw, data={"dog_name": dog_name})
+    assert res.status_code == 201
+
+
+def test_sync_route_with_firewall():
+    res = requests.get(sync_route_fw)
+    assert res.status_code == 200
+
+
+def test_sync_route_without_firewall():
+    res = requests.get(sync_route_nofw)
+    assert res.status_code == 200
+
+
+def test_routes_discovered_in_heartbeat():
+    # This test verifies that FastAPI APIRoute endpoints are discovered via route discovery.
+    # Without patching fastapi.routing.request_response, the heartbeat will report
+    # current_routes: {} even with active traffic, because the post_response hook never fires.
+    time.sleep(55)  # Wait for first heartbeat (fires ~60s after start)
+
+    events = fetch_events_from_mock("http://localhost:5000")
+    heartbeat_events = filter_on_event_type(events, "heartbeat")
+    assert len(heartbeat_events) >= 1
+
+    routes = heartbeat_events[0]["routes"]
+    route_map = {(r["method"], r["path"]): r for r in routes}
+
+    assert ("POST", "/create") in route_map
+    create_route = route_map[("POST", "/create")]
+    assert create_route["apispec"] == {
+        "auth": None,
+        "body": {
+            "schema": {
+                "properties": {"dog_name": {"type": "string"}},
+                "type": "object",
+            },
+            "type": "form-urlencoded",
+        },
+        "query": None,
+    }
+    assert create_route["hits"] == 1
+
+    assert ("GET", "/sync_route") in route_map

end2end/fastapi_postgres_uvicorn_test.py

@@ -0,0 +1,107 @@
+import time
+import pytest
+import requests
+from .server.check_events_from_mock import (
+    fetch_events_from_mock,
+    validate_started_event,
+    validate_heartbeat,
+    filter_on_event_type,
+)
+
+post_url_fw = "http://localhost:8112/create"
+post_url_nofw = "http://localhost:8113/create"
+sync_route_fw = "http://localhost:8112/sync_route"
+sync_route_nofw = "http://localhost:8113/sync_route"
+
+
+def test_firewall_started_okay():
+    events = fetch_events_from_mock("http://localhost:5000")
+    started_events = filter_on_event_type(events, "started")
+    assert len(started_events) == 1
+    validate_started_event(started_events[0], None)
+
+
+def test_safe_response_with_firewall():
+    res = requests.post(post_url_fw, data={"dog_name": "Bobby Tables"})
+    assert res.status_code == 201
+
+
+def test_safe_response_without_firewall():
+    res = requests.post(post_url_nofw, data={"dog_name": "Bobby Tables"})
+    assert res.status_code == 201
+
+
+def test_dangerous_response_with_firewall():
+    dog_name = "Dangerous Bobby', TRUE); -- "
+    res = requests.post(post_url_fw, data={"dog_name": dog_name})
+    assert res.status_code == 500
+
+    time.sleep(5)  # Wait for attack to be reported
+    events = fetch_events_from_mock("http://localhost:5000")
+    attacks = filter_on_event_type(events, "detected_attack")
+
+    assert len(attacks) == 1
+    del attacks[0]["attack"]["stack"]
+    assert attacks[0]["attack"]["blocked"] == True
+    assert attacks[0]["attack"]["kind"] == "sql_injection"
+    assert attacks[0]["attack"]["metadata"]["sql"] == "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Dangerous Bobby', TRUE); -- ', FALSE)"
+    assert attacks[0]["attack"]["metadata"]["dialect"] == "postgres"
+    assert attacks[0]["attack"]["operation"] == "asyncpg.connection.Connection.execute"
+    assert attacks[0]["attack"]["pathToPayload"] == ".dog_name"
+    assert attacks[0]["attack"]["payload"] == "\"Dangerous Bobby', TRUE); -- \""
+    assert attacks[0]["attack"]["source"] == "body"
+    assert attacks[0]["attack"]["user"]["id"] == "user123"
+    assert attacks[0]["attack"]["user"]["name"] == "John Doe"
+
+    # These assertions verify the pre/post response hooks fired for FastAPI APIRoute endpoints.
+    # Without patching fastapi.routing.request_response, route will be None/empty and
+    # source will not reflect the fastapi framework context.
+    assert attacks[0]["request"]["route"] == "/create"
+    assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
+
+
+def test_dangerous_response_without_firewall():
+    dog_name = "Dangerous Bobby', TRUE); -- "
+    res = requests.post(post_url_nofw, data={"dog_name": dog_name})
+    assert res.status_code == 201
+
+
+def test_sync_route_with_firewall():
+    res = requests.get(sync_route_fw)
+    assert res.status_code == 200
+
+
+def test_sync_route_without_firewall():
+    res = requests.get(sync_route_nofw)
+    assert res.status_code == 200
+
+
+def test_routes_discovered_in_heartbeat():
+    # This test verifies that FastAPI APIRoute endpoints are discovered via route discovery.
+    # Without patching fastapi.routing.request_response, the heartbeat will report
+    # current_routes: {} even with active traffic, because the post_response hook never fires.
+    time.sleep(55)  # Wait for first heartbeat (fires ~60s after start)
+
+    events = fetch_events_from_mock("http://localhost:5000")
+    heartbeat_events = filter_on_event_type(events, "heartbeat")
+    assert len(heartbeat_events) >= 1
+
+    routes = heartbeat_events[0]["routes"]
+    route_map = {(r["method"], r["path"]): r for r in routes}
+
+    assert ("POST", "/create") in route_map
+    create_route = route_map[("POST", "/create")]
+    assert create_route["apispec"] == {
+        "auth": None,
+        "body": {
+            "schema": {
+                "properties": {"dog_name": {"type": "string"}},
+                "type": "object",
+            },
+            "type": "form-urlencoded",
+        },
+        "query": None,
+    }
+    assert create_route["hits"] == 1
+
+    assert ("GET", "/sync_route") in route_map

sample-apps/fastapi-postgres-uvicorn-latest/Makefile

@@ -0,0 +1,16 @@
+include ../common.mk
+
+PORT = 8116
+PORT_DISABLED = 8117
+
+.PHONY: run
+run: install
+	@echo "Running sample app fastapi-postgres-uvicorn-latest with Zen on port $(PORT)"
+	$(AIKIDO_ENV_COMMON) \
+	poetry run uvicorn app:app --host 0.0.0.0 --port $(PORT) --workers 4
+
+.PHONY: runZenDisabled
+runZenDisabled: install
+	@echo "Running sample app fastapi-postgres-uvicorn-latest without Zen on port $(PORT_DISABLED)"
+	$(AIKIDO_ENV_DISABLED) \
+	poetry run uvicorn app:app --host 0.0.0.0 --port $(PORT_DISABLED) --workers 4

sample-apps/fastapi-postgres-uvicorn-latest/app.py

@@ -0,0 +1,91 @@
+import aikido_zen  # Aikido package import
+aikido_zen.protect()
+
+import time
+import asyncpg
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.templating import Jinja2Templates
+from fastapi.middleware.cors import CORSMiddleware
+from aikido_zen.middleware import AikidoFastAPIMiddleware
+
+templates = Jinja2Templates(directory="templates")
+
+app = FastAPI()
+
+# CORS middleware (optional, depending on your needs)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Adjust this as needed
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+@app.middleware("http")
+async def set_user_middleware(request: Request, call_next):
+    aikido_zen.set_user({"id": "user123", "name": "John Doe"})
+    return await call_next(request)
+
+app.add_middleware(AikidoFastAPIMiddleware)
+
+async def get_db_connection():
+    return await asyncpg.connect(
+        host="localhost",
+        database="db",
+        user="user",
+        password="password"
+    )
+
+@app.get("/", response_class=HTMLResponse)
+async def homepage(request: Request):
+    conn = await get_db_connection()
+    dogs = await conn.fetch("SELECT * FROM dogs")
+    await conn.close()
+    return templates.TemplateResponse('index.html', {"request": request, "title": 'Homepage', "dogs": dogs})
+
+@app.get("/dogpage/{dog_id:int}", response_class=HTMLResponse)
+async def get_dogpage(request: Request, dog_id: int):
+    conn = await get_db_connection()
+    dog = await conn.fetchrow("SELECT * FROM dogs WHERE id = $1", dog_id)
+    await conn.close()
+    if dog is None:
+        raise HTTPException(status_code=404, detail="Dog not found")
+    return templates.TemplateResponse('dogpage.html', {"request": request, "title": 'Dog', "dog": dog, "isAdmin": "Yes" if dog[2] else "No"})
+
+@app.get("/create", response_class=HTMLResponse)
+async def show_create_dog_form(request: Request):
+    return templates.TemplateResponse('create_dog.html', {"request": request})
+
+@app.post("/create")
+async def create_dog(request: Request):
+    data = await request.form()
+    dog_name = data.get('dog_name')
+
+    if not dog_name:
+        return JSONResponse({"error": "dog_name is required"}, status_code=400)
+
+    conn = await get_db_connection()
+    try:
+        await conn.execute(f"INSERT INTO dogs (dog_name, isAdmin) VALUES ('%s', FALSE)" % (dog_name))
+    finally:
+        await conn.close()
+
+    return JSONResponse({"message": f'Dog {dog_name} created successfully'}, status_code=201)
+
+@app.get("/just")
+async def just():
+    return JSONResponse({"message": "Empty Page"})
+
+@app.get("/test_ratelimiting_1")
+async def just():
+    return JSONResponse({"message": "Empty Page"})
+
+@app.get("/delayed_route")
+async def delayed_route():
+    time.sleep(1/1000)  # Note: This will block the event loop; consider using asyncio.sleep instead
+    return JSONResponse({"message": "Empty Page"})
+
+@app.get("/sync_route")
+def sync_route():
+    data = {"message": "This is a non-async route!"}
+    return JSONResponse(data)