fixed bug of stack_alignment_substitution analysis (#370)

src/cwe_checker_lib/src/analysis/stack_alignment_substitution/mod.rs

@@ -10,6 +10,8 @@
 //! - the argument for the AND operation is not a constant
 //! - an operation alters the stack pointer, which can not be journaled.
 
+use std::collections::HashSet;
+
 use anyhow::{anyhow, Result};
 use apint::ApInt;
 use itertools::Itertools;
@@ -86,6 +88,59 @@ fn journal_sp_value(
     }
 }
 
+/// Returns the tid of the target of the first Jmp::Branch of the provided block.
+fn get_first_branch_tid(blk: &Term<Blk>) -> Option<&Tid> {
+    if let Some(jmp) = blk.term.jmps.get(0) {
+        if let Jmp::Branch(jump_to_blk) = &jmp.term {
+            return Some(jump_to_blk);
+        }
+    }
+    None
+}
+
+/// Returns the index of the first block with non-empty defs.
+/// Blocks are iterated according by considering their first `Jmp::Branch`.
+/// If a block is revisited, `None` is returned.
+fn get_first_blk_with_defs(sub: &Sub) -> Option<usize> {
+    let blocks = &sub.blocks;
+    if let Some(start_blk) = blocks.first() {
+        let mut visited = HashSet::new();
+        let mut blk = start_blk;
+
+        'search_loop: while blk.term.defs.is_empty() {
+            if let Some(target_tid) = get_first_branch_tid(blk) {
+                if !visited.contains(&blk.tid) {
+                    visited.insert(&blk.tid);
+
+                    // try find this target
+                    for (index, target_blk) in blocks.iter().enumerate() {
+                        if &target_blk.tid == target_tid {
+                            if !target_blk.term.defs.is_empty() {
+                                return Some(index);
+                            } else {
+                                // continue with new block
+                                blk = target_blk;
+                                continue 'search_loop;
+                            }
+                        }
+                    }
+                    // did not find target
+                    return None;
+                } else {
+                    // busy loop
+                    return None;
+                }
+            } else {
+                // did not find branch in block
+                return None;
+            }
+        }
+        // first block was not empty
+        return Some(0);
+    }
+    None
+}
+
 /// Substitutes logical AND on the stackpointer register by SUB.
 /// Expressions are changed to use constants w.r.t the provided bit mask.
 pub fn substitute_and_on_stackpointer(project: &mut Project) -> Option<Vec<LogMessage>> {
@@ -101,8 +156,8 @@ pub fn substitute_and_on_stackpointer(project: &mut Project) -> Option<Vec<LogMe
 
     'sub_loop: for sub in project.program.term.subs.values_mut() {
         let journaled_sp: &mut i64 = &mut 0;
-        // only for the first block SP can be reasonable tracked
-        if let Some(blk) = sub.term.blocks.first_mut() {
+        if let Some(index) = get_first_blk_with_defs(&sub.term) {
+            let blk = &mut sub.term.blocks[index];
             for def in blk.term.defs.iter_mut() {
                 if let Def::Assign { var, value } = &mut def.term {
                     if *var == project.stack_pointer_register {

src/cwe_checker_lib/src/analysis/stack_alignment_substitution/tests.rs

@@ -382,3 +382,112 @@ fn supports_commutative_and() {
         }
     }
 }
+#[test]
+/// Some functions have leading blocks without any defs. This might be due to `endbr`-like instructions.
+/// We skip those empty blocks and start substituting for rhe first non-empty block.
+fn skips_empty_blocks() {
+    let sub_from_sp = Def::assign(
+        "tid_alter_sp",
+        Project::mock_x64().stack_pointer_register.clone(),
+        Expression::minus(
+            Expression::Var(Project::mock_x64().stack_pointer_register.clone()),
+            Expression::const_from_apint(ApInt::from_u64(1)),
+        ),
+    );
+
+    let byte_alignment_as_and = Def::assign(
+        "tid_to_be_substituted",
+        Project::mock_x64().stack_pointer_register.clone(),
+        Expression::BinOp {
+            op: BinOpType::IntAnd,
+            lhs: Box::new(Expression::Var(
+                Project::mock_x64().stack_pointer_register.clone(),
+            )),
+            rhs: Box::new(Expression::const_from_apint(ApInt::from_u64(
+                0xFFFFFFFF_FFFFFFFF << 4, // 16 Byte alignment
+            ))),
+        },
+    );
+    // get project with empty block
+    let mut proj = setup(vec![], true);
+    // add jmp
+    proj.program
+        .term
+        .subs
+        .get_mut(&Tid::new("sub_tid"))
+        .unwrap()
+        .term
+        .blocks[0]
+        .term
+        .jmps
+        .push(Term {
+            tid: Tid::new("tid"),
+            term: Jmp::Branch(Tid::new("not_empty_blk")),
+        });
+
+    let mut blk = Blk::mock_with_tid("not_empty_blk");
+    blk.term.defs.push(sub_from_sp.clone());
+    blk.term.defs.push(byte_alignment_as_and.clone());
+
+    // add block with substitutional def
+    proj.program
+        .term
+        .subs
+        .get_mut(&Tid::new("sub_tid"))
+        .unwrap()
+        .term
+        .blocks
+        .push(blk);
+
+    let expected_def = Def::assign(
+        "tid_to_be_substituted",
+        Project::mock_x64().stack_pointer_register.clone(),
+        Expression::minus(
+            Expression::Var(Project::mock_x64().stack_pointer_register.clone()),
+            Expression::const_from_apint(ApInt::from_u64(15)),
+        ),
+    );
+
+    substitute_and_on_stackpointer(&mut proj);
+
+    assert_eq!(
+        proj.program
+            .term
+            .subs
+            .get(&Tid::new("sub_tid"))
+            .unwrap()
+            .term
+            .blocks[1]
+            .term
+            .defs,
+        vec![sub_from_sp.clone(), expected_def]
+    );
+}
+
+#[test]
+fn skip_busy_loop() {
+    let mut proj = setup(vec![], true);
+    proj.program
+        .term
+        .subs
+        .get_mut(&Tid::new("sub_tid"))
+        .unwrap()
+        .term
+        .blocks[0]
+        .term
+        .jmps
+        .push(Jmp::branch("jmp_to_1st_blk", "block"));
+
+    assert_eq!(
+        get_first_blk_with_defs(
+            &proj
+                .program
+                .term
+                .subs
+                .get_mut(&Tid::new("sub_tid"))
+                .unwrap()
+                .term
+        ),
+        None
+    );
+}