add support for subclasses of EventEmitter

javascript/ql/src/semmle/javascript/frameworks/NodeJSLib.qll

@@ -894,17 +894,52 @@ module NodeJSLib {
     DataFlow::SourceNode ref() { result = EventEmitter::trackEventEmitter(this) }
   }
 
+  /**
+   * Gets an import of the NodeJS EventEmitter.
+   */
+  private DataFlow::SourceNode getAnEventEmitterImport() {
+    result = DataFlow::moduleImport("events") or
+    result = DataFlow::moduleMember("events", "EventEmitter")
+  }
+
   /**
    * An instance of an EventEmitter that is imported through the 'events' module.
    */
   private class ImportedNodeJSEventEmitter extends NodeJSEventEmitter {
     ImportedNodeJSEventEmitter() {
-      exists(DataFlow::SourceNode clazz |
-        clazz = DataFlow::moduleImport("events") or
-        clazz = DataFlow::moduleMember("events", "EventEmitter")
-      |
-        this = clazz.getAnInstantiation()
-      )
+      this = getAnEventEmitterImport().getAnInstantiation()
+    }
+  }
+
+  /**
+   * A class that extends EventEmitter.
+   */
+  private class EventEmitterSubClass extends DataFlow::ClassNode {
+    EventEmitterSubClass() {
+      this.getASuperClassNode().getALocalSource() = getAnEventEmitterImport() or
+      this.getADirectSuperClass() instanceof EventEmitterSubClass
+    }
+
+    private DataFlow::SourceNode ref(DataFlow::TypeTracker t) {
+      t.start() and result = this
+      or
+      exists (DataFlow::TypeTracker t2 | result = ref(t2).track(t2, t))
+    }
+
+    /**
+     * Gets a reference to this class.
+     */
+    DataFlow::SourceNode ref() { result = ref(DataFlow::TypeTracker::end()) }
+  }
+
+  /**
+   * An instantiation of a class that extends EventEmitter. 
+   * 
+   * By extending `NodeJSEventEmitter' we get data-flow on the events passing through this EventEmitter.
+   */
+  private class CustomEventEmitter extends NodeJSEventEmitter {
+    CustomEventEmitter() {
+      this = any(EventEmitterSubClass clazz).ref().getAnInstantiation()
     }
   }
 

javascript/ql/test/library-tests/frameworks/EventEmitter/test.expected

@@ -4,3 +4,5 @@
 | tst.js:20:17:20:21 | "foo" | tst.js:19:16:19:18 | foo |
 | tst.js:21:17:21:21 | "bar" | tst.js:19:39:19:41 | bar |
 | tst.js:28:17:28:22 | "blab" | tst.js:25:16:25:20 | event |
+| tst.js:34:18:34:22 | "BOH" | tst.js:33:17:33:17 | x |
+| tst.js:40:20:40:27 | "yabity" | tst.js:39:19:39:19 | x |

javascript/ql/test/library-tests/frameworks/EventEmitter/tst.js

@@ -26,3 +26,15 @@ em3.on("bla", (event) => {
 	event.returnValue = "foo"
 });
 em3.emit("bla", "blab");
+
+class MyEventEmitter extends emitter {};
+
+var em4 = new MyEventEmitter();
+em4.on("blab", (x) => {});
+em4.emit("blab", "BOH");
+
+class ExtendsMyCustomEmitter extends MyEventEmitter{}
+
+var em5 = new ExtendsMyCustomEmitter();
+em5.on("yibity", (x) => {});
+em5.emit("yibity", "yabity");