fix stack register value for calls on x86

src/cwe_checker_lib/src/abstract_domain/identifier/location.rs

@@ -35,13 +35,12 @@ pub enum AbstractLocation {
 impl std::fmt::Display for AbstractLocation {
     fn fmt(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
-            Self::Register(var) => write!(formatter, "{}", var.name),
-            Self::GlobalAddress { address, size: _ } => write!(formatter, "0x{address:x}"),
-            Self::Pointer(var, location) => write!(formatter, "{}->{}", var.name, location),
-            Self::GlobalPointer(address, location) => {
-                write!(formatter, "0x{address:x}->{location}")
-            }
-        }
+            Self::Register(var) => write!(formatter, "{}", var.name)?,
+            Self::GlobalAddress { address, size: _ } => write!(formatter, "0x{address:x}")?,
+            Self::Pointer(var, location) => write!(formatter, "{}{}", var.name, location)?,
+            Self::GlobalPointer(address, location) => write!(formatter, "0x{address:x}{location}")?,
+        };
+        write!(formatter, ":i{}", self.bytesize().as_bit_length())
     }
 }
 

src/cwe_checker_lib/src/abstract_domain/identifier/mem_location.rs

@@ -128,8 +128,8 @@ impl AbstractMemoryLocation {
 impl std::fmt::Display for AbstractMemoryLocation {
     fn fmt(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
-            Self::Location { offset, .. } => write!(formatter, "({offset})"),
-            Self::Pointer { offset, target } => write!(formatter, "({offset})->{target}"),
+            Self::Location { offset, .. } => write!(formatter, "[{offset}]"),
+            Self::Pointer { offset, target } => write!(formatter, "[{offset}]{target}"),
         }
     }
 }

src/cwe_checker_lib/src/analysis/function_signature/context/mod.rs

@@ -1,7 +1,7 @@
 use super::*;
 use crate::abstract_domain::{
-    AbstractDomain, AbstractIdentifier, AbstractLocation, BitvectorDomain, DataDomain, SizedDomain,
-    TryToBitvec,
+    AbstractDomain, AbstractIdentifier, AbstractLocation, BitvectorDomain, DataDomain,
+    RegisterDomain as _, SizedDomain, TryToBitvec,
 };
 use crate::utils::arguments;
 use crate::{
@@ -305,6 +305,34 @@ impl<'a> Context<'a> {
         }
         None
     }
+
+    /// Adjust the stack register after a call to a function.
+    ///
+    /// On x86, this removes the return address from the stack
+    /// (other architectures pass the return address in a register, not on the stack).
+    /// On other architectures the stack register retains the value it had before the call.
+    /// Note that in some calling conventions the callee also clears function parameters from the stack.
+    /// We do not detect and handle these cases yet.
+    fn adjust_stack_register_on_return_from_call(
+        &self,
+        state_before_call: &State,
+        new_state: &mut State,
+    ) {
+        let stack_register = &self.project.stack_pointer_register;
+        let stack_pointer = state_before_call.get_register(stack_register);
+        match self.project.cpu_architecture.as_str() {
+            "x86" | "x86_32" | "x86_64" => {
+                let offset = Bitvector::from_u64(stack_register.size.into())
+                    .into_truncate(apint::BitWidth::from(stack_register.size))
+                    .unwrap();
+                new_state.set_register(
+                    stack_register,
+                    stack_pointer.bin_op(BinOpType::IntAdd, &offset.into()),
+                );
+            }
+            _ => new_state.set_register(stack_register, stack_pointer),
+        }
+    }
 }
 
 impl<'a> forward_interprocedural_fixpoint::Context<'a> for Context<'a> {
@@ -412,13 +440,15 @@ impl<'a> forward_interprocedural_fixpoint::Context<'a> for Context<'a> {
                         cconv,
                         &self.project.runtime_memory_image,
                     );
+                    self.adjust_stack_register_on_return_from_call(state, &mut new_state);
                     return Some(new_state);
                 }
             }
             Jmp::Call { target, .. } => {
                 if let Some(extern_symbol) = self.project.program.term.extern_symbols.get(target) {
                     self.handle_extern_symbol_call(&mut new_state, extern_symbol, &call.tid);
                     if !extern_symbol.no_return {
+                        self.adjust_stack_register_on_return_from_call(state, &mut new_state);
                         return Some(new_state);
                     }
                 } else if let Some(cconv) = self.project.get_standard_calling_convention() {
@@ -427,6 +457,7 @@ impl<'a> forward_interprocedural_fixpoint::Context<'a> for Context<'a> {
                         cconv,
                         &self.project.runtime_memory_image,
                     );
+                    self.adjust_stack_register_on_return_from_call(state, &mut new_state);
                     return Some(new_state);
                 }
             }
@@ -474,6 +505,7 @@ impl<'a> forward_interprocedural_fixpoint::Context<'a> for Context<'a> {
         for (var, value) in return_value_list {
             new_state.set_register(var, value);
         }
+        self.adjust_stack_register_on_return_from_call(state_before_call, &mut new_state);
         Some(new_state)
     }
 

src/cwe_checker_lib/src/analysis/function_signature/context/tests.rs

@@ -119,6 +119,51 @@ fn test_call_stub_handling() {
     assert_eq!(params.len(), 5);
 }
 
+#[test]
+fn test_stack_register_adjustment_after_call() {
+    let project = Project::mock_x64();
+    let graph = crate::analysis::graph::get_program_cfg(&project.program);
+    let context = Context::new(&project, &graph);
+    let mut state_before_call = State::mock_x64("mock_fn");
+    let stack_id = AbstractIdentifier::mock("mock_fn", "RSP", 8);
+    state_before_call.set_register(
+        &variable!("RSP:8"),
+        DataDomain::from_target(stack_id.clone(), bitvec!("0x-20:8").into()),
+    );
+    let call_term = Term {
+        tid: Tid::new("call_tid"),
+        term: Jmp::CallInd {
+            target: Expression::Var(variable!("R15:8")),
+            return_: Some(Tid::new("return_")),
+        },
+    };
+    // Test adjustment on extern calls
+    let state_after_call = context
+        .update_call_stub(&state_before_call, &call_term)
+        .unwrap();
+    let adjusted_sp = state_after_call.get_register(&variable!("RSP:8"));
+    assert_eq!(
+        adjusted_sp,
+        DataDomain::from_target(stack_id.clone(), bitvec!("0x-18:8").into())
+    );
+    // Test adjustment on intern calls
+    let state_before_return = State::mock_x64("callee");
+    let state_after_call = context
+        .update_return(
+            Some(&state_before_return),
+            Some(&state_before_call),
+            &call_term,
+            &call_term,
+            &None,
+        )
+        .unwrap();
+    let adjusted_sp = state_after_call.get_register(&variable!("RSP:8"));
+    assert_eq!(
+        adjusted_sp,
+        DataDomain::from_target(stack_id.clone(), bitvec!("0x-18:8").into())
+    );
+}
+
 #[test]
 fn test_get_global_mem_address() {
     let project = Project::mock_arm32();