From b0f38ccad3a1c35385e2b6d3dee3d2a3ad53430d Mon Sep 17 00:00:00 2001
From: "Victor M. Alvarez" <vmalvarez@virustotal.com>
Date: Tue, 1 Aug 2023 23:05:01 +0200
Subject: [PATCH] fix: `for .. of ..` loops not working correctly.

---
 yara-x/src/compiler/emit.rs      |  3 +--
 yara-x/src/compiler/ir/ast2ir.rs | 11 ++++++++---
 yara-x/src/compiler/ir/mod.rs    |  1 +
 yara-x/src/tests/mod.rs          | 26 ++++++++++++++++++++++++++
 4 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/yara-x/src/compiler/emit.rs b/yara-x/src/compiler/emit.rs
index fd1c0c14d..176167557 100644
--- a/yara-x/src/compiler/emit.rs
+++ b/yara-x/src/compiler/emit.rs
@@ -1310,7 +1310,6 @@ fn emit_of_expr_tuple(
     of: &mut Of,
 ) {
     let expressions = cast!(&mut of.items, OfItems::BoolExprTuple);
-
     let next_item = of.stack_frame.new_var(Type::Bool);
     let num_expressions = expressions.len();
     let mut expressions = expressions.iter_mut();
@@ -1357,7 +1356,7 @@ fn emit_for_of_pattern_set(
 ) {
     let num_patterns = for_of.pattern_set.len();
     let mut pattern_ids = for_of.pattern_set.iter();
-    let next_pattern_id = for_of.stack_frame.new_var(Type::Integer);
+    let next_pattern_id = for_of.variable;
 
     emit_for(
         ctx,
diff --git a/yara-x/src/compiler/ir/ast2ir.rs b/yara-x/src/compiler/ir/ast2ir.rs
index f64c0b18c..998ac0f73 100644
--- a/yara-x/src/compiler/ir/ast2ir.rs
+++ b/yara-x/src/compiler/ir/ast2ir.rs
@@ -561,7 +561,9 @@ fn of_expr_from_ast(
     of: &ast::Of,
 ) -> Result<Expr, CompileError> {
     let quantifier = quantifier_from_ast(ctx, &of.quantifier)?;
-
+    // Create new stack frame with 5 slots:
+    //   1 slot for the loop variable, a bool in this case.
+    //   4 up to slots used for loop control variables (see: emit::emit_for)
     let stack_frame = ctx.vars.new_frame(5);
 
     let (items, num_items) = match &of.items {
@@ -665,8 +667,10 @@ fn for_of_expr_from_ast(
 ) -> Result<Expr, CompileError> {
     let quantifier = quantifier_from_ast(ctx, &for_of.quantifier)?;
     let pattern_set = pattern_set_from_ast(ctx, &for_of.pattern_set);
-
-    let mut stack_frame = ctx.vars.new_frame(4);
+    // Create new stack frame with 5 slots:
+    //   1 slot for the loop variable, a pattern ID in this case
+    //   4 up to slots used for loop control variables (see: emit::emit_for)
+    let mut stack_frame = ctx.vars.new_frame(5);
     let next_pattern_id = stack_frame.new_var(Type::Integer);
     let mut loop_vars = SymbolTable::new();
 
@@ -692,6 +696,7 @@ fn for_of_expr_from_ast(
         pattern_set,
         condition,
         stack_frame,
+        variable: next_pattern_id,
     })))
 }
 
diff --git a/yara-x/src/compiler/ir/mod.rs b/yara-x/src/compiler/ir/mod.rs
index 8bfb616ec..7c14f70eb 100644
--- a/yara-x/src/compiler/ir/mod.rs
+++ b/yara-x/src/compiler/ir/mod.rs
@@ -411,6 +411,7 @@ pub(in crate::compiler) struct Of {
 /// `for 1 of ($a,$b) : (..)`)
 pub(in crate::compiler) struct ForOf {
     pub quantifier: Quantifier,
+    pub variable: Var,
     pub pattern_set: Vec<PatternId>,
     pub condition: Expr,
     pub stack_frame: VarStackFrame,
diff --git a/yara-x/src/tests/mod.rs b/yara-x/src/tests/mod.rs
index d9ab83880..f0ff36a7b 100644
--- a/yara-x/src/tests/mod.rs
+++ b/yara-x/src/tests/mod.rs
@@ -2100,6 +2100,32 @@ fn for_of() {
         "#,
         b"foobar"
     );
+
+    rule_true!(
+        r#"
+        rule test {
+          strings:
+            $a = "foo"
+            $b = "bar"
+          condition:
+            for 1 of them : ( # == 2 )
+        }
+        "#,
+        b"foobarbar"
+    );
+
+    rule_true!(
+        r#"
+        rule test {
+          strings:
+            $a = "foo"
+            $b = "bar"
+          condition:
+            for 1 of them : ( @ > 0 )
+        }
+        "#,
+        b"foobarbar"
+    );
 }
 
 #[test]