fix(transforms): handle immutable variables and inherited contract patterns

crates/core/src/cfg_ir/mod.rs

@@ -751,6 +751,216 @@ impl CfgIrBundle {
         Ok(())
     }
 
+    /// Remap return-address PUSH instructions in Solidity internal function call patterns.
+    ///
+    /// After `reindex_pcs` shifts PCs, `write_symbolic_immediates` and `patch_jump_immediates`
+    /// update PUSH values that feed directly into JUMP/JUMPI. However, return addresses pushed
+    /// earlier in a block (the `PUSH ret_addr` in `PUSH ret_addr; PUSH func_entry; JUMP`) are
+    /// not part of any recognized jump pattern and become stale. This pass finds those specific
+    /// return-address PUSHes and remaps them.
+    pub fn remap_orphan_jump_pushes(
+        &mut self,
+        pc_mapping: &HashMap<usize, usize>,
+        old_runtime_bounds: Option<(usize, usize)>,
+    ) -> Result<(), Error> {
+        let old_runtime_start = old_runtime_bounds.map(|(s, _)| s);
+        let new_runtime_start = self.runtime_bounds.map(|(s, _)| s);
+
+        // Build set of old JUMPDEST PCs so we can verify candidates are real jump targets.
+        let inverse: HashMap<usize, usize> =
+            pc_mapping.iter().map(|(&old, &new)| (new, old)).collect();
+        let mut old_jumpdest_pcs: HashSet<usize> = HashSet::new();
+        for node in self.cfg.node_indices() {
+            if let Some(Block::Body(body)) = self.cfg.node_weight(node) {
+                for instr in &body.instructions {
+                    if matches!(instr.op, Opcode::JUMPDEST)
+                        && let Some(&old_pc) = inverse.get(&instr.pc)
+                    {
+                        old_jumpdest_pcs.insert(old_pc);
+                    }
+                }
+            }
+        }
+
+        if old_jumpdest_pcs.is_empty() {
+            return Ok(());
+        }
+
+        // First pass: collect (node, instruction_index) pairs that need remapping.
+        // We look for the internal call pattern: PUSH ret_addr; PUSH func_entry; JUMP
+        // The return address PUSH is at push_idx - 1 in a Direct pattern.
+        let nodes: Vec<_> = self.cfg.node_indices().collect();
+        let mut edits: Vec<(NodeIndex, usize, usize)> = Vec::new(); // (node, instr_idx, new_value)
+
+        for &node in &nodes {
+            let Some(Block::Body(body)) = self.cfg.node_weight(node) else {
+                continue;
+            };
+
+            let in_runtime = body.is_runtime(self.runtime_bounds);
+
+            // compute the remapped value for a PUSH instruction
+            let try_remap = |push_value: usize| -> Option<usize> {
+                let old_pc_abs = if in_runtime {
+                    old_runtime_start.unwrap_or(0).saturating_add(push_value)
+                } else {
+                    push_value
+                };
+                if !old_jumpdest_pcs.contains(&old_pc_abs) {
+                    return None;
+                }
+                let &new_pc_abs = pc_mapping.get(&old_pc_abs)?;
+                let new_value = if in_runtime {
+                    new_runtime_start
+                        .map(|s| new_pc_abs.saturating_sub(s))
+                        .unwrap_or(new_pc_abs)
+                } else {
+                    new_pc_abs
+                };
+                if push_value != new_value {
+                    Some(new_value)
+                } else {
+                    None
+                }
+            };
+
+            // Find the terminal jump pattern
+            let pattern = detect_jump_pattern(&body.instructions);
+
+            if let Some(ref pat) = pattern {
+                // check the PUSH immediately before the jump pattern
+                let pattern_first_push_idx = match pat {
+                    JumpPattern::Direct { push_idx } => *push_idx,
+                    JumpPattern::SplitAdd { push_a_idx, .. } => *push_a_idx,
+                    JumpPattern::PcRelative { push_idx, .. } => *push_idx,
+                };
+
+                if pattern_first_push_idx > 0 {
+                    let ret_idx = pattern_first_push_idx - 1;
+                    let ret_instr = &body.instructions[ret_idx];
+                    if matches!(ret_instr.op, Opcode::PUSH(_))
+                        && let Some(imm) = &ret_instr.imm
+                        && let Ok(push_value) = usize::from_str_radix(imm, 16)
+                        && let Some(new_value) = try_remap(push_value)
+                    {
+                        let old_pc_abs = if in_runtime {
+                            old_runtime_start.unwrap_or(0).saturating_add(push_value)
+                        } else {
+                            push_value
+                        };
+                        let new_pc_abs = pc_mapping.get(&old_pc_abs).copied().unwrap_or(0);
+                        tracing::debug!(
+                            "remap_orphan_jump_pushes: block {} instr {} at pc=0x{:x}: \
+                             0x{:x} -> 0x{:x} (abs: 0x{:x} -> 0x{:x})",
+                            node.index(),
+                            ret_idx,
+                            ret_instr.pc,
+                            push_value,
+                            new_value,
+                            old_pc_abs,
+                            new_pc_abs,
+                        );
+                        edits.push((node, ret_idx, new_value));
+                    }
+                }
+            }
+
+            // Extended scan: for blocks ending with JUMP/JUMPI (regardless of pattern),
+            // scan ALL PUSH instructions for values matching old JUMPDEST PCs.
+            //
+            // Solidity contracts with inheritance (e.g. EscrowERC20 + EscrowBase) emit
+            // internal function call patterns where the return address PUSH is separated
+            // from the terminal JUMP by several instructions:
+            //
+            //   PUSH ret_addr   ← return address, not adjacent to JUMP
+            //   DUP3
+            //   PUSH2 value
+            //   SWAP4
+            //   PUSH0
+            //   SSTORE
+            //   PUSH1 slot
+            //   SSTORE
+            //   JUMP            ← uses ret_addr still on the stack
+            //
+            // The standard path above only checks the PUSH immediately before a recognized
+            // jump pattern (PUSH+JUMP). This extended scan catches return addresses at
+            // arbitrary positions within the block.
+            let last = body.instructions.last();
+            let ends_with_jump = last.is_some_and(|i| matches!(i.op, Opcode::JUMP | Opcode::JUMPI));
+            if ends_with_jump {
+                // Determine which indices are already part of the recognized jump pattern
+                // to avoid double-remapping
+                let pattern_indices: HashSet<usize> = match &pattern {
+                    Some(JumpPattern::Direct { push_idx }) => {
+                        [*push_idx, push_idx.wrapping_sub(1)].into_iter().collect()
+                    }
+                    Some(JumpPattern::SplitAdd {
+                        push_a_idx,
+                        push_b_idx,
+                    }) => [*push_a_idx, *push_b_idx, push_a_idx.wrapping_sub(1)]
+                        .into_iter()
+                        .collect(),
+                    Some(JumpPattern::PcRelative { push_idx, .. }) => {
+                        [*push_idx, push_idx.wrapping_sub(1)].into_iter().collect()
+                    }
+                    None => HashSet::new(),
+                };
+
+                for (idx, instr) in body.instructions.iter().enumerate() {
+                    if pattern_indices.contains(&idx) {
+                        continue;
+                    }
+                    if !matches!(instr.op, Opcode::PUSH(_)) {
+                        continue;
+                    }
+                    let Some(imm) = &instr.imm else {
+                        continue;
+                    };
+                    let Ok(push_value) = usize::from_str_radix(imm, 16) else {
+                        continue;
+                    };
+                    if let Some(new_value) = try_remap(push_value) {
+                        let old_pc_abs = if in_runtime {
+                            old_runtime_start.unwrap_or(0).saturating_add(push_value)
+                        } else {
+                            push_value
+                        };
+                        let new_pc_abs = pc_mapping.get(&old_pc_abs).copied().unwrap_or(0);
+                        tracing::debug!(
+                            "remap_orphan_jump_pushes: block {} instr {} at pc=0x{:x}: \
+                             0x{:x} -> 0x{:x} (abs: 0x{:x} -> 0x{:x}) [extended scan]",
+                            node.index(),
+                            idx,
+                            instr.pc,
+                            push_value,
+                            new_value,
+                            old_pc_abs,
+                            new_pc_abs,
+                        );
+                        edits.push((node, idx, new_value));
+                    }
+                }
+            }
+        }
+
+        // Second pass: apply the edits
+        let total_remapped = edits.len();
+        for (node, instr_idx, new_value) in edits {
+            if let Some(Block::Body(body)) = self.cfg.node_weight_mut(node) {
+                apply_immediate(&mut body.instructions[instr_idx], new_value)?;
+            }
+        }
+
+        if total_remapped > 0 {
+            tracing::debug!(
+                "remap_orphan_jump_pushes: remapped {} internal-call return address PUSHes",
+                total_remapped
+            );
+        }
+
+        Ok(())
+    }
+
     /// Remap all stored metadata that references absolute PCs using the supplied mapping.
     ///
     /// This should be called any time a transform invokes `reindex_pcs` directly so that

crates/core/src/strip.rs

@@ -398,6 +398,115 @@ impl CleanReport {
         Ok(())
     }
 
+    /// Patch immutable reference offsets in the init code.
+    ///
+    /// The Solidity compiler's init code writes immutable variable values into the
+    /// runtime bytecode at hardcoded byte offsets. When obfuscation transforms change
+    /// the runtime layout (e.g., PushSplit growing blocks), these offsets become stale.
+    /// This method detects the pattern `PUSH2 <offset>; ... ADD` in the init code and
+    /// updates each offset using the supplied byte-offset remapping closure.
+    ///
+    /// The `remap` closure takes an old byte offset within the runtime and returns the
+    /// new byte offset, or `None` if no mapping is available.
+    pub fn patch_init_immutable_refs(
+        &mut self,
+        remap: &dyn Fn(usize) -> Option<usize>,
+    ) -> Result<(), String> {
+        let runtime_start = self
+            .runtime_layout
+            .iter()
+            .map(|span| span.offset)
+            .min()
+            .ok_or("No runtime layout found")?;
+        let runtime_end = runtime_start + self.clean_len;
+
+        let init_section = self
+            .removed
+            .iter_mut()
+            .find(|r| matches!(r.kind, SectionKind::Init))
+            .ok_or("No Init section found")?;
+
+        let mut init_bytes = init_section.data.to_vec();
+        let mut patched = 0usize;
+        let mut idx = 0usize;
+
+        while idx < init_bytes.len() {
+            let opcode = init_bytes[idx];
+            if !(0x60..=0x7f).contains(&opcode) {
+                idx += 1;
+                continue;
+            }
+
+            let width = (opcode - 0x60 + 1) as usize;
+            if idx + 1 + width > init_bytes.len() {
+                idx += 1;
+                continue;
+            }
+
+            let mut value = 0usize;
+            for &byte in &init_bytes[idx + 1..idx + 1 + width] {
+                value = (value << 8) | byte as usize;
+            }
+
+            // Check if the next non-stack-manipulation opcode is ADD (0x01).
+            // The pattern is: PUSH2 <offset>; (DUP/SWAP ops); ADD
+            let after = idx + 1 + width;
+            let is_add_target = if after < init_bytes.len() {
+                init_bytes[after] == 0x01 // ADD immediately follows
+            } else {
+                false
+            };
+
+            // Only remap values that look like runtime offsets followed by ADD
+            if is_add_target
+                && value >= 1
+                && value < runtime_end.saturating_sub(runtime_start)
+                && let Some(new_value) = remap(value)
+                && new_value != value
+            {
+                // Check that new value fits in the same width
+                let max = if width >= std::mem::size_of::<usize>() {
+                    usize::MAX
+                } else {
+                    (1usize << (width * 8)) - 1
+                };
+                if new_value > max {
+                    tracing::warn!(
+                        "Immutable ref at init offset 0x{:x}: new value 0x{:x} exceeds \
+                         PUSH{} capacity",
+                        idx,
+                        new_value,
+                        width
+                    );
+                } else {
+                    for j in 0..width {
+                        let shift = (width - 1 - j) * 8;
+                        init_bytes[idx + 1 + j] = ((new_value >> shift) & 0xff) as u8;
+                    }
+                    tracing::debug!(
+                        "Patched immutable ref at init offset 0x{:x}: 0x{:x} -> 0x{:x}",
+                        idx,
+                        value,
+                        new_value
+                    );
+                    patched += 1;
+                }
+            }
+
+            idx += 1 + width;
+        }
+
+        if patched > 0 {
+            tracing::debug!(
+                "Patched {} immutable reference offsets in init code",
+                patched
+            );
+            init_section.data = Bytes::from(init_bytes);
+        }
+
+        Ok(())
+    }
+
     /// Reassemble bytecode by placing the clean runtime at original offsets
     /// and filling removed sections with their original data.
     pub fn reassemble(&mut self, clean: &[u8]) -> Vec<u8> {