From 8a56b48357d7a6a59a44bf476a820ed516966444 Mon Sep 17 00:00:00 2001
From: Rasmus Wriedt Larsen <rasmuswl@github.com>
Date: Mon, 24 Oct 2022 15:35:53 +0200
Subject: [PATCH] Python: Support `super().__new__(cls)`

---
 .../dataflow/new/internal/DataFlowDispatch.qll  | 17 +++++++++++++++++
 .../CallGraph/InlineCallGraphTest.expected      |  3 +++
 .../CallGraph/code/class_construction.py        |  6 +++---
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll
index d0016c28673..5f957553a76 100644
--- a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll
+++ b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll
@@ -468,6 +468,13 @@ private TypeTrackingNode classInstanceTracker(TypeTracker t, Class cls) {
   t.start() and
   resolveClassCall(result.(CallCfgNode).asCfgNode(), cls)
   or
+  // result of `super().__new__` as used in a `__new__` method implementation
+  t.start() and
+  exists(Class classUsedInSuper |
+    fromSuperNewCall(result.(CallCfgNode).asCfgNode(), classUsedInSuper, _, _) and
+    classUsedInSuper = getADirectSuperclass*(cls)
+  )
+  or
   exists(TypeTracker t2 | result = classInstanceTracker(t2, cls).track(t2, t)) and
   not result.(ParameterNodeImpl).isParameterOf(_, any(ParameterPosition pp | pp.isSelf()))
 }
@@ -856,6 +863,16 @@ private module MethodCalls {
     attr.accesses(self, functionName)
   }
 
+  /**
+   * Like `fromSuper`, but only for `__new__`, and without requirement for being able to
+   * resolve the call to a known target (since the only super class might be the
+   * builtin `object`, so we never have the implementation of `__new__` in the DB).
+   */
+  predicate fromSuperNewCall(CallNode call, Class classUsedInSuper, AttrRead attr, Node self) {
+    fromSuper_join(call, "__new__", classUsedInSuper, attr, self) and
+    self in [classTracker(_), clsTracker(_)]
+  }
+
   /**
    * Holds if `call` is a call to a method `target`, derived from a use of `super`, either
    * as:
diff --git a/python/ql/test/experimental/library-tests/CallGraph/InlineCallGraphTest.expected b/python/ql/test/experimental/library-tests/CallGraph/InlineCallGraphTest.expected
index b3bc72cca1e..2478db5a060 100644
--- a/python/ql/test/experimental/library-tests/CallGraph/InlineCallGraphTest.expected
+++ b/python/ql/test/experimental/library-tests/CallGraph/InlineCallGraphTest.expected
@@ -21,6 +21,9 @@ pointsTo_found_typeTracker_notFound
 | code/type_tracking_limitation.py:8:1:8:3 | ControlFlowNode for x() | my_func |
 typeTracker_found_pointsTo_notFound
 | code/callable_as_argument.py:29:5:29:12 | ControlFlowNode for Attribute() | test_class.InsideTestFunc.sm |
+| code/class_construction.py:44:9:44:26 | ControlFlowNode for Attribute() | WithNew.some_method |
+| code/class_construction.py:61:9:61:26 | ControlFlowNode for Attribute() | WithNew.some_method |
+| code/class_construction.py:75:9:75:27 | ControlFlowNode for Attribute() | ExtraCallToInit.__init__ |
 | code/class_special_methods.py:22:9:22:16 | ControlFlowNode for self() | Base.__call__ |
 | code/class_special_methods.py:22:9:22:16 | ControlFlowNode for self() | Sub.__call__ |
 | code/class_special_methods.py:33:1:33:5 | ControlFlowNode for b() | Base.__call__ |
diff --git a/python/ql/test/experimental/library-tests/CallGraph/code/class_construction.py b/python/ql/test/experimental/library-tests/CallGraph/code/class_construction.py
index 06669902714..1ae696edf61 100644
--- a/python/ql/test/experimental/library-tests/CallGraph/code/class_construction.py
+++ b/python/ql/test/experimental/library-tests/CallGraph/code/class_construction.py
@@ -41,7 +41,7 @@ class WithNew(object):
         print("WithNew.__new__", arg)
         inst = super().__new__(cls)
         assert isinstance(inst, cls)
-        inst.some_method() # $ MISSING: pt,tt=WithNew.some_method
+        inst.some_method() # $ tt=WithNew.some_method
         return inst
 
     def __init__(self, arg=None):
@@ -58,7 +58,7 @@ class WithNewSub(WithNew):
         print("WithNewSub.__new__")
         inst = super().__new__(cls, 44.1) # $ pt,tt=WithNew.__new__
         assert isinstance(inst, cls)
-        inst.some_method() # $ MISSING: pt,tt=WithNew.some_method
+        inst.some_method() # $ tt=WithNew.some_method
         return inst
 
 WithNewSub() # $ tt=WithNewSub.__new__ tt=WithNew.__init__
@@ -72,7 +72,7 @@ class ExtraCallToInit(object):
         inst = super().__new__(cls)
         assert isinstance(inst, cls)
         # you're not supposed to do this, since it will cause the __init__ method will be run twice.
-        inst.__init__(1001) # $ MISSING: pt,tt=ExtraCallToInit.__init__
+        inst.__init__(1001) # $ tt=ExtraCallToInit.__init__
         return inst
 
     def __init__(self, arg):