Add test for fastapi-route-discovery

.github/workflows/end2end.yml

@@ -9,7 +9,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Install poetry
-        run: pip install poetry
+        run: pip install -c pip-constraints.txt poetry
       - name: Build
         run: make build
       - name: Upload artifacts
@@ -37,6 +37,7 @@ 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 }
         python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Install packages

.github/workflows/publish.yml

@@ -56,7 +56,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Install poetry
-      run: pip install poetry
+      run: pip install -c pip-constraints.txt poetry
     - name: Set the version for this release
       run: |
         TAG_NAME=${GITHUB_REF##*/}

.github/workflows/qa-tests.yml

@@ -30,7 +30,7 @@ jobs:
       - name: Setup safe-chain
         run: curl -fsSL https://github.com/AikidoSec/safe-chain/releases/latest/download/install-safe-chain.sh | sh -s -- --ci
       - name: Install poetry
-        run: pip install poetry
+        run: pip install -c firewall-python/pip-constraints.txt poetry
 
       - name: Build firewall-python dev package
         run: |

.github/workflows/test-publish.yml

@@ -57,7 +57,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Install poetry
-      run: pip install poetry
+      run: pip install -c pip-constraints.txt poetry
     - name: Set the version for this release
       run: |
         TAG_NAME=${GITHUB_REF##*/}

Makefile

@@ -19,7 +19,7 @@ lint:
 	poetry run black aikido_zen/
 	poetry run pylint aikido_zen/
 install: check_binaries
-	pip install poetry
+	pip install -c pip-constraints.txt poetry
 	poetry install
 .PHONY: dev_install
 dev_install: install

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.py

@@ -0,0 +1,14 @@
+from aikido_zen.sinks import on_import, patch_function
+from aikido_zen.sources.starlette.starlette_routing import _request_response
+
+
+@on_import("fastapi.routing", "fastapi")
+def patch(m):
+    """
+    FastAPI defines its own request_response in fastapi/routing.py rather than
+    importing it from starlette.routing. Patching only starlette.routing.request_response
+    leaves all FastAPI APIRoute endpoints unwrapped, so pre_response and post_response
+    hooks never fire for those routes.
+    (github src: https://github.com/tiangolo/fastapi/blob/master/fastapi/routing.py)
+    """
+    patch_function(m, "request_response", _request_response)

end2end/fastapi_postgres_uvicorn_test.py

@@ -0,0 +1,117 @@
+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
+
+    validate_heartbeat(
+        heartbeat_events[0],
+        routes=[
+            {
+                "apispec": {
+                    "auth": None,
+                    "body": {
+                        "schema": {
+                            "properties": {"dog_name": {"type": "string"}},
+                            "type": "object",
+                        },
+                        "type": "form-urlencoded",
+                    },
+                    "query": None,
+                },
+                "hits": 2,
+                "hits_delta_since_sync": 0,
+                "method": "POST",
+                "path": "/create",
+            },
+            {
+                "apispec": {},
+                "hits": 1,
+                "hits_delta_since_sync": 0,
+                "method": "GET",
+                "path": "/sync_route",
+            },
+        ],
+    )

pip-constraints.txt

@@ -0,0 +1,6 @@
+charset-normalizer==3.3.2
+more-itertools==10.8.0
+certifi==2024.8.30
+urllib3==2.2.3
+idna==3.8
+requests==2.32.3