diff --git a/python/ql/test/library-tests/frameworks/rest_framework/taint_test.py b/python/ql/test/library-tests/frameworks/rest_framework/taint_test.py index ff30b00ca99..31daba8d6c6 100644 --- a/python/ql/test/library-tests/frameworks/rest_framework/taint_test.py +++ b/python/ql/test/library-tests/frameworks/rest_framework/taint_test.py @@ -107,12 +107,30 @@ class MyClass(APIView): return Response("ok") # $ HttpResponse +# Viewsets +# see https://www.django-rest-framework.org/api-guide/viewsets/ + +class MyModelViewSet(viewsets.ModelViewSet): + def retrieve(self, request, *args, **kwargs): # $ requestHandler + ensure_tainted( + request, # $ tainted + request.GET, # $ tainted + request.GET.get("pk"), # $ tainted + ) + + ensure_tainted( + kwargs, # $ tainted + kwargs["pk"], # $ tainted + kwargs.get("pk"), # $ tainted + ) + return Response("retrieve") # $ HttpResponse # fake setup, you can't actually run this urlpatterns = [ path("test-taint/", test_taint), # $ routeSetup="test-taint/" - path("ClassView/", MyClass.as_view()), # $ routeSetup="ClassView/" + path("ClassView/", MyClass.as_view()), # $ routeSetup="ClassView/", + path("MyModelViewSet/", MyModelViewSet.as_view()) # $ routeSetup="MyModelViewSet/", ] # tests with no route-setup, but we can still tell that these are using Django REST