JS: Track functions with methods

Author: asgerf

javascript/ql/lib/semmle/javascript/dataflow/internal/CallGraphs.qll

@@ -79,9 +79,9 @@ module CallGraph {
       cls.getAClassReference(t.continue()) = result
     )
     or
-    exists(DataFlow::ObjectLiteralNode object, string prop |
+    exists(DataFlow::SourceNode object, string prop |
       function = object.getAPropertySource(prop) and
-      result = getAnObjectLiteralRef(object).getAPropertyRead(prop) and
+      result = getAnAllocationSiteRef(object).getAPropertyRead(prop) and
       t.start()
     )
     or
@@ -203,21 +203,26 @@ module CallGraph {
     )
     or
     exists(DataFlow::ObjectLiteralNode object, string name |
-      ref = getAnObjectLiteralRef(object).getAPropertyRead(name) and
+      ref = getAnAllocationSiteRef(object).getAPropertyRead(name) and
       result = object.getPropertyGetter(name)
       or
-      ref = getAnObjectLiteralRef(object).getAPropertyWrite(name) and
+      ref = getAnAllocationSiteRef(object).getAPropertyWrite(name) and
       result = object.getPropertySetter(name)
     )
   }
 
-  private predicate shouldTrackObjectLiteral(DataFlow::ObjectLiteralNode node) {
+  private predicate shouldTrackObjectWithMethods(DataFlow::SourceNode node) {
     (
+      (
+        node instanceof DataFlow::ObjectLiteralNode
+        or
+        node instanceof DataFlow::FunctionNode
+      ) and
       node.getAPropertySource() instanceof DataFlow::FunctionNode
       or
-      exists(node.getPropertyGetter(_))
+      exists(node.(DataFlow::ObjectLiteralNode).getPropertyGetter(_))
       or
-      exists(node.getPropertySetter(_))
+      exists(node.(DataFlow::ObjectLiteralNode).getPropertySetter(_))
     ) and
     not node.getTopLevel().isExterns()
   }
@@ -228,14 +233,14 @@ module CallGraph {
    * To avoid false flow from callbacks passed in via "named parameters", we only track object
    * literals out of returns, not into calls.
    */
-  private StepSummary objectLiteralStep() { result = LevelStep() or result = ReturnStep() }
+  private StepSummary objectWithMethodsStep() { result = LevelStep() or result = ReturnStep() }
 
-  /** Gets a node that refers to the given object literal, via a limited form of type tracking. */
+  /** Gets a node that refers to the given object, via a limited form of type tracking. */
   cached
-  DataFlow::SourceNode getAnObjectLiteralRef(DataFlow::ObjectLiteralNode node) {
-    shouldTrackObjectLiteral(node) and
+  DataFlow::SourceNode getAnAllocationSiteRef(DataFlow::SourceNode node) {
+    shouldTrackObjectWithMethods(node) and
     result = node
     or
-    StepSummary::step(getAnObjectLiteralRef(node), result, objectLiteralStep())
+    StepSummary::step(getAnAllocationSiteRef(node), result, objectWithMethodsStep())
   }
 }

javascript/ql/test/library-tests/TaintTracking/BasicTaintTracking.expected

@@ -83,6 +83,8 @@ typeInferenceMismatch
 | exceptions.js:158:13:158:20 | source() | exceptions.js:161:10:161:10 | e |
 | factory-function.js:21:13:21:20 | source() | factory-function.js:7:10:7:12 | obj |
 | factory-function.js:22:13:22:20 | source() | factory-function.js:7:10:7:12 | obj |
+| factory-function.js:26:7:26:14 | source() | factory-function.js:16:14:16:16 | obj |
+| factory-function.js:27:7:27:14 | source() | factory-function.js:16:14:16:16 | obj |
 | getters-and-setters.js:6:20:6:27 | source() | getters-and-setters.js:9:10:9:18 | new C().x |
 | getters-and-setters.js:6:20:6:27 | source() | getters-and-setters.js:13:18:13:20 | c.x |
 | getters-and-setters.js:27:15:27:22 | source() | getters-and-setters.js:23:18:23:18 | v |

javascript/ql/test/library-tests/TaintTracking/DataFlowTracking.expected

@@ -45,6 +45,8 @@
 | exceptions.js:158:13:158:20 | source() | exceptions.js:161:10:161:10 | e |
 | factory-function.js:21:13:21:20 | source() | factory-function.js:7:10:7:12 | obj |
 | factory-function.js:22:13:22:20 | source() | factory-function.js:7:10:7:12 | obj |
+| factory-function.js:26:7:26:14 | source() | factory-function.js:16:14:16:16 | obj |
+| factory-function.js:27:7:27:14 | source() | factory-function.js:16:14:16:16 | obj |
 | getters-and-setters.js:6:20:6:27 | source() | getters-and-setters.js:9:10:9:18 | new C().x |
 | getters-and-setters.js:6:20:6:27 | source() | getters-and-setters.js:13:18:13:20 | c.x |
 | getters-and-setters.js:27:15:27:22 | source() | getters-and-setters.js:23:18:23:18 | v |