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()
     }
   }