fix: wrong code generation

VirusTotal · Nov 22, 2023 · 9a1e350 · 9a1e350
1 parent f5f0602
commit 9a1e350
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 37 deletions.
diff --git a/yara-x/src/compiler/emit.rs b/yara-x/src/compiler/emit.rs
@@ -25,7 +25,7 @@ use crate::compiler::ir::{
 use crate::compiler::{LiteralId, RegexpId, RuleId, Var, VarStackFrame};
 use crate::string_pool::{BStringPool, StringPool};
 use crate::symbols::SymbolKind;
-use crate::types::{Array, Map, Type, TypeValue, Value};
+use crate::types::{Array, Func, Map, Type, TypeValue, Value};
 use crate::utils::cast;
 use crate::wasm;
 use crate::wasm::builder::WasmModuleBuilder;
@@ -326,11 +326,13 @@ fn emit_expr(
         Expr::Filesize { .. } => {
             instr.global_get(ctx.wasm_symbols.filesize);
         }
+
         Expr::Entrypoint { .. } => {
             // TODO
             // todo!()
             instr.i64_const(0);
         }
+
         Expr::Ident { symbol } => {
             match symbol.kind() {
                 SymbolKind::Rule(rule_id) => {
@@ -349,43 +351,45 @@ fn emit_expr(
                     ctx.lookup_start = Some(*var);
                 }
                 SymbolKind::Func(func) => {
-                    let signature =
-                        &func.signatures()[ctx.current_signature.unwrap()];
-
-                    if signature.result_may_be_undef {
-                        emit_call_and_handle_undef(
-                            ctx,
-                            instr,
-                            ctx.function_id(signature.mangled_name.as_str()),
-                        );
-                    } else {
-                        instr.call(
-                            ctx.function_id(signature.mangled_name.as_str()),
-                        );
-                    }
+                    emit_func_call(ctx, instr, func);
                 }
                 SymbolKind::FieldIndex(index) => {
                     ctx.lookup_stack.push_back((*index).try_into().unwrap());
 
-                    match symbol.type_value().ty() {
-                        Type::Integer => {
+                    match symbol.type_value() {
+                        TypeValue::Integer(_) => {
                             emit_lookup_integer(ctx, instr);
+                            assert!(ctx.lookup_stack.is_empty());
                         }
-                        Type::Float => {
+                        TypeValue::Float(_) => {
                             emit_lookup_float(ctx, instr);
+                            assert!(ctx.lookup_stack.is_empty());
                         }
-                        Type::Bool => {
+                        TypeValue::Bool(_) => {
                             emit_lookup_bool(ctx, instr);
+                            assert!(ctx.lookup_stack.is_empty());
                         }
-                        Type::String => {
+                        TypeValue::String(_) => {
                             emit_lookup_string(ctx, instr);
+                            assert!(ctx.lookup_stack.is_empty());
                         }
-                        Type::Struct | Type::Array | Type::Map => {
+                        TypeValue::Struct(_)
+                        | TypeValue::Array(_)
+                        | TypeValue::Map(_) => {
                             // Do nothing. For structs, arrays and maps pushing
                             // the field index in `lookup_stack` is enough. We
                             // don't need to emit a call for retrieving a value.
                         }
-                        _ => {
+                        TypeValue::Func(func) => {
+                            emit_func_call(ctx, instr, func);
+                            ctx.lookup_stack.clear();
+                        }
+                        TypeValue::Regexp(_) => {
+                            // The value of an identifier can't be a regular
+                            // expression.
+                            unreachable!();
+                        }
+                        TypeValue::Unknown => {
                             // This point should not be reached. The type of
                             // identifiers must be known during code emitting
                             // because they are resolved during the semantic
@@ -2396,6 +2400,24 @@ fn emit_bool_expr(
     }
 }
 
+/// Emit function call.
+fn emit_func_call(
+    ctx: &EmitContext,
+    instr: &mut InstrSeqBuilder,
+    func: &Func,
+) {
+    let signature = &func.signatures()[ctx.current_signature.unwrap()];
+    if signature.result_may_be_undef {
+        emit_call_and_handle_undef(
+            ctx,
+            instr,
+            ctx.function_id(signature.mangled_name.as_str()),
+        );
+    } else {
+        instr.call(ctx.function_id(signature.mangled_name.as_str()));
+    }
+}
+
 /// Calls a function that may return an undefined value.
 ///
 /// Some functions in YARA can return undefined values, for example the

diff --git a/yara-x/src/compiler/tests/mod.rs b/yara-x/src/compiler/tests/mod.rs
@@ -497,3 +497,26 @@ fn issue() {
     let mut scanner = Scanner::new(&rules);
     let _ = scanner.scan(b"foobar").unwrap();
 }
+
+#[test]
+fn issue2() {
+    let mut compiler = Compiler::new();
+
+    compiler
+        .add_source(
+            r#"
+            import "hash"
+            import "pe"
+            rule winti_blackfly_rich_header
+            {
+             condition:
+                hash.md5(pe.rich_signature.clear_data) =="6e89f2fff33a5c1cd4c94c13f54f9e2c" or
+                hash.md5(pe.rich_signature.clear_data) =="52667216065ac5c098808eedfec7d70b"
+            }"#,
+        )
+        .unwrap();
+
+    let rules = compiler.build();
+    let mut scanner = Scanner::new(&rules);
+    let _ = scanner.scan(b"foobar").unwrap();
+}
diff --git a/yara-x/src/symbols/mod.rs b/yara-x/src/symbols/mod.rs
@@ -105,21 +105,10 @@ impl SymbolLookup for Option<Symbol> {
 
 impl SymbolLookup for Struct {
     fn lookup(&self, ident: &str) -> Option<Symbol> {
-        let field = self.field_by_name(ident)?;
-
-        let symbol = if let TypeValue::Func(func) = &field.type_value {
-            Symbol::new(
-                field.type_value.clone(),
-                SymbolKind::Func(func.clone()),
-            )
-        } else {
-            Symbol::new(
-                field.type_value.clone(),
-                SymbolKind::FieldIndex(self.index_of(ident)),
-            )
-        };
-
-        Some(symbol)
+        Some(Symbol::new(
+            self.field_by_name(ident)?.type_value.clone(),
+            SymbolKind::FieldIndex(self.index_of(ident)),
+        ))
     }
 }