From 79c0ed607487f7c59d903113f576dfdfde60bcca Mon Sep 17 00:00:00 2001
From: Joe Farebrother <joefarebrother@github.com>
Date: Fri, 21 Jun 2024 14:01:33 +0100
Subject: [PATCH] Add additional fastapi mheader write models

---
 .../lib/semmle/python/frameworks/FastApi.qll  | 28 +++++++++++++++++++
 .../frameworks/fastapi/response_test.py       |  4 +--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/python/ql/lib/semmle/python/frameworks/FastApi.qll b/python/ql/lib/semmle/python/frameworks/FastApi.qll
index 423f8580a5b..028c5f26601 100644
--- a/python/ql/lib/semmle/python/frameworks/FastApi.qll
+++ b/python/ql/lib/semmle/python/frameworks/FastApi.qll
@@ -383,5 +383,33 @@ module FastApi {
 
       override predicate valueAllowsNewline() { none() }
     }
+
+    class HeaderSubscriptWrite extends Http::Server::ResponseHeaderWrite::Range {
+      DataFlow::Node index;
+      DataFlow::Node value;
+
+      HeaderSubscriptWrite() {
+        exists(SubscriptNode subscript, DataFlow::AttrRead headerLookup |
+          // To give `this` a value, we need to choose between either LHS or RHS,
+          // and just go with the LHS
+          this.asCfgNode() = subscript
+        |
+          headerLookup.accesses(instance(), "headers") and
+          exists(DataFlow::Node subscriptObj | subscriptObj.asCfgNode() = subscript.getObject() |
+            headerLookup.flowsTo(subscriptObj)
+          ) and
+          value.asCfgNode() = subscript.(DefinitionNode).getValue() and
+          index.asCfgNode() = subscript.getIndex()
+        )
+      }
+
+      override DataFlow::Node getNameArg() { result = index }
+
+      override DataFlow::Node getValueArg() { result = value }
+
+      override predicate nameAllowsNewline() { none() }
+
+      override predicate valueAllowsNewline() { none() }
+    }
   }
 }
diff --git a/python/ql/test/library-tests/frameworks/fastapi/response_test.py b/python/ql/test/library-tests/frameworks/fastapi/response_test.py
index 44582d6cd6e..2bc26c15fda 100644
--- a/python/ql/test/library-tests/frameworks/fastapi/response_test.py
+++ b/python/ql/test/library-tests/frameworks/fastapi/response_test.py
@@ -13,7 +13,7 @@ async def response_parameter(response: Response): # $ requestHandler
     response.set_cookie(key="key", value="value") # $ CookieWrite CookieName="key" CookieValue="value"
     response.headers.append("Set-Cookie", "key2=value2") # $ headerWriteName="Set-Cookie" headerWriteValue="key2=value2" CookieWrite CookieRawHeader="key2=value2"
     response.headers.append(key="Set-Cookie", value="key2=value2") # $ headerWriteName="Set-Cookie" headerWriteValue="key2=value2" CookieWrite CookieRawHeader="key2=value2"
-    response.headers["X-MyHeader"] = "header-value" # $ MISSING: headerWriteName="X-MyHeader" headerWriteValue="header-value"
+    response.headers["X-MyHeader"] = "header-value" # $ headerWriteName="X-MyHeader" headerWriteValue="header-value"
     response.status_code = 418
     return {"message": "response as parameter"} # $ HttpResponse mimetype=application/json responseBody=Dict
 
@@ -45,7 +45,7 @@ async def response_parameter_custom_type(response: MyXmlResponse): # $ requestHa
     print(type(response))
     assert type(response) == fastapi.responses.Response
     response.set_cookie("key", "value") # $ CookieWrite CookieName="key" CookieValue="value"
-    response.headers["Custom-Response-Type"] = "yes, but only after function has run" # $ MISSING: headerWriteName="Custom-Response-Typer" headerWriteValue="yes, but only after function has run"
+    response.headers["Custom-Response-Type"] = "yes, but only after function has run" # $ headerWriteName="Custom-Response-Type" headerWriteValue="yes, but only after function has run"
     xml_data = "<foo>FOO</foo>"
     return xml_data # $ HttpResponse responseBody=xml_data mimetype=application/xml