Python: ORM: add flow to base-class

RasmusWL · RasmusWL · commit 48fba87273ab · 2022-02-28T16:38:40.000+01:00
diff --git a/python/ql/lib/semmle/python/frameworks/Django.qll b/python/ql/lib/semmle/python/frameworks/Django.qll
@@ -826,6 +826,32 @@ module PrivateDjango {
             API::Node getModelClass() { result = modelClass }
           }
 
+          /**
+           * Gets a synthetic node where the data in the attribute `fieldName` can flow
+           * to, when a DB store is made on `subModel`, taking ORM inheritance into
+           * account.
+           *
+           * If `fieldName` is defined in class `base`, the results will include the
+           * synthetic node for `base` itself, the synthetic node for `subModel`, as
+           * well as all the classes in-between (in the class hierarchy).
+           */
+          SyntheticDjangoOrmModelNode nodeToStoreIn(API::Node subModel, string fieldName) {
+            exists(Class base, API::Node baseModel, API::Node resultModel |
+              baseModel = Model::subclassRef() and
+              resultModel = Model::subclassRef() and
+              baseModel.getASubclass*() = subModel and
+              base = getModelClassClass(baseModel) and
+              exists(Variable v |
+                base.getBody().getAnItem().(AssignStmt).defines(v) and
+                v.getId() = fieldName
+              )
+            |
+              baseModel.getASubclass*() = resultModel and
+              resultModel.getASubclass*() = subModel and
+              result.getModelClass() = resultModel
+            )
+          }
+
           /** Additional data-flow steps for Django ORM models. */
           class DjangOrmSteps extends AdditionalOrmSteps {
             override predicate storeStep(
@@ -867,7 +893,7 @@ module PrivateDjango {
                   )
                   or
                   // -> DB store on synthetic node
-                  nodeTo.(SyntheticDjangoOrmModelNode).getModelClass() = modelClass
+                  nodeTo = nodeToStoreIn(modelClass, fieldName)
                 )
               )
               or
@@ -877,7 +903,7 @@ module PrivateDjango {
                 call = [manager(modelClass), querySet(modelClass)].getMember("update").getACall() and
                 nodeFrom = call.getArgByName(fieldName) and
                 c.(DataFlow::AttributeContent).getAttribute() = fieldName and
-                nodeTo.(SyntheticDjangoOrmModelNode).getModelClass() = modelClass
+                nodeTo = nodeToStoreIn(modelClass, fieldName)
               )
               or
               // synthetic -> method-call that returns collection of ORM models (all/filter/...)
@@ -900,7 +926,12 @@ module PrivateDjango {
             override predicate jumpStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
               // save -> synthetic
               exists(API::Node modelClass, DataFlow::MethodCallNode saveCall |
-                nodeTo.(SyntheticDjangoOrmModelNode).getModelClass() = modelClass and
+                // TODO: The `nodeTo` should be restricted more, such that flow to
+                // base-classes are only for the fields that are defined in the
+                // base-class... but only passing on flow for a specific attribute requires flow-summaries,
+                // so we can do
+                // `obj (in obj.save call) ==read of attr==> synthetic attr on base-class ==store of attr==> synthetic for base-class`
+                nodeTo = nodeToStoreIn(modelClass, _) and
                 saveCall.calls(Model::instance(modelClass), "save") and
                 nodeFrom = saveCall.getObject()
               )
diff --git a/python/ql/test/library-tests/frameworks/django-orm/testapp/orm_inheritance.py b/python/ql/test/library-tests/frameworks/django-orm/testapp/orm_inheritance.py
@@ -52,7 +52,7 @@ def fetch_book(id):
     try:
         # This sink should have 2 sources, from `save_base_book` and
         # `save_physical_book`
-        SINK(book.title) # $ flow="SOURCE, l:-10 -> book.title"
+        SINK(book.title) # $ flow="SOURCE, l:-10 -> book.title" flow="SOURCE, l:+21 -> book.title"
     # The sink assertion will fail for the EBook, which we handle. The title attribute
     # of a Book could be tainted, so we want this to be a sink in general.
     except AssertionError:
@@ -140,7 +140,7 @@ def poly_fetch_book(id, test_for_subclass=True):
     try:
         # This sink should have 2 sources, from `poly_save_base_book` and
         # `poly_save_physical_book`
-        SINK(book.title) # $ flow="SOURCE, l:-10 -> book.title"
+        SINK(book.title) # $ flow="SOURCE, l:-10 -> book.title" flow="SOURCE, l:+24 -> book.title"
     # The sink assertion will fail for the PolyEBook, which we handle. The title
     # attribute of a PolyBook could be tainted, so we want this to be a sink in general.
     except AssertionError:
@@ -153,11 +153,11 @@ def poly_fetch_book(id, test_for_subclass=True):
         assert isinstance(book, PolyPhysicalBook) or isinstance(book, PolyEBook)
 
     if isinstance(book, PolyPhysicalBook):
-        SINK(book.title) # $ MISSING: flow="SOURCE, l:+11 -> book.title" SPURIOUS: flow="SOURCE, l:-23 -> book.title"
+        SINK(book.title) # $ flow="SOURCE, l:+11 -> book.title" SPURIOUS: flow="SOURCE, l:-23 -> book.title"
         SINK(book.physical_location) # $ MISSING: flow
         SINK(book.same_name_different_value) # $ MISSING: flow
     elif isinstance(book, PolyEBook):
-        SINK_F(book.title) # $ SPURIOUS: flow="SOURCE, l:-27 -> book.title"
+        SINK_F(book.title) # $ SPURIOUS: flow="SOURCE, l:-27 -> book.title" flow="SOURCE, l:+7 -> book.title"
         SINK_F(book.download_link)
         SINK_F(book.same_name_different_value)
 
