Skip to content

Commit 5014770

Browse files
RasmusWLRasmusWL
committed
Python: FastAPI: Model response classes
Figuring out how to do the `media_type` tracking was quite difficult.
1 parent eef946a commit 5014770

2 files changed

Lines changed: 116 additions & 17 deletions

File tree

‎python/ql/lib/semmle/python/frameworks/FastApi.qll‎

Lines changed: 107 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -79,14 +79,80 @@ private module FastApi {
7979
// Response modeling
8080
// ---------------------------------------------------------------------------
8181
/**
82-
* Provides models for the `fastapi.Response` class
82+
* Provides models for the `fastapi.Response` class and subclasses.
83+
*
84+
* See https://fastapi.tiangolo.com/advanced/custom-response/#response.
8385
*/
8486
module Response {
85-
/** Gets a reference to the `fastapi.Response` class. */
86-
private API::Node classRef() { result = API::moduleImport("fastapi").getMember("Response") }
87+
/**
88+
* Gets the `API::Node` for the manually modeled response classes called `name`.
89+
*/
90+
private API::Node getModeledResponseClass(string name) {
91+
name = "Response" and
92+
result = API::moduleImport("fastapi").getMember("Response")
93+
or
94+
// see https://github.com/tiangolo/fastapi/blob/master/fastapi/responses.py
95+
name in [
96+
"Response", "HTMLResponse", "PlainTextResponse", "JSONResponse", "UJSONResponse",
97+
"ORJSONResponse", "RedirectResponse", "StreamingResponse", "FileResponse"
98+
] and
99+
result = API::moduleImport("fastapi").getMember("responses").getMember(name)
100+
}
101+
102+
/**
103+
* Gets the default MIME type for a FastAPI response class (defined with the
104+
* `media_type` class-attribute).
105+
*
106+
* Also models user-defined classes and tries to take inheritance into account.
107+
*
108+
* TODO: build easy way to solve problems like this, like we used to have the
109+
* `ClassValue.lookup` predicate.
110+
*/
111+
private string getDefaultMimeType(API::Node responseClass) {
112+
exists(string name | responseClass = getModeledResponseClass(name) |
113+
// no defaults for these.
114+
name in ["Response", "RedirectResponse", "StreamingResponse"] and
115+
none()
116+
or
117+
// For `FileResponse` the code will guess what mimetype
118+
// to use, or fall back to "text/plain", but claiming that all responses will
119+
// have "text/plain" per default is also highly inaccurate, so just going to not
120+
// do anything about this.
121+
name = "FileResponse" and
122+
none()
123+
or
124+
name = "HTMLResponse" and
125+
result = "text/html"
126+
or
127+
name = "PlainTextResponse" and
128+
result = "text/plain"
129+
or
130+
name in ["JSONResponse", "UJSONResponse", "ORJSONResponse"] and
131+
result = "application/json"
132+
)
133+
or
134+
// user-defined subclasses
135+
exists(Class cls, API::Node base |
136+
cls.getABase() = base.getAUse().asExpr() and
137+
// since we _know_ that the base has an API Node, that means there is a subclass
138+
// edge leading to the `ClassExpr` for the class.
139+
responseClass.getAnImmediateUse().asExpr().(ClassExpr) = cls.getParent()
140+
|
141+
exists(Assign assign | assign = cls.getAStmt() |
142+
assign.getATarget().(Name).getId() = "media_type" and
143+
result = assign.getValue().(StrConst).getText()
144+
)
145+
or
146+
// TODO: this should use a proper MRO calculation instead
147+
not exists(Assign assign | assign = cls.getAStmt() |
148+
assign.getATarget().(Name).getId() = "media_type"
149+
) and
150+
result = getDefaultMimeType(base)
151+
)
152+
}
87153

88154
/**
89-
* A source of instances of `fastapi.Response`, extend this class to model new instances.
155+
* A source of instances of `fastapi.Response` and its' subclasses, extend this class to model new instances.
90156
*
91157
* This can include instantiations of the class, return values from function
92158
* calls, or a special parameter that will be set when functions are called by an external
@@ -96,9 +162,41 @@ private module FastApi {
96162
*/
97163
abstract class InstanceSource extends DataFlow::LocalSourceNode { }
98164

99-
/** A direct instantiation of `fastapi.Response`. */
100-
private class ClassInstantiation extends InstanceSource, DataFlow::CallCfgNode {
101-
ClassInstantiation() { this = classRef().getACall() }
165+
/** A direct instantiation of a response class. */
166+
private class ResponseInstantiation extends InstanceSource, HTTP::Server::HttpResponse::Range,
167+
DataFlow::CallCfgNode {
168+
API::Node baseApiNode;
169+
API::Node responseClass;
170+
171+
ResponseInstantiation() {
172+
baseApiNode = getModeledResponseClass(_) and
173+
responseClass = baseApiNode.getASubclass*() and
174+
this = responseClass.getACall()
175+
}
176+
177+
override DataFlow::Node getBody() {
178+
not baseApiNode = getModeledResponseClass(["RedirectResponse", "FileResponse"]) and
179+
result in [this.getArg(0), this.getArgByName("content")]
180+
}
181+
182+
override DataFlow::Node getMimetypeOrContentTypeArg() {
183+
not baseApiNode = getModeledResponseClass("RedirectResponse") and
184+
result in [this.getArg(3), this.getArgByName("media_type")]
185+
}
186+
187+
override string getMimetypeDefault() { result = getDefaultMimeType(responseClass) }
188+
}
189+
190+
/**
191+
* A direct instantiation of a redirect response.
192+
*/
193+
private class RedirectResponseInstantiation extends ResponseInstantiation,
194+
HTTP::Server::HttpRedirectResponse::Range {
195+
RedirectResponseInstantiation() { baseApiNode = getModeledResponseClass("RedirectResponse") }
196+
197+
override DataFlow::Node getRedirectLocation() {
198+
result in [this.getArg(0), this.getArgByName("url")]
199+
}
102200
}
103201

104202
/**
@@ -109,7 +207,8 @@ private module FastApi {
109207
*/
110208
class RequestHandlerParam extends InstanceSource, DataFlow::ParameterNode {
111209
RequestHandlerParam() {
112-
this.getParameter().getAnnotation() = classRef().getAUse().asExpr() and
210+
this.getParameter().getAnnotation() =
211+
getModeledResponseClass(_).getASubclass*().getAUse().asExpr() and
113212
any(FastApiRouteSetup rs).getARequestHandler().getArgByName(_) = this.getParameter()
114213
}
115214
}

‎python/ql/test/library-tests/frameworks/fastapi/response_test.py‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -37,14 +37,14 @@ class MyXmlResponse(fastapi.responses.Response):
3737

3838

3939
@app.get("/response_parameter_custom_type", response_class=MyXmlResponse) # $ routeSetup="/response_parameter_custom_type"
40-
async def response_parameter_custom_type(response: MyXmlResponse): # $ requestHandler SPURIOUS: routedParameter=response
40+
async def response_parameter_custom_type(response: MyXmlResponse): # $ requestHandler
4141
# NOTE: This is a contrived example of using a wrong annotation for the response
4242
# parameter. It will be passed a `fastapi.responses.Response` value when handling an
4343
# incoming request, so NOT a `MyXmlResponse` value. Cookies/Headers are still
4444
# propagated to the final response though.
4545
print(type(response))
4646
assert type(response) == fastapi.responses.Response
47-
response.set_cookie("key", "value") # $ MISSING: CookieWrite CookieName="key" CookieValue="value"
47+
response.set_cookie("key", "value") # $ CookieWrite CookieName="key" CookieValue="value"
4848
response.headers["Custom-Response-Type"] = "yes, but only after function has run"
4949
xml_data = "<foo>FOO</foo>"
5050
return xml_data # $ HttpResponse responseBody=xml_data SPURIOUS: mimetype=application/json MISSING: mimetype=application/xml
@@ -60,8 +60,8 @@ async def response_parameter_custom_type(response: MyXmlResponse): # $ requestHa
6060
@app.get("/direct_response") # $ routeSetup="/direct_response"
6161
async def direct_response(): # $ requestHandler
6262
xml_data = "<foo>FOO</foo>"
63-
resp = fastapi.responses.Response(xml_data, 200, None, "application/xml") # $ MISSING: HttpResponse mimetype=application/xml responseBody=xml_data
64-
resp = fastapi.responses.Response(content=xml_data, media_type="application/xml") # $ MISSING: HttpResponse mimetype=application/xml responseBody=xml_data
63+
resp = fastapi.responses.Response(xml_data, 200, None, "application/xml") # $ HttpResponse mimetype=application/xml responseBody=xml_data
64+
resp = fastapi.responses.Response(content=xml_data, media_type="application/xml") # $ HttpResponse mimetype=application/xml responseBody=xml_data
6565
return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
6666

6767

@@ -74,7 +74,7 @@ async def direct_response2(): # $ requestHandler
7474
@app.get("/my_xml_response") # $ routeSetup="/my_xml_response"
7575
async def my_xml_response(): # $ requestHandler
7676
xml_data = "<foo>FOO</foo>"
77-
resp = MyXmlResponse(content=xml_data) # $ MISSING: HttpResponse mimetype=application/xml responseBody=xml_data
77+
resp = MyXmlResponse(content=xml_data) # $ HttpResponse mimetype=application/xml responseBody=xml_data
7878
return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
7979

8080

@@ -87,7 +87,7 @@ async def my_xml_response2(): # $ requestHandler
8787
@app.get("/html_response") # $ routeSetup="/html_response"
8888
async def html_response(): # $ requestHandler
8989
hello_world = "<h1>Hello World!</h1>"
90-
resp = fastapi.responses.HTMLResponse(hello_world) # $ MISSING: HttpResponse mimetype=text/html responseBody=hello_world
90+
resp = fastapi.responses.HTMLResponse(hello_world) # $ HttpResponse mimetype=text/html responseBody=hello_world
9191
return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
9292

9393

@@ -100,7 +100,7 @@ async def html_response2(): # $ requestHandler
100100
@app.get("/redirect") # $ routeSetup="/redirect"
101101
async def redirect(): # $ requestHandler
102102
next = "https://www.example.com"
103-
resp = fastapi.responses.RedirectResponse(next) # $ MISSING: HttpResponse HttpRedirectResponse redirectLocation=next
103+
resp = fastapi.responses.RedirectResponse(next) # $ HttpResponse HttpRedirectResponse redirectLocation=next
104104
return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
105105

106106

@@ -121,7 +121,7 @@ async def content():
121121
await asyncio.sleep(0.5)
122122
yield b"!"
123123

124-
resp = fastapi.responses.StreamingResponse(content()) # $ MISSING: HttpResponse responseBody=content()
124+
resp = fastapi.responses.StreamingResponse(content()) # $ HttpResponse responseBody=content()
125125
return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
126126

127127

@@ -136,7 +136,7 @@ async def file_response(): # $ requestHandler
136136

137137
# We don't really have any good QL modeling of passing a file-path, whose content
138138
# will be returned as part of the response... so will leave this as a TODO for now.
139-
resp = fastapi.responses.FileResponse(__file__) # $ MISSING: HttpResponse
139+
resp = fastapi.responses.FileResponse(__file__) # $ HttpResponse
140140
return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
141141

142142

0 commit comments

Comments
 (0)