From 05460f6444e165ac87324f9441a26a23a5632d76 Mon Sep 17 00:00:00 2001
From: Anna Railton <annarailton@github.com>
Date: Thu, 18 Nov 2021 11:58:31 +0000
Subject: [PATCH] Add check for neighborhood node being top-level enclosing
 function

---
 .../adaptivethreatmodeling/CodeToFeatures.qll      |  2 ++
 .../adaptivethreatmodeling/EndpointFeatures.qll    | 14 +++++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll
index 9d8b1632c32..dfdfbed3562 100644
--- a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll
+++ b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll
@@ -49,6 +49,8 @@ module Raw {
 
     AstNode getParentNode() { result = TAstNode(rawNode.getParent()) }
 
+    raw::ASTNode getNode() { result = rawNode }
+
     raw::StmtContainer getContainer() { result = rawNode.getContainer() }
 
     /**
diff --git a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll
index f2ff59db46d..0bb2cf66940 100644
--- a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll
+++ b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll
@@ -184,7 +184,10 @@ module NeighborhoodBodies {
    * in the subtree.
    */
   Raw::AstNode getNeighborhoodAstNode(Raw::AstNode node) {
-    if getNumDescendents(node.getParentNode()) > maxNumDescendants()
+    if
+      // `node` will always have a parent as we start at and endpoint
+      node.getParentNode() = getOutermostEnclosingFunction(node) or
+      getNumDescendents(node.getParentNode()) > maxNumDescendants()
     then result = node
     else result = getNeighborhoodAstNode(node.getParentNode())
   }
@@ -192,6 +195,15 @@ module NeighborhoodBodies {
   /** Count number of descendants of an AST node */
   int getNumDescendents(Raw::AstNode node) { result = count(node.getAChildNode*()) }
 
+  private ASTNode getContainer(ASTNode node) {
+    result = node.getContainer()
+  }
+
+  /** Return the AST node that is outermost enclosing function (as an AST Node) */
+  Raw::AstNode getOutermostEnclosingFunction(Raw::AstNode node) {
+    result = Raw::astNode(getContainer*(node.getNode())) and result.getContainer() instanceof TopLevel
+  }
+
   /**
    * Holds if `childNode` is an AST node under `rootNode` and `token` is a node attribute associated
    * with `childNode`. Note that only AST leaves have node attributes.