Rust: CFG edges for break and continue with labels

Author: paldepind

rust/ql/lib/codeql/rust/controlflow/internal/Completion.qll

@@ -1,3 +1,4 @@
+private import codeql.util.Option
 private import codeql.util.Boolean
 private import codeql.rust.controlflow.ControlFlowGraph
 private import rust
@@ -7,9 +8,9 @@ private newtype TCompletion =
   TSimpleCompletion() or
   TBooleanCompletion(Boolean b) or
   TMatchCompletion(Boolean isMatch) or
-  TBreakCompletion() or
-  TContinueCompletion() or
-  TReturnCompletion()
+  TLoopCompletion(TLoopJumpType kind, TLabelType label) or
+  TReturnCompletion() or
+  TDivergeCompletion() // A completion that never reaches the successor (e.g. by panicking or spinning)
 
 /** A completion of a statement or an expression. */
 abstract class Completion extends TCompletion {
@@ -105,25 +106,44 @@ class MatchCompletion extends TMatchCompletion, ConditionalCompletion {
 }
 
 /**
- * A completion that represents a break.
+ * A completion that represents a break or a continue.
  */
-class BreakCompletion extends TBreakCompletion, Completion {
-  override BreakSuccessor getAMatchingSuccessorType() { any() }
+class LoopJumpCompletion extends TLoopCompletion, Completion {
+  override LoopJumpSuccessor getAMatchingSuccessorType() {
+    result = TLoopSuccessor(this.getKind(), this.getLabelType())
+  }
 
-  override predicate isValidForSpecific(AstNode e) { e instanceof BreakExpr }
+  final TLoopJumpType getKind() { this = TLoopCompletion(result, _) }
 
-  override string toString() { result = "break" }
-}
+  final TLabelType getLabelType() { this = TLoopCompletion(_, result) }
 
-/**
- * A completion that represents a continue.
- */
-class ContinueCompletion extends TContinueCompletion, Completion {
-  override ContinueSuccessor getAMatchingSuccessorType() { any() }
+  final predicate hasLabel() { this.getLabelType() = TLabel(_) }
+
+  final string getLabelName() { TLabel(result) = this.getLabelType() }
+
+  final predicate isContinue() { this.getKind() = TContinueJump() }
+
+  final predicate isBreak() { this.getKind() = TBreakJump() }
 
-  override predicate isValidForSpecific(AstNode e) { e instanceof ContinueExpr }
+  override predicate isValidForSpecific(AstNode e) {
+    this.isBreak() and
+    e instanceof BreakExpr and
+    (
+      not e.(BreakExpr).hasLabel() and not this.hasLabel()
+      or
+      e.(BreakExpr).getLabel().getName() = this.getLabelName()
+    )
+    or
+    this.isContinue() and
+    e instanceof ContinueExpr and
+    (
+      not e.(ContinueExpr).hasLabel() and not this.hasLabel()
+      or
+      e.(ContinueExpr).getLabel().getName() = this.getLabelName()
+    )
+  }
 
-  override string toString() { result = "continue" }
+  override string toString() { result = this.getAMatchingSuccessorType().toString() }
 }
 
 /**
@@ -145,9 +165,3 @@ predicate completionIsSimple(Completion c) { c instanceof SimpleCompletion }
 
 /** Holds if `c` is a valid completion for `n`. */
 predicate completionIsValidFor(Completion c, AstNode n) { c.isValidFor(n) }
-
-/** Holds if `c` is a completion that interacts with a loop such as `loop`, `for`, `while`. */
-predicate isLoopCompletion(Completion c) {
-  c instanceof BreakCompletion or
-  c instanceof ContinueCompletion
-}

rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll

@@ -303,24 +303,37 @@ class LoopExprTree extends PostOrderTree instanceof LoopExpr {
 
   override predicate first(AstNode node) { first(super.getBody(), node) }
 
+  /** Whether this `LoopExpr` captures a completion for a `break`/`continue`. */
+  predicate capturesLoopJumpCompletion(LoopJumpCompletion c) {
+    not c.hasLabel()
+    or
+    c.getLabelName() = super.getLabel().getName()
+  }
+
   override predicate succ(AstNode pred, AstNode succ, Completion c) {
     // Edge back to the start for final expression and continue expressions
     last(super.getBody(), pred, c) and
-    (completionIsNormal(c) or c instanceof ContinueCompletion) and
+    (
+      completionIsNormal(c)
+      or
+      c.(LoopJumpCompletion).isContinue() and this.capturesLoopJumpCompletion(c)
+    ) and
     this.first(succ)
     or
     // Edge for exiting the loop with a break expressions
     last(super.getBody(), pred, c) and
-    c instanceof BreakCompletion and
+    c.(LoopJumpCompletion).isBreak() and
+    this.capturesLoopJumpCompletion(c) and
     succ = this
   }
 
   override predicate last(AstNode last, Completion c) {
     super.last(last, c)
     or
+    // Any abnormal completions that this loop does not capture should propagate
     last(super.getBody(), last, c) and
     not completionIsNormal(c) and
-    not isLoopCompletion(c)
+    not this.capturesLoopJumpCompletion(c)
   }
 }
 

rust/ql/lib/codeql/rust/controlflow/internal/SuccessorType.qll

@@ -1,14 +1,25 @@
+private import rust
+private import codeql.rust.generated.Raw
 private import codeql.util.Boolean
+private import codeql.util.Option
+
+newtype TLoopJumpType =
+  TContinueJump() or
+  TBreakJump()
+
+newtype TLabelType =
+  TLabel(string s) { any(Label l).getName() = s } or
+  TNoLabel()
 
 cached
 newtype TSuccessorType =
   TSuccessorSuccessor() or
   TBooleanSuccessor(Boolean b) or
   TMatchSuccessor(Boolean b) or
-  TBreakSuccessor() or
-  TContinueSuccessor() or
+  TLoopSuccessor(TLoopJumpType kind, TLabelType label) or
   TReturnSuccessor()
 
+// class TBreakSuccessor = TUnlabeledBreakSuccessor or TLabeledBreakSuccessor;
 /** The type of a control flow successor. */
 abstract private class SuccessorTypeImpl extends TSuccessorType {
   /** Gets a textual representation of successor type. */
@@ -51,14 +62,29 @@ final class MatchSuccessor extends ConditionalSuccessor, TMatchSuccessor {
   }
 }
 
-/** A `break` control flow successor. */
-final class BreakSuccessor extends SuccessorTypeImpl, TBreakSuccessor {
-  final override string toString() { result = "break" }
-}
+/**
+ * A control flow successor of a loop control flow expression, `continue` or `break`.
+ */
+final class LoopJumpSuccessor extends SuccessorTypeImpl, TLoopSuccessor {
+  final private TLoopJumpType getKind() { this = TLoopSuccessor(result, _) }
+
+  final private TLabelType getLabelType() { this = TLoopSuccessor(_, result) }
 
-/** A `continue` control flow successor. */
-final class ContinueSuccessor extends SuccessorTypeImpl, TContinueSuccessor {
-  final override string toString() { result = "continue" }
+  final predicate hasLabel() { this.getLabelType() = TLabel(_) }
+
+  final string getLabelName() { this = TLoopSuccessor(_, TLabel(result)) }
+
+  final predicate isContinue() { this.getKind() = TContinueJump() }
+
+  final predicate isBreak() { this.getKind() = TBreakJump() }
+
+  final override string toString() {
+    exists(string kind, string label |
+      (if this.isContinue() then kind = "continue" else kind = "break") and
+      (if this.hasLabel() then label = "(" + this.getLabelName() + ")" else label = "") and
+      result = kind + label
+    )
+  }
 }
 
 /**