Python: ORM: Handle load of PolymorphicModels

Author: RasmusWL

python/ql/lib/semmle/python/frameworks/Django.qll

@@ -852,6 +852,26 @@ module PrivateDjango {
             )
           }
 
+          /**
+           * Gets the synthetic node where data could be loaded from, when a fetch is
+           * made on `modelClass`.
+           *
+           * In vanilla Django inheritance, this is simply the model itself, but if a
+           * model is based on `polymorphic.models.PolymorphicModel`, a fetch of the
+           * base-class can also yield instances of its subclasses.
+           */
+          SyntheticDjangoOrmModelNode nodeToLoadFrom(API::Node modelClass) {
+            result.getModelClass() = modelClass
+            or
+            exists(API::Node polymorphicModel |
+              polymorphicModel =
+                API::moduleImport("polymorphic").getMember("models").getMember("PolymorphicModel")
+            |
+              polymorphicModel.getASubclass+() = modelClass and
+              modelClass.getASubclass+() = result.getModelClass()
+            )
+          }
+
           /** Additional data-flow steps for Django ORM models. */
           class DjangOrmSteps extends AdditionalOrmSteps {
             override predicate storeStep(
@@ -908,15 +928,15 @@ module PrivateDjango {
               or
               // synthetic -> method-call that returns collection of ORM models (all/filter/...)
               exists(API::Node modelClass |
-                nodeFrom.(SyntheticDjangoOrmModelNode).getModelClass() = modelClass and
+                nodeFrom = nodeToLoadFrom(modelClass) and
                 nodeTo.(Model::QuerySetMethodInstanceCollection).getModelClass() = modelClass and
                 nodeTo.(Model::QuerySetMethodInstanceCollection).isDbFetch() and
                 c instanceof DataFlow::ListElementContent
               )
               or
               // synthetic -> method-call that returns dictionary with ORM models as values
               exists(API::Node modelClass |
-                nodeFrom.(SyntheticDjangoOrmModelNode).getModelClass() = modelClass and
+                nodeFrom = nodeToLoadFrom(modelClass) and
                 nodeTo.(Model::QuerySetMethodInstanceDictValue).getModelClass() = modelClass and
                 nodeTo.(Model::QuerySetMethodInstanceDictValue).isDbFetch() and
                 c instanceof DataFlow::DictionaryElementAnyContent
@@ -938,9 +958,9 @@ module PrivateDjango {
               or
               // synthetic -> method-call that returns single ORM model (get/first/...)
               exists(API::Node modelClass |
+                nodeFrom = nodeToLoadFrom(modelClass) and
                 nodeTo.(Model::InstanceSource).getModelClass() = modelClass and
-                nodeTo.(Model::InstanceSource).isDbFetch() and
-                nodeFrom.(SyntheticDjangoOrmModelNode).getModelClass() = modelClass
+                nodeTo.(Model::InstanceSource).isDbFetch()
               )
             }
           }

python/ql/test/library-tests/frameworks/django-orm/testapp/orm_inheritance.py

@@ -154,12 +154,12 @@ def poly_fetch_book(id, test_for_subclass=True):
 
     if isinstance(book, PolyPhysicalBook):
         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
+        SINK(book.physical_location) # $ flow="SOURCE, l:+11 -> book.physical_location"
+        SINK(book.same_name_different_value) # $ flow="SOURCE, l:+11 -> book.same_name_different_value"
     elif isinstance(book, PolyEBook):
         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)
+        SINK_F(book.same_name_different_value) # $ SPURIOUS: flow="SOURCE, l:+7 -> book.same_name_different_value"
 
 
 def poly_save_physical_book():