Python: ORM: Add tests for inheritance

python/ql/test/library-tests/frameworks/django-orm/README.md

@@ -1,13 +1,16 @@
-The main test files are:
+The interesting ORM tests files can be found under `testapp/orm_*.py`. These are set up to be executed by the [testapp/tests.py](testapp/tests.py) file.
+
+List of interesting tests files (that might go out of date if it is forgotten :flushed:):
 
 - [testapp/orm_tests.py](testapp/orm_tests.py): which tests flow from source to sink
 - [testapp/orm_security_tests.py](testapp/orm_form_test.py): shows how forms can be used to save Models to the DB
 - [testapp/orm_security_tests.py](testapp/orm_security_tests.py): which highlights some interesting interactions with security queries
+- [testapp/orm_inheritance.py](testapp/orm_inheritance.py): which highlights how inheritance of ORM models works
 
 ## Setup
 
 ```
-pip install django pytest pytest-django
+pip install django pytest pytest-django django-polymorphic
 ```
 
 ## Run server

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

@@ -4,3 +4,4 @@ from django.db import models
 from .orm_tests import *
 from .orm_security_tests import *
 from .orm_form_test import *
+from .orm_inheritance import *

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

@@ -0,0 +1,194 @@
+from django.db import models
+
+SOURCE = "source"
+
+
+def SINK(arg):
+    print(arg)
+    assert arg == SOURCE
+
+
+def SINK_F(arg):
+    print(arg)
+    assert arg != SOURCE
+
+
+# ==============================================================================
+# Inheritance
+#
+# If base class defines a field, there can be
+# 1. flow when field is assigned on subclass construction to lookup of base class
+# 2. no flow from field assignment on subclass A to lookup of sibling subclass B
+# 3. no flow from field assignment on base class to lookup of subclass
+# ==============================================================================
+
+# ------------------------------------------------------------------------------
+# Inheritance with vanilla Django
+# ------------------------------------------------------------------------------
+
+class Book(models.Model):
+    title = models.CharField(max_length=256)
+
+
+class PhysicalBook(Book):
+    physical_location = models.CharField(max_length=256)
+    same_name_different_value = models.CharField(max_length=256)
+
+
+class EBook(Book):
+    download_link = models.CharField(max_length=256)
+    same_name_different_value = models.CharField(max_length=256)
+
+
+def save_base_book():
+    return Book.objects.create(
+        title=SOURCE,
+    )
+
+
+def fetch_book(id):
+    book = Book.objects.get(id=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"
+    # 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:
+        if book.title == "safe ebook":
+            pass
+        else:
+            raise
+
+    assert not isinstance(book, PhysicalBook)
+    assert not isinstance(book, EBook)
+
+    try:
+        SINK_F(book.physical_location)
+        raise Exception("This field is not available with vanilla Django")
+    except AttributeError:
+        pass
+
+
+def save_physical_book():
+    return PhysicalBook.objects.create(
+        title=SOURCE,
+        physical_location=SOURCE,
+        same_name_different_value=SOURCE,
+    )
+
+
+def fetch_physical_book(id):
+    book = PhysicalBook.objects.get(id=id)
+
+    # This sink should have only 1 sources, from `save_physical_book`
+    SINK(book.title) # $ flow="SOURCE, l:-10 -> book.title"
+    SINK(book.physical_location) # $ flow="SOURCE, l:-10 -> book.physical_location"
+    SINK(book.same_name_different_value) # $ flow="SOURCE, l:-10 -> book.same_name_different_value"
+
+
+def save_ebook():
+        return EBook.objects.create(
+        title="safe ebook",
+        download_link="safe",
+        same_name_different_value="safe",
+    )
+
+
+def fetch_ebook(id):
+    book = EBook.objects.get(id=id)
+
+    SINK_F(book.title)
+    SINK_F(book.download_link)
+    SINK_F(book.same_name_different_value)
+
+
+# ------------------------------------------------------------------------------
+# Inheritance with `django-polymorphic`, which automatically turns lookups on the
+# base class into the right subclass
+#
+# see https://django-polymorphic.readthedocs.io/en/stable/quickstart.html
+# ------------------------------------------------------------------------------
+
+from polymorphic.models import PolymorphicModel
+
+
+class PolyBook(PolymorphicModel):
+    title = models.CharField(max_length=256)
+
+
+class PolyPhysicalBook(PolyBook):
+    physical_location = models.CharField(max_length=256)
+    same_name_different_value = models.CharField(max_length=256)
+
+
+class PolyEBook(PolyBook):
+    download_link = models.CharField(max_length=256)
+    same_name_different_value = models.CharField(max_length=256)
+
+
+def poly_save_base_book():
+    return PolyBook.objects.create(
+        title=SOURCE
+    )
+
+
+def poly_fetch_book(id, test_for_subclass=True):
+    book = PolyBook.objects.get(id=id)
+
+    try:
+        # This sink should have 2 sources, from `poly_save_base_book` and
+        # `poly_save_physical_book`
+        SINK(book.title) # $ MISSING: flow
+    # 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:
+        if book.title == "safe ebook":
+            pass
+        else:
+            raise
+
+    if test_for_subclass:
+        assert isinstance(book, PolyPhysicalBook) or isinstance(book, PolyEBook)
+
+    if isinstance(book, PolyPhysicalBook):
+        SINK(book.title) # $ MISSING: flow
+        SINK(book.physical_location) # $ MISSING: flow
+        SINK(book.same_name_different_value) # $ MISSING: flow
+    elif isinstance(book, PolyEBook):
+        SINK_F(book.title)
+        SINK_F(book.download_link)
+        SINK_F(book.same_name_different_value)
+
+
+def poly_save_physical_book():
+    return PolyPhysicalBook.objects.create(
+        title=SOURCE,
+        physical_location=SOURCE,
+        same_name_different_value=SOURCE,
+    )
+
+
+def poly_fetch_physical_book(id):
+    book = PolyPhysicalBook.objects.get(id=id)
+
+    SINK(book.title) # $ MISSING: flow
+    SINK(book.physical_location) # $ MISSING: flow
+    SINK(book.same_name_different_value) # $ MISSING: flow
+
+
+def poly_save_ebook():
+        return PolyEBook.objects.create(
+        title="safe ebook",
+        download_link="safe",
+        same_name_different_value="safe",
+    )
+
+
+def poly_fetch_ebook(id):
+    book = PolyEBook.objects.get(id=id)
+
+    SINK_F(book.title)
+    SINK_F(book.download_link)
+    SINK_F(book.same_name_different_value)

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

@@ -83,3 +83,61 @@ def test_none_all():
     assert len(MyModel.objects.all()) == 1
     assert len(MyModel.objects.none().all()) == 0
     assert len(MyModel.objects.all().none()) == 0
+
+
+@pytest.mark.django_db
+def test_orm_inheritance():
+    from .orm_inheritance import (save_physical_book, save_ebook, save_base_book,
+        fetch_book, fetch_physical_book, fetch_ebook,
+        PhysicalBook, EBook,
+    )
+
+    base = save_base_book()
+    physical = save_physical_book()
+    ebook = save_ebook()
+
+    fetch_book(base.id)
+    fetch_book(physical.id)
+    fetch_book(ebook.id)
+
+    fetch_physical_book(physical.id)
+    fetch_ebook(ebook.id)
+
+    try:
+        fetch_physical_book(base.id)
+    except PhysicalBook.DoesNotExist:
+        pass
+
+    try:
+        fetch_ebook(ebook.id)
+    except EBook.DoesNotExist:
+        pass
+
+
+@pytest.mark.django_db
+def test_poly_orm_inheritance():
+    from .orm_inheritance import (poly_save_physical_book, poly_save_ebook, poly_save_base_book,
+        poly_fetch_book, poly_fetch_physical_book, poly_fetch_ebook,
+        PolyPhysicalBook, PolyEBook,
+    )
+
+    base = poly_save_base_book()
+    physical = poly_save_physical_book()
+    ebook = poly_save_ebook()
+
+    poly_fetch_book(base.id, test_for_subclass=False)
+    poly_fetch_book(physical.id)
+    poly_fetch_book(ebook.id)
+
+    poly_fetch_physical_book(physical.id)
+    poly_fetch_ebook(ebook.id)
+
+    try:
+        poly_fetch_physical_book(base.id)
+    except PolyPhysicalBook.DoesNotExist:
+        pass
+
+    try:
+        poly_fetch_ebook(ebook.id)
+    except PolyEBook.DoesNotExist:
+        pass