Skip to content

Commit d66b0ff

Browse files
committed
Add new sample app for latest version of fastapi
1 parent b56d7c5 commit d66b0ff

9 files changed

Lines changed: 1318 additions & 0 deletions

File tree

.github/workflows/end2end.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ jobs:
3838
- { name: quart-postgres-uvicorn, testfile: end2end/quart_postgres_uvicorn_test.py }
3939
- { name: starlette-postgres-uvicorn, testfile: end2end/starlette_postgres_uvicorn_test.py }
4040
- { name: fastapi-postgres-uvicorn, testfile: end2end/fastapi_postgres_uvicorn_test.py }
41+
- { name: fastapi-postgres-uvicorn-latest, testfile: end2end/fastapi_postgres_uvicorn_latest_test.py }
4142
python-version: ["3.10", "3.11", "3.12", "3.13"]
4243
steps:
4344
- name: Install packages
Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
import time
2+
import pytest
3+
import requests
4+
from .server.check_events_from_mock import (
5+
fetch_events_from_mock,
6+
validate_started_event,
7+
validate_heartbeat,
8+
filter_on_event_type,
9+
)
10+
11+
post_url_fw = "http://localhost:8116/create"
12+
post_url_nofw = "http://localhost:8117/create"
13+
sync_route_fw = "http://localhost:8116/sync_route"
14+
sync_route_nofw = "http://localhost:8117/sync_route"
15+
16+
17+
def test_firewall_started_okay():
18+
events = fetch_events_from_mock("http://localhost:5000")
19+
started_events = filter_on_event_type(events, "started")
20+
assert len(started_events) == 1
21+
validate_started_event(started_events[0], None)
22+
23+
24+
def test_safe_response_with_firewall():
25+
res = requests.post(post_url_fw, data={"dog_name": "Bobby Tables"})
26+
assert res.status_code == 201
27+
28+
29+
def test_safe_response_without_firewall():
30+
res = requests.post(post_url_nofw, data={"dog_name": "Bobby Tables"})
31+
assert res.status_code == 201
32+
33+
34+
def test_dangerous_response_with_firewall():
35+
dog_name = "Dangerous Bobby', TRUE); -- "
36+
res = requests.post(post_url_fw, data={"dog_name": dog_name})
37+
assert res.status_code == 500
38+
39+
time.sleep(5) # Wait for attack to be reported
40+
events = fetch_events_from_mock("http://localhost:5000")
41+
attacks = filter_on_event_type(events, "detected_attack")
42+
43+
assert len(attacks) == 1
44+
del attacks[0]["attack"]["stack"]
45+
assert attacks[0]["attack"]["blocked"] == True
46+
assert attacks[0]["attack"]["kind"] == "sql_injection"
47+
assert attacks[0]["attack"]["metadata"]["sql"] == "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Dangerous Bobby', TRUE); -- ', FALSE)"
48+
assert attacks[0]["attack"]["metadata"]["dialect"] == "postgres"
49+
assert attacks[0]["attack"]["operation"] == "asyncpg.connection.Connection.execute"
50+
assert attacks[0]["attack"]["pathToPayload"] == ".dog_name"
51+
assert attacks[0]["attack"]["payload"] == "\"Dangerous Bobby', TRUE); -- \""
52+
assert attacks[0]["attack"]["source"] == "body"
53+
assert attacks[0]["attack"]["user"]["id"] == "user123"
54+
assert attacks[0]["attack"]["user"]["name"] == "John Doe"
55+
56+
# These assertions verify the pre/post response hooks fired for FastAPI APIRoute endpoints.
57+
# Without patching fastapi.routing.request_response, route will be None/empty and
58+
# source will not reflect the fastapi framework context.
59+
assert attacks[0]["request"]["route"] == "/create"
60+
assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
61+
62+
63+
def test_dangerous_response_without_firewall():
64+
dog_name = "Dangerous Bobby', TRUE); -- "
65+
res = requests.post(post_url_nofw, data={"dog_name": dog_name})
66+
assert res.status_code == 201
67+
68+
69+
def test_sync_route_with_firewall():
70+
res = requests.get(sync_route_fw)
71+
assert res.status_code == 200
72+
73+
74+
def test_sync_route_without_firewall():
75+
res = requests.get(sync_route_nofw)
76+
assert res.status_code == 200
77+
78+
79+
def test_routes_discovered_in_heartbeat():
80+
# This test verifies that FastAPI APIRoute endpoints are discovered via route discovery.
81+
# Without patching fastapi.routing.request_response, the heartbeat will report
82+
# current_routes: {} even with active traffic, because the post_response hook never fires.
83+
time.sleep(55) # Wait for first heartbeat (fires ~60s after start)
84+
85+
events = fetch_events_from_mock("http://localhost:5000")
86+
heartbeat_events = filter_on_event_type(events, "heartbeat")
87+
assert len(heartbeat_events) >= 1
88+
89+
routes = heartbeat_events[0]["routes"]
90+
route_map = {(r["method"], r["path"]): r for r in routes}
91+
92+
assert ("POST", "/create") in route_map
93+
create_route = route_map[("POST", "/create")]
94+
assert create_route["apispec"] == {
95+
"auth": None,
96+
"body": {
97+
"schema": {
98+
"properties": {"dog_name": {"type": "string"}},
99+
"type": "object",
100+
},
101+
"type": "form-urlencoded",
102+
},
103+
"query": None,
104+
}
105+
assert create_route["hits"] == 1
106+
107+
assert ("GET", "/sync_route") in route_map
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
include ../common.mk
2+
3+
PORT = 8116
4+
PORT_DISABLED = 8117
5+
6+
.PHONY: run
7+
run: install
8+
@echo "Running sample app fastapi-postgres-uvicorn-latest with Zen on port $(PORT)"
9+
$(AIKIDO_ENV_COMMON) \
10+
poetry run uvicorn app:app --host 0.0.0.0 --port $(PORT) --workers 4
11+
12+
.PHONY: runZenDisabled
13+
runZenDisabled: install
14+
@echo "Running sample app fastapi-postgres-uvicorn-latest without Zen on port $(PORT_DISABLED)"
15+
$(AIKIDO_ENV_DISABLED) \
16+
poetry run uvicorn app:app --host 0.0.0.0 --port $(PORT_DISABLED) --workers 4
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import aikido_zen # Aikido package import
2+
aikido_zen.protect()
3+
4+
import time
5+
import asyncpg
6+
from fastapi import FastAPI, Request, HTTPException
7+
from fastapi.responses import HTMLResponse, JSONResponse
8+
from fastapi.templating import Jinja2Templates
9+
from fastapi.middleware.cors import CORSMiddleware
10+
from aikido_zen.middleware import AikidoFastAPIMiddleware
11+
12+
templates = Jinja2Templates(directory="templates")
13+
14+
app = FastAPI()
15+
16+
# CORS middleware (optional, depending on your needs)
17+
app.add_middleware(
18+
CORSMiddleware,
19+
allow_origins=["*"], # Adjust this as needed
20+
allow_credentials=True,
21+
allow_methods=["*"],
22+
allow_headers=["*"],
23+
)
24+
@app.middleware("http")
25+
async def set_user_middleware(request: Request, call_next):
26+
aikido_zen.set_user({"id": "user123", "name": "John Doe"})
27+
return await call_next(request)
28+
29+
app.add_middleware(AikidoFastAPIMiddleware)
30+
31+
async def get_db_connection():
32+
return await asyncpg.connect(
33+
host="localhost",
34+
database="db",
35+
user="user",
36+
password="password"
37+
)
38+
39+
@app.get("/", response_class=HTMLResponse)
40+
async def homepage(request: Request):
41+
conn = await get_db_connection()
42+
dogs = await conn.fetch("SELECT * FROM dogs")
43+
await conn.close()
44+
return templates.TemplateResponse('index.html', {"request": request, "title": 'Homepage', "dogs": dogs})
45+
46+
@app.get("/dogpage/{dog_id:int}", response_class=HTMLResponse)
47+
async def get_dogpage(request: Request, dog_id: int):
48+
conn = await get_db_connection()
49+
dog = await conn.fetchrow("SELECT * FROM dogs WHERE id = $1", dog_id)
50+
await conn.close()
51+
if dog is None:
52+
raise HTTPException(status_code=404, detail="Dog not found")
53+
return templates.TemplateResponse('dogpage.html', {"request": request, "title": 'Dog', "dog": dog, "isAdmin": "Yes" if dog[2] else "No"})
54+
55+
@app.get("/create", response_class=HTMLResponse)
56+
async def show_create_dog_form(request: Request):
57+
return templates.TemplateResponse('create_dog.html', {"request": request})
58+
59+
@app.post("/create")
60+
async def create_dog(request: Request):
61+
data = await request.form()
62+
dog_name = data.get('dog_name')
63+
64+
if not dog_name:
65+
return JSONResponse({"error": "dog_name is required"}, status_code=400)
66+
67+
conn = await get_db_connection()
68+
try:
69+
await conn.execute(f"INSERT INTO dogs (dog_name, isAdmin) VALUES ('%s', FALSE)" % (dog_name))
70+
finally:
71+
await conn.close()
72+
73+
return JSONResponse({"message": f'Dog {dog_name} created successfully'}, status_code=201)
74+
75+
@app.get("/just")
76+
async def just():
77+
return JSONResponse({"message": "Empty Page"})
78+
79+
@app.get("/test_ratelimiting_1")
80+
async def just():
81+
return JSONResponse({"message": "Empty Page"})
82+
83+
@app.get("/delayed_route")
84+
async def delayed_route():
85+
time.sleep(1/1000) # Note: This will block the event loop; consider using asyncio.sleep instead
86+
return JSONResponse({"message": "Empty Page"})
87+
88+
@app.get("/sync_route")
89+
def sync_route():
90+
data = {"message": "This is a non-async route!"}
91+
return JSONResponse(data)

0 commit comments

Comments
 (0)