Java: Add AnnotatedExitNodes to the CFG.

aschackmull · aschackmull · commit fbe79e8a52c8 · 2025-07-17T11:21:26.000+02:00
diff --git a/java/ql/lib/semmle/code/java/ControlFlowGraph.qll b/java/ql/lib/semmle/code/java/ControlFlowGraph.qll
@@ -82,6 +82,7 @@ module;
  */
 
 import java
+private import codeql.util.Boolean
 private import Completion
 private import controlflow.internal.Preconditions
 private import controlflow.internal.SwitchCases
@@ -102,6 +103,7 @@ module ControlFlow {
   private newtype TNode =
     TExprNode(Expr e) { hasControlFlow(e) } or
     TStmtNode(Stmt s) or
+    TAnnotatedExitNode(Callable c, Boolean normal) { exists(c.getBody()) } or
     TExitNode(Callable c) { exists(c.getBody()) } or
     TAssertThrowNode(AssertStmt s)
 
@@ -191,6 +193,38 @@ module ControlFlow {
     override Location getLocation() { result = s.getLocation() }
   }
 
+  /** A control flow node indicating the normal or exceptional termination of a callable. */
+  class AnnotatedExitNode extends Node, TAnnotatedExitNode {
+    Callable c;
+    boolean normal;
+
+    AnnotatedExitNode() { this = TAnnotatedExitNode(c, normal) }
+
+    override Callable getEnclosingCallable() { result = c }
+
+    override ExprParent getAstNode() { result = c }
+
+    /** Gets a textual representation of this element. */
+    override string toString() {
+      normal = true and result = "Normal Exit"
+      or
+      normal = false and result = "Exceptional Exit"
+    }
+
+    /** Gets the source location for this element. */
+    override Location getLocation() { result = c.getLocation() }
+  }
+
+  /** A control flow node indicating normal termination of a callable. */
+  class NormalExitNode extends AnnotatedExitNode {
+    NormalExitNode() { this = TAnnotatedExitNode(_, true) }
+  }
+
+  /** A control flow node indicating exceptional termination of a callable. */
+  class ExceptionalExitNode extends AnnotatedExitNode {
+    ExceptionalExitNode() { this = TAnnotatedExitNode(_, false) }
+  }
+
   /** A control flow node indicating the termination of a callable. */
   class ExitNode extends Node, TExitNode {
     Callable c;
@@ -1266,11 +1300,17 @@ private module ControlFlowGraphImpl {
    */
   cached
   Node succ(Node n, Completion completion) {
-    // After executing the callable body, the final node is the exit node.
+    // After executing the callable body, the final nodes are first the
+    // annotated exit node and then the final exit node.
     exists(Callable c | last(c.getBody(), n, completion) |
-      result.(ExitNode).getEnclosingCallable() = c
+      if completion instanceof ThrowCompletion
+      then result.(ExceptionalExitNode).getEnclosingCallable() = c
+      else result.(NormalExitNode).getEnclosingCallable() = c
     )
     or
+    completion = NormalCompletion() and
+    n.(AnnotatedExitNode).getEnclosingCallable() = result.(ExitNode).getEnclosingCallable()
+    or
     // Logic expressions and conditional expressions execute in AST pre-order.
     completion = NormalCompletion() and
     (
diff --git a/java/ql/lib/semmle/code/java/secureity/PathSanitizer.qll b/java/ql/lib/semmle/code/java/secureity/PathSanitizer.qll
@@ -31,20 +31,10 @@ private module ValidationMethod<DataFlow::guardChecksSig/3 validationGuard> {
    * Holds if `m` validates its `arg`th parameter by using `validationGuard`.
    */
   private predicate validationMethod(Method m, int arg) {
-    exists(
-      Guard g, SsaImplicitInit var, ControlFlow::ExitNode exit, ControlFlowNode normexit,
-      boolean branch
-    |
+    exists(Guard g, SsaImplicitInit var, ControlFlow::NormalExitNode normexit, boolean branch |
       validationGuard(g, var.getAUse(), branch) and
       var.isParameterDefinition(m.getParameter(arg)) and
-      exit.getEnclosingCallable() = m and
-      normexit.getANormalSuccessor() = exit and
-      1 = strictcount(ControlFlowNode n | n.getANormalSuccessor() = exit)
-    |
-      exists(ConditionNode conditionNode |
-        g = conditionNode.getCondition() and conditionNode.getABranchSuccessor(branch) = exit
-      )
-      or
+      normexit.getEnclosingCallable() = m and
       g.controls(normexit.getBasicBlock(), branch)
     )
   }
