JS: More data flow through classes

javascript/ql/src/semmle/javascript/dataflow/Configuration.qll

@@ -605,11 +605,16 @@ private predicate storeStep(
   basicStoreStep(pred, succ, prop) and
   summary = PathSummary::level()
   or
-  exists(Function f, DataFlow::Node mid, DataFlow::Node base |
+  exists(Function f, DataFlow::Node mid |
     // `f` stores its parameter `pred` in property `prop` of a value that it returns,
     // and `succ` is an invocation of `f`
     reachableFromInput(f, succ, pred, mid, cfg, summary) and
-    returnedPropWrite(f, base, prop, mid)
+    (
+      returnedPropWrite(f, _, prop, mid)
+      or
+      succ instanceof DataFlow::NewNode and
+      receiverPropWrite(f, prop, mid)
+    )
   )
 }
 
@@ -621,6 +626,13 @@ predicate returnedPropWrite(Function f, DataFlow::SourceNode base, string prop,
   base.flowsToExpr(f.getAReturnedExpr())
 }
 
+/**
+ * Holds if `f` may return `base`, which has a write of property `prop` with right-hand side `rhs`.
+ */
+predicate receiverPropWrite(Function f, string prop, DataFlow::Node rhs) {
+  DataFlow::thisNode(f).hasPropertyWrite(prop, rhs)
+}
+
 /**
  * Holds if `rhs` is the right-hand side of a write to property `prop`, and `nd` is reachable
  * from the base of that write under configuration `cfg` (possibly through callees) along a

javascript/ql/src/semmle/javascript/dataflow/DataFlow.qll

@@ -802,6 +802,13 @@ module DataFlow {
    */
   predicate thisNode(DataFlow::Node node, StmtContainer container) { node = TThisNode(container) }
 
+  /**
+   * Gets the node representing the receiver of the given function, or `this` in the given top-level.
+   *
+   * Has no result if `container` is an arrow function.
+   */
+  DataFlow::ThisNode thisNode(StmtContainer container) { result = TThisNode(container) }
+
   /**
    * A classification of flows that are not modeled, or only modeled incompletely, by
    * `DataFlowNode`:

javascript/ql/src/semmle/javascript/dataflow/Nodes.qll

@@ -462,3 +462,138 @@ DataFlow::SourceNode moduleMember(string path, string m) {
     result = DataFlow::ssaDefinitionNode(ssa)
   )
 }
+
+/**
+ * A data flow node corresponding to a class definition or a function definition
+ * acting as a class.
+ *
+ * The following patterns are recognized as classes and methods:
+ * ```
+ * class C {
+ *   method()
+ * }
+ * 
+ * function F() {}
+ *
+ * F.prototype.method = function() {}
+ * 
+ * F.prototype = {
+ *   method: function() {}
+ * }
+ * 
+ * extend(F.prototype, {
+ *   method: function() {}
+ * }); 
+ * ```
+ */
+class ClassNode extends DataFlow::SourceNode, DataFlow::ValueNode {
+  ClassNode() {
+    astNode instanceof ClassDefinition
+    or
+    astNode instanceof Function and
+    exists(getAPropertyReference("prototype"))
+  }
+
+  /**
+   * Gets the name of the class, if it has one.
+   */
+  string getName() {
+    result = astNode.(ClassDefinition).getName()
+    or
+    result = astNode.(Function).getName()
+  }
+
+  /**
+   * Gets a description of the class.
+   */
+  string describe() {
+    result = astNode.(ClassDefinition).describe()
+    or
+    result = astNode.(Function).describe()
+  }
+
+  /**
+   * Gets the constructor function of this class.
+   */
+  FunctionNode getConstructor() {
+    result = astNode.(ClassDefinition).getConstructor().getBody().flow()
+    or
+    result = this
+  }
+
+  /**
+   * Gets an instance method with the given name, if any.
+   */
+  FunctionNode getInstanceMethod(string name) {
+    exists (MethodDeclaration method |
+      method = astNode.(ClassDefinition).getMethod(name) and
+      not method.isStatic() and
+      not method.isAmbient() and
+      not method instanceof ConstructorDeclaration and
+      result = method.getBody().flow()
+    )
+    or
+    result = getAPrototypeReference().getAPropertyWrite(name).getRhs().getALocalSource()
+  }
+
+  /**
+   * Gets an instance method of this class.
+   *
+   * The constructor is not considered an instance method.
+   */
+  FunctionNode getAnInstanceMethod() {
+    exists (MethodDeclaration method |
+      method = astNode.(ClassDefinition).getAMethod() and
+      not method.isStatic() and
+      not method.isAmbient() and
+      not method instanceof ConstructorDeclaration and
+      result = method.getBody().flow()
+    )
+    or
+    result = getAPrototypeReference().getAPropertyWrite().getRhs().getALocalSource()
+  }
+
+  /**
+   * Gets the static method of this class with the given name.
+   */
+  FunctionNode getStaticMethod(string name) {
+    exists (MethodDeclaration method |
+      method = astNode.(ClassDefinition).getMethod(name) and
+      method.isStatic() and
+      not method.isAmbient() and
+      result = method.getBody().flow()
+    )
+    or
+    result = getAPropertyWrite(name).getRhs().getALocalSource()
+  }
+
+  /**
+   * Gets a static method of this class.
+   *
+   * The constructor is not considered a static method.
+   */
+  FunctionNode getAStaticMethod() {
+    exists (MethodDeclaration method |
+      method = astNode.(ClassDefinition).getAMethod() and
+      method.isStatic() and
+      not method.isAmbient() and
+      result = method.getBody().flow()
+    )
+    or
+    result = getAPropertyWrite().getRhs().getALocalSource()
+  }
+
+  /**
+   * Gets a reference to the prototype of this class.
+   */
+  private DataFlow::SourceNode getAPrototypeReference() {
+    result = getAPropertyRead("prototype")
+    or
+    result = getAPropertySource("prototype")
+    or
+    exists (ExtendCall call |
+      call.getDestinationOperand() = getAPropertyRead("prototype") and
+      result = call.getASourceOperand()
+    )
+  }
+}

javascript/ql/src/semmle/javascript/dataflow/internal/FlowSteps.qll

@@ -115,16 +115,31 @@ predicate callStep(DataFlow::Node pred, DataFlow::Node succ) {
     argumentPassing(_, pred, _, parm) and
     succ = DataFlow::parameterNode(parm)
   )
+  or
+  constructorReceiverToMethodStep(pred, succ)
+}
+
+/**
+ * Holds if `pred` is the receiver of a class constructor and `succ` is the receiver of one
+ * of the methods of that class.
+ */
+predicate constructorReceiverToMethodStep(DataFlow::Node pred, DataFlow::Node succ) {
+  exists (DataFlow::ClassNode class_ |
+    pred = class_.getConstructor().getReceiver() and
+    succ = class_.getAnInstanceMethod().getReceiver()
+  )
 }
 
 /**
  * Holds if there is a flow step from `pred` to `succ` through returning
  * from a function call.
  */
 predicate returnStep(DataFlow::Node pred, DataFlow::Node succ) {
-  exists(Function f |
-    returnExpr(f, pred, _) and
-    calls(succ, f)
+  exists(Function f | calls(succ, f) |
+    returnExpr(f, pred, _)
+    or
+    succ instanceof DataFlow::NewNode and
+    DataFlow::thisNode(pred, f)
   )
 }
 

javascript/ql/test/library-tests/Classes/ClassNodeConstructor.expected

@@ -0,0 +1,7 @@
+| fields.js:1:1:4:1 | class C ...  = 42\\n} | fields.js:1:9:1:8 | () {} |
+| points.js:1:1:18:1 | class P ... ;\\n  }\\n} | points.js:2:14:5:3 | (x, y)  ...  y;\\n  } |
+| points.js:20:1:33:1 | class C ... ;\\n  }\\n} | points.js:21:14:24:3 | (x, y,  ...  c;\\n  } |
+| staticConstructor.js:1:1:3:1 | class M ... r"; }\\n} | staticConstructor.js:1:15:1:14 | () {} |
+| tst.js:1:9:4:1 | class { ...  */ }\\n} | tst.js:2:16:2:50 | () { /* ... r. */ } |
+| tst.js:6:1:8:1 | class B ... t); }\\n} | tst.js:7:14:7:38 | () { su ... get); } |
+| tst.js:11:1:14:1 | class C ... () {}\\n} | tst.js:11:9:11:8 | () {} |

javascript/ql/test/library-tests/Classes/ClassNodeConstructor.ql

@@ -0,0 +1,4 @@
+import javascript
+
+from DataFlow::ClassNode class_
+select class_, class_.getConstructor()

javascript/ql/test/library-tests/Classes/ClassNodeInstanceMethod.expected

@@ -0,0 +1,5 @@
+| points.js:1:1:18:1 | class P ... ;\\n  }\\n} | dist | points.js:7:11:9:3 | () {\\n   ... y);\\n  } |
+| points.js:1:1:18:1 | class P ... ;\\n  }\\n} | toString | points.js:11:11:13:3 | () {\\n   ... )";\\n  } |
+| points.js:20:1:33:1 | class C ... ;\\n  }\\n} | toString | points.js:26:11:28:3 | () {\\n   ... ur;\\n  } |
+| tst.js:1:9:4:1 | class { ...  */ }\\n} | constructor | tst.js:3:18:3:56 | () { /* ... r. */ } |
+| tst.js:11:1:14:1 | class C ... () {}\\n} | m | tst.js:12:4:12:8 | () {} |

javascript/ql/test/library-tests/Classes/ClassNodeInstanceMethod.ql

@@ -0,0 +1,4 @@
+import javascript
+
+from DataFlow::ClassNode class_, string name
+select class_, name, class_.getInstanceMethod(name)

javascript/ql/test/library-tests/Classes/ClassNodeStaticMethod.expected

@@ -0,0 +1,3 @@
+| points.js:1:1:18:1 | class P ... ;\\n  }\\n} | className | points.js:15:19:17:3 | () {\\n   ... t";\\n  } |
+| points.js:20:1:33:1 | class C ... ;\\n  }\\n} | className | points.js:30:19:32:3 | () {\\n   ... t";\\n  } |
+| staticConstructor.js:1:1:3:1 | class M ... r"; }\\n} | constructor | staticConstructor.js:2:21:2:59 | () { re ... tor"; } |

javascript/ql/test/library-tests/Classes/ClassNodeStaticMethod.ql

@@ -0,0 +1,4 @@
+import javascript
+
+from DataFlow::ClassNode class_, string name
+select class_, name, class_.getStaticMethod(name)

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

@@ -1,4 +1,18 @@
 | advanced-callgraph.js:2:13:2:20 | source() | advanced-callgraph.js:6:22:6:22 | v |
+| class-fields.js:6:18:6:25 | source() | class-fields.js:11:10:11:19 | this.taint |
+| class-fields.js:24:25:24:32 | source() | class-fields.js:10:10:10:19 | this.param |
+| class-fields.js:32:16:32:23 | source() | class-fields.js:37:8:37:17 | this.taint |
+| class-fields.js:49:31:49:38 | source() | class-fields.js:36:8:36:17 | this.param |
+| class-fields.js:57:16:57:23 | source() | class-fields.js:63:10:63:19 | this.taint |
+| class-fields.js:76:36:76:43 | source() | class-fields.js:62:10:62:19 | this.param |
+| class-fields.js:84:16:84:23 | source() | class-fields.js:91:10:91:19 | this.taint |
+| class-fields.js:104:35:104:42 | source() | class-fields.js:90:10:90:19 | this.param |
+| constructor-calls.js:4:18:4:25 | source() | constructor-calls.js:18:8:18:14 | c.taint |
+| constructor-calls.js:4:18:4:25 | source() | constructor-calls.js:22:8:22:19 | c_safe.taint |
+| constructor-calls.js:10:16:10:23 | source() | constructor-calls.js:26:8:26:14 | d.taint |
+| constructor-calls.js:10:16:10:23 | source() | constructor-calls.js:30:8:30:19 | d_safe.taint |
+| constructor-calls.js:14:15:14:22 | source() | constructor-calls.js:17:8:17:14 | c.param |
+| constructor-calls.js:14:15:14:22 | source() | constructor-calls.js:25:8:25:14 | d.param |
 | partialCalls.js:4:17:4:24 | source() | partialCalls.js:17:14:17:14 | x |
 | partialCalls.js:4:17:4:24 | source() | partialCalls.js:20:14:20:14 | y |
 | partialCalls.js:4:17:4:24 | source() | partialCalls.js:30:14:30:20 | x.value |

javascript/ql/test/library-tests/TaintTracking/class-fields.js

@@ -0,0 +1,107 @@
+import * as dummy from 'dummy'; // treat as module
+
+class EcmaClass {
+  constructor(param) {
+    this.param = param;
+    this.taint = source();
+  }
+
+  fieldSinks() {
+    sink(this.param); // NOT OK
+    sink(this.taint); // NOT OK
+  }
+
+  getParam() {
+    return this.param;
+  }
+
+  getTaint() {
+    return this.taint;
+  }
+}
+
+function testEcmaClass() {
+  let c = new EcmaClass(source());
+  sink(c.getParam()); // NOT OK
+  sink(c.getTaint()); // NOT OK
+}
+
+
+function ProtoStyleClass(param) {
+  this.param = param;
+  this.taint = source();
+}
+
+ProtoStyleClass.prototype.fieldSinks = function() {
+  sink(this.param); // NOT OK
+  sink(this.taint); // NOT OK
+};
+
+ProtoStyleClass.prototype.getParam = function() {
+  return this.param;
+};
+
+ProtoStyleClass.prototype.getTaint = function() {
+  return this.taint;
+};
+
+function testProtoStyleClass() {
+  let c = new ProtoStyleClass(source());
+  sink(c.getParam()); // NOT OK
+  sink(c.getTaint()); // NOT OK
+}
+
+
+function ProtoAssignmentClass(param) {
+  this.param = param;
+  this.taint = source();
+}
+
+ProtoAssignmentClass.prototype = {
+  fieldSinks() {
+    sink(this.param); // NOT OK
+    sink(this.taint); // NOT OK
+  },
+
+  getParam() {
+    return this.param;
+  },
+
+  getTaint() {
+    return this.taint;
+  }
+};
+
+function testProtoAssignmentClass() {
+  let c = new ProtoAssignmentClass(source());
+  sink(c.getParam()); // NOT OK
+  sink(c.getTaint()); // NOT OK
+}
+
+
+function ProtoExtensionClass(param) {
+  this.param = param;
+  this.taint = source();
+}
+
+let extend = require('extend');
+extend(ProtoExtensionClass.prototype, {
+  fieldSinks() {
+    sink(this.param); // NOT OK
+    sink(this.taint); // NOT OK
+  },
+
+  getParam() {
+    return this.param;
+  },
+
+  getTaint() {
+    return this.taint;
+  }
+});
+
+function testProtoExtensionClass() {
+  let c = new ProtoExtensionClass(source());
+  sink(c.getParam()); // NOT OK
+  sink(c.getTaint()); // NOT OK
+}