From 2f238280dcfdd10c05ae9bfc109793f5ff08cf78 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Mon, 30 Nov 2020 16:06:13 +0100
Subject: [PATCH] CFG: model if-modifier and unless

---
 .../controlflow/internal/AstNodes.qll         | 23 ++++++++++++++++++-
 .../internal/ControlFlowGraphImpl.qll         |  4 ++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/ql/src/codeql_ruby/controlflow/internal/AstNodes.qll b/ql/src/codeql_ruby/controlflow/internal/AstNodes.qll
index 937d676a739..6451ee330ec 100644
--- a/ql/src/codeql_ruby/controlflow/internal/AstNodes.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/AstNodes.qll
@@ -41,7 +41,8 @@ class LogicalOrAstNode extends Binary {
   AstNode getAnOperand() { result in [left, right] }
 }
 
-private class If_or_elisif = @if or @elsif or @conditional;
+private class If_or_elisif =
+  @if or @elsif or @conditional or @if_modifier or @unless or @unless_modifier;
 
 class IfElsifAstNode extends AstNode, If_or_elisif {
   AstNode getConditionNode() { none() }
@@ -75,6 +76,26 @@ private class ConditionalAstNode extends IfElsifAstNode, Conditional {
   override AstNode getAlternativeNode() { result = this.getAlternative() }
 }
 
+private class IfModifierAstNode extends IfElsifAstNode, IfModifier {
+  override AstNode getConditionNode() { result = this.getCondition() }
+
+  override AstNode getConsequenceNode() { result = this.getBody() }
+}
+
+private class UnlessAstNode extends IfElsifAstNode, Unless {
+  override AstNode getConditionNode() { result = this.getCondition() }
+
+  override AstNode getConsequenceNode() { result = this.getAlternative() }
+
+  override AstNode getAlternativeNode() { result = this.getConsequence() }
+}
+
+private class UnlessModifierAstNode extends IfElsifAstNode, UnlessModifier {
+  override AstNode getConditionNode() { result = this.getCondition() }
+
+  override AstNode getAlternativeNode() { result = this.getBody() }
+}
+
 private class CondLoop = @while or @while_modifier or @until or @until_modifier;
 
 class ConditionalLoopAstNode extends AstNode, CondLoop {
diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index a9db11ec8d0..a842532eeb1 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -405,6 +405,10 @@ private module Trees {
       c instanceof FalseCompletion and
       not exists(this.getAlternativeNode())
       or
+      last(this.getConditionNode(), last, c) and
+      c instanceof TrueCompletion and
+      not exists(this.getConsequenceNode())
+      or
       last(this.getConsequenceNode(), last, c)
       or
       last(this.getAlternativeNode(), last, c)