From 63617f0fdead9a2705b0ece8d45f3f1d7372c947 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Thu, 15 Jan 2026 10:59:19 +0100 Subject: [PATCH 1/2] feat(decompilation): add labels to identify each changed caused by a transformation, init impl. --- Cargo.lock | 7 - Cargo.toml | 10 + crates/analysis/src/decompile_diff/mod.rs | 1 + crates/cli/src/commands/decompile_diff.rs | 425 +++++++++++++++++++++- vendor/heimdall-rs | 1 + 5 files changed, 424 insertions(+), 20 deletions(-) create mode 160000 vendor/heimdall-rs diff --git a/Cargo.lock b/Cargo.lock index c06b79ee..89a93df4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2789,7 +2789,6 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "heimdall-cache" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "bincode", "clap", @@ -2802,7 +2801,6 @@ dependencies = [ [[package]] name = "heimdall-common" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "alloy", "alloy-dyn-abi", @@ -2837,7 +2835,6 @@ dependencies = [ [[package]] name = "heimdall-config" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "clap", "heimdall-common", @@ -2852,7 +2849,6 @@ dependencies = [ [[package]] name = "heimdall-decoder" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "alloy", "alloy-dyn-abi", @@ -2874,7 +2870,6 @@ dependencies = [ [[package]] name = "heimdall-decompiler" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "alloy", "alloy-dyn-abi", @@ -2902,7 +2897,6 @@ dependencies = [ [[package]] name = "heimdall-disassembler" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "clap", "eyre", @@ -2917,7 +2911,6 @@ dependencies = [ [[package]] name = "heimdall-vm" version = "0.9.0" -source = "git+https://github.com/Jon-Becker/heimdall-rs?tag=0.9.0#173cfc90c1b9b9b01809627148a83dcdbf35b00f" dependencies = [ "alloy", "async-openai", diff --git a/Cargo.toml b/Cargo.toml index 1d19b31a..3de4b981 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = [ "crates/*", "tests", "examples"] +exclude = ["vendor/heimdall-rs"] resolver = "2" [workspace.package] @@ -93,3 +94,12 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } # Formal verification z3 = "0.12.1" + +[patch."https://github.com/Jon-Becker/heimdall-rs"] +heimdall-decompiler = { path = "vendor/heimdall-rs/crates/decompile" } +heimdall-disassembler = { path = "vendor/heimdall-rs/crates/disassemble" } +heimdall-common = { path = "vendor/heimdall-rs/crates/common" } +heimdall-config = { path = "vendor/heimdall-rs/crates/config" } +heimdall-cache = { path = "vendor/heimdall-rs/crates/cache" } +heimdall-decoder = { path = "vendor/heimdall-rs/crates/decode" } +heimdall-vm = { path = "vendor/heimdall-rs/crates/vm" } diff --git a/crates/analysis/src/decompile_diff/mod.rs b/crates/analysis/src/decompile_diff/mod.rs index 574db1ee..5642e4a5 100644 --- a/crates/analysis/src/decompile_diff/mod.rs +++ b/crates/analysis/src/decompile_diff/mod.rs @@ -399,6 +399,7 @@ pub async fn decompile(target: Bytes) -> Result { .target(target.to_string()) .output("print".into()) .include_solidity(true) + .annotate_pc(true) .build() .unwrap(); let result = heimdall_decompiler::decompile(args).await?; diff --git a/crates/cli/src/commands/decompile_diff.rs b/crates/cli/src/commands/decompile_diff.rs index 7630e5a3..7622068f 100644 --- a/crates/cli/src/commands/decompile_diff.rs +++ b/crates/cli/src/commands/decompile_diff.rs @@ -11,10 +11,12 @@ use async_trait::async_trait; use azoth_analysis::decompile_diff::{self, DiffStats, StructureKind, StructuredDiffResult}; +use azoth_core::cfg_ir::{CfgIrDiff, OperationKind, TraceEvent}; +use azoth_core::decoder::Instruction; use azoth_transform::obfuscator::{obfuscate_bytecode, ObfuscationConfig}; -use clap::Args; +use clap::{ArgAction, Args}; use owo_colors::OwoColorize; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::error::Error; use std::fs; use std::path::PathBuf; @@ -55,6 +57,10 @@ pub struct DecompileDiffArgs { /// Only show items that have changes. #[arg(long)] pub changed_only: bool, + + /// Disable heuristic cause annotations in diff output. + #[arg(long = "no-annotate-causes", action = ArgAction::SetFalse, default_value_t = true)] + pub annotate_causes: bool, } /// Aggregated statistics across multiple structured diff runs. @@ -66,6 +72,19 @@ struct AggregatedStructuredStats { sample_count: usize, } +#[derive(Debug, Clone)] +struct DiffRun { + diff: StructuredDiffResult, + attribution: PcAttribution, +} + +#[derive(Debug, Clone, Default)] +struct PcAttribution { + pre: HashMap>, + post: HashMap>, + observed: HashSet, +} + impl AggregatedStructuredStats { fn new(first: &DiffStats) -> Self { Self { @@ -131,7 +150,7 @@ impl super::Command for DecompileDiffArgs { let pre_bytes = Arc::new(pre_bytes); let passes = Arc::new(self.passes.clone()); - let mut join_set: JoinSet> = JoinSet::new(); + let mut join_set: JoinSet> = JoinSet::new(); for _ in 0..self.iterations { let input_bytecode = Arc::clone(&input_bytecode); @@ -171,7 +190,11 @@ impl super::Command for DecompileDiffArgs { .await .map_err(|e| format!("decompile: {e}"))?; - Ok(diff_result) + let attribution = build_pc_attribution(&obf_result.trace); + Ok(DiffRun { + diff: diff_result, + attribution, + }) }); } @@ -184,13 +207,15 @@ impl super::Command for DecompileDiffArgs { // Aggregate statistics let mut aggregated: Option = None; let mut first_result: Option = None; + let mut first_attribution: Option = None; for (i, result) in results.into_iter().enumerate() { - let diff_result = result?; - let stats = diff_result.aggregate_stats(); + let diff_run = result?; + let stats = diff_run.diff.aggregate_stats(); if i == 0 { - first_result = Some(diff_result); + first_result = Some(diff_run.diff); + first_attribution = Some(diff_run.attribution); } match &mut aggregated { @@ -201,13 +226,14 @@ impl super::Command for DecompileDiffArgs { let aggregated = aggregated.expect("at least one iteration"); let first_result = first_result.expect("at least one iteration"); + let first_attribution = first_attribution.expect("at least one iteration"); // Write diff to file if requested, otherwise print to terminal if let Some(output_path) = &self.output { - let output = self.format_diff_output(&first_result); + let output = self.format_diff_output(&first_result, &first_attribution); fs::write(output_path, output)?; } else { - self.print_structured_diff(&first_result); + self.print_structured_diff(&first_result, &first_attribution); } // Always print summary and statistics @@ -219,7 +245,11 @@ impl super::Command for DecompileDiffArgs { impl DecompileDiffArgs { /// Formats the diff only (no stats) as plain text for file output. - fn format_diff_output(&self, result: &StructuredDiffResult) -> String { + fn format_diff_output( + &self, + result: &StructuredDiffResult, + attribution: &PcAttribution, + ) -> String { let mut output = String::new(); for item in &result.items { @@ -230,7 +260,11 @@ impl DecompileDiffArgs { output.push_str(&format!("─── {} ───\n", item.kind)); if item.has_changes() { - output.push_str(&item.diff.unified_diff); + if self.annotate_causes { + output.push_str(&annotate_diff_plain(item, attribution)); + } else { + output.push_str(&render_clean_diff(item, false)); + } } else { output.push_str("(no changes)\n"); } @@ -241,7 +275,11 @@ impl DecompileDiffArgs { } /// Prints the structured diff to stdout with colors (no summary, just diffs). - fn print_structured_diff(&self, result: &StructuredDiffResult) { + fn print_structured_diff( + &self, + result: &StructuredDiffResult, + attribution: &PcAttribution, + ) { // Each item for item in &result.items { if self.changed_only && !item.has_changes() { @@ -295,7 +333,15 @@ impl DecompileDiffArgs { "added".dimmed() ); println!(); - print!("{}", item.diff.colored_diff); + if self.annotate_causes { + let summary = format_cause_summary(item, attribution); + if !summary.is_empty() { + println!(" {} {}", "Causes:".dimmed(), summary); + } + print!("{}", annotate_diff_colored(item, attribution)); + } else { + print!("{}", render_clean_diff(item, true)); + } } else { println!(" {}", "(no changes)".dimmed()); } @@ -369,3 +415,356 @@ impl DecompileDiffArgs { ); } } + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +enum DiffCause { + FunctionDispatcher, + SlotShuffle, + PushSplit, +} + +fn format_cause_summary( + item: &decompile_diff::StructuredDiffItem, + attribution: &PcAttribution, +) -> String { + let causes = collect_item_causes(item, attribution); + if causes.is_empty() { + return String::new(); + } + let labels: Vec<&'static str> = causes + .into_iter() + .map(|cause| match cause { + DiffCause::FunctionDispatcher => "FunctionDispatcher", + DiffCause::SlotShuffle => "SlotShuffle", + DiffCause::PushSplit => "PushSplit", + }) + .collect(); + labels.join(", ") +} + +fn annotate_diff_plain( + item: &decompile_diff::StructuredDiffItem, + attribution: &PcAttribution, +) -> String { + render_annotated_diff(item, attribution, false) +} + +fn annotate_diff_colored( + item: &decompile_diff::StructuredDiffItem, + attribution: &PcAttribution, +) -> String { + render_annotated_diff(item, attribution, true) +} + +fn render_annotated_diff( + item: &decompile_diff::StructuredDiffItem, + attribution: &PcAttribution, + colored: bool, +) -> String { + let mut out = String::new(); + for raw_line in item.diff.unified_diff.lines() { + let mut suffix = String::new(); + let mut display_line = raw_line.to_string(); + if let Some((prefix, content)) = split_diff_line(raw_line) { + let (clean_content, pc) = strip_pc_tag(content); + display_line = format!("{prefix}{clean_content}"); + let causes = causes_for_line(prefix, &clean_content, pc, item, attribution); + if !causes.is_empty() { + let tags = causes + .iter() + .map(|cause| match cause { + DiffCause::FunctionDispatcher => "FD", + DiffCause::SlotShuffle => "SS", + DiffCause::PushSplit => "PS", + }) + .collect::>() + .join(","); + suffix = format!(" [{}]", tags); + } + } else if let Some(stripped) = strip_context_pc(raw_line) { + display_line = stripped; + } + if colored { + display_line = colorize_diff_line(&display_line); + } + if !suffix.is_empty() { + if colored { + display_line.push_str(&suffix.dimmed().to_string()); + } else { + display_line.push_str(&suffix); + } + } + out.push_str(&display_line); + out.push('\n'); + } + out +} + +fn render_clean_diff(item: &decompile_diff::StructuredDiffItem, colored: bool) -> String { + let mut out = String::new(); + for raw_line in item.diff.unified_diff.lines() { + let mut display_line = raw_line.to_string(); + if let Some((prefix, content)) = split_diff_line(raw_line) { + let (clean_content, _) = strip_pc_tag(content); + display_line = format!("{prefix}{clean_content}"); + } else if let Some(stripped) = strip_context_pc(raw_line) { + display_line = stripped; + } + if colored { + display_line = colorize_diff_line(&display_line); + } + out.push_str(&display_line); + out.push('\n'); + } + out +} + +fn split_diff_line(line: &str) -> Option<(char, &str)> { + let mut chars = line.chars(); + let prefix = chars.next()?; + if prefix != '+' && prefix != '-' { + return None; + } + let content = chars.as_str(); + if content.starts_with("+++") || content.starts_with("---") { + return None; + } + Some((prefix, content)) +} + +fn strip_context_pc(line: &str) -> Option { + if !line.starts_with(' ') { + return None; + } + let content = &line[1..]; + let (clean_content, _) = strip_pc_tag(content); + Some(format!(" {}", clean_content)) +} + +fn strip_pc_tag(content: &str) -> (String, Option) { + let trimmed = content.trim_end(); + if let Some(start) = trimmed.rfind("/*pc=0x") { + if let Some(end) = trimmed[start..].find("*/") { + let hex_start = start + "/*pc=0x".len(); + let hex_end = start + end; + let hex = &trimmed[hex_start..hex_end]; + if let Ok(pc) = usize::from_str_radix(hex, 16) { + let mut clean = trimmed[..start].trim_end().to_string(); + if content.ends_with('\n') { + clean.push('\n'); + } + return (clean, Some(pc)); + } + } + } + (content.to_string(), None) +} + +fn colorize_diff_line(line: &str) -> String { + if line.starts_with("@@") { + return line.cyan().to_string(); + } + if line.starts_with('+') { + return line.green().to_string(); + } + if line.starts_with('-') { + return line.red().to_string(); + } + line.to_string() +} + +fn push_cause(causes: &mut Vec, cause: DiffCause) { + if !causes.contains(&cause) { + causes.push(cause); + } +} + +fn collect_item_causes( + item: &decompile_diff::StructuredDiffItem, + attribution: &PcAttribution, +) -> Vec { + let mut causes = Vec::new(); + for raw_line in item.diff.unified_diff.lines() { + if let Some((prefix, content)) = split_diff_line(raw_line) { + let (clean_content, pc) = strip_pc_tag(content); + let line_causes = causes_for_line(prefix, &clean_content, pc, item, attribution); + for cause in line_causes { + push_cause(&mut causes, cause); + } + } + } + causes +} + +fn causes_for_line( + prefix: char, + content: &str, + pc: Option, + item: &decompile_diff::StructuredDiffItem, + attribution: &PcAttribution, +) -> Vec { + let mut causes = Vec::new(); + if let Some(pc) = pc { + let target = match prefix { + '-' => attribution.pre.get(&pc), + '+' => attribution.post.get(&pc), + _ => None, + }; + if let Some(set) = target { + if set.len() == 1 { + if let Some(cause) = set.iter().copied().next() { + push_cause(&mut causes, cause); + } + } + } + return causes; + } + + if attribution + .observed + .contains(&DiffCause::FunctionDispatcher) + && should_tag_dispatcher_line(content, item) + { + push_cause(&mut causes, DiffCause::FunctionDispatcher); + } + + causes +} + +fn should_tag_dispatcher_line(content: &str, item: &decompile_diff::StructuredDiffItem) -> bool { + let trimmed = content.trim_start(); + let is_header = trimmed.contains("@custom:selector") + || trimmed.contains("@custom:signature") + || trimmed.starts_with("function "); + if !is_header { + return false; + } + match &item.kind { + StructureKind::Function { + original_selector, + obfuscated_selector, + .. + } => original_selector != obfuscated_selector, + StructureKind::UnmatchedOriginal { .. } | StructureKind::UnmatchedObfuscated { .. } => true, + _ => false, + } +} + +fn build_pc_attribution(trace: &[TraceEvent]) -> PcAttribution { + let mut attribution = PcAttribution::default(); + let mut active: Option = None; + let (pc_origin, origin_to_final) = build_pc_origin_maps(trace); + + for event in trace { + match &event.kind { + OperationKind::TransformStart { name } => { + active = cause_from_name(name); + if let Some(cause) = active { + attribution.observed.insert(cause); + } + } + OperationKind::TransformEnd { .. } => { + active = None; + } + _ => {} + } + + match &event.diff { + CfgIrDiff::BlockChanges(changes) => { + let Some(cause) = active else { + continue; + }; + for change in &changes.changes { + let (before_changed, after_changed) = diff_instruction_pcs( + &change.before.instructions, + &change.after.instructions, + &pc_origin, + ); + for origin in before_changed { + attribution.pre.entry(origin).or_default().insert(cause); + } + for origin in after_changed { + let final_pc = origin_to_final.get(&origin).copied().unwrap_or(origin); + attribution.post.entry(final_pc).or_default().insert(cause); + } + } + } + _ => {} + } + } + + attribution +} + +fn build_pc_origin_maps(trace: &[TraceEvent]) -> (HashMap, HashMap) { + let mut pc_origin: HashMap = HashMap::new(); + let mut origin_to_final: HashMap = HashMap::new(); + + for event in trace { + if let CfgIrDiff::PcsRemapped { instructions, .. } = &event.diff { + for instr in instructions { + let origin = pc_origin + .get(&instr.old_pc) + .copied() + .unwrap_or(instr.old_pc); + pc_origin.insert(instr.old_pc, origin); + pc_origin.insert(instr.new_pc, origin); + origin_to_final.insert(origin, instr.new_pc); + } + } + } + + (pc_origin, origin_to_final) +} + +fn diff_instruction_pcs( + before: &[Instruction], + after: &[Instruction], + pc_origin: &HashMap, +) -> (HashSet, HashSet) { + let mut before_changed = HashSet::new(); + let mut after_changed = HashSet::new(); + + let mut before_by_origin: HashMap = HashMap::new(); + for instr in before { + let origin = pc_origin.get(&instr.pc).copied().unwrap_or(instr.pc); + before_by_origin.insert(origin, instr); + } + + let mut after_by_origin: HashMap = HashMap::new(); + for instr in after { + let origin = pc_origin.get(&instr.pc).copied().unwrap_or(instr.pc); + after_by_origin.insert(origin, instr); + } + + for (origin, before_instr) in &before_by_origin { + match after_by_origin.get(origin) { + Some(after_instr) => { + if before_instr.op != after_instr.op || before_instr.imm != after_instr.imm { + before_changed.insert(*origin); + after_changed.insert(*origin); + } + } + None => { + before_changed.insert(*origin); + } + } + } + + for (origin, _) in &after_by_origin { + if !before_by_origin.contains_key(origin) { + after_changed.insert(*origin); + } + } + + (before_changed, after_changed) +} + +fn cause_from_name(name: &str) -> Option { + match name { + "FunctionDispatcher" => Some(DiffCause::FunctionDispatcher), + "SlotShuffle" => Some(DiffCause::SlotShuffle), + "PushSplit" => Some(DiffCause::PushSplit), + _ => None, + } +} diff --git a/vendor/heimdall-rs b/vendor/heimdall-rs new file mode 160000 index 00000000..173cfc90 --- /dev/null +++ b/vendor/heimdall-rs @@ -0,0 +1 @@ +Subproject commit 173cfc90c1b9b9b01809627148a83dcdbf35b00f From 17557a693651008d4e100c2520f569975b74ede0 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Thu, 15 Jan 2026 11:00:28 +0100 Subject: [PATCH 2/2] fix: clippy error --- crates/cli/src/commands/decompile_diff.rs | 43 ++++++++++------------- 1 file changed, 18 insertions(+), 25 deletions(-) diff --git a/crates/cli/src/commands/decompile_diff.rs b/crates/cli/src/commands/decompile_diff.rs index 7622068f..8a148ff3 100644 --- a/crates/cli/src/commands/decompile_diff.rs +++ b/crates/cli/src/commands/decompile_diff.rs @@ -275,11 +275,7 @@ impl DecompileDiffArgs { } /// Prints the structured diff to stdout with colors (no summary, just diffs). - fn print_structured_diff( - &self, - result: &StructuredDiffResult, - attribution: &PcAttribution, - ) { + fn print_structured_diff(&self, result: &StructuredDiffResult, attribution: &PcAttribution) { // Each item for item in &result.items { if self.changed_only && !item.has_changes() { @@ -669,27 +665,24 @@ fn build_pc_attribution(trace: &[TraceEvent]) -> PcAttribution { _ => {} } - match &event.diff { - CfgIrDiff::BlockChanges(changes) => { - let Some(cause) = active else { - continue; - }; - for change in &changes.changes { - let (before_changed, after_changed) = diff_instruction_pcs( - &change.before.instructions, - &change.after.instructions, - &pc_origin, - ); - for origin in before_changed { - attribution.pre.entry(origin).or_default().insert(cause); - } - for origin in after_changed { - let final_pc = origin_to_final.get(&origin).copied().unwrap_or(origin); - attribution.post.entry(final_pc).or_default().insert(cause); - } + if let CfgIrDiff::BlockChanges(changes) = &event.diff { + let Some(cause) = active else { + continue; + }; + for change in &changes.changes { + let (before_changed, after_changed) = diff_instruction_pcs( + &change.before.instructions, + &change.after.instructions, + &pc_origin, + ); + for origin in before_changed { + attribution.pre.entry(origin).or_default().insert(cause); + } + for origin in after_changed { + let final_pc = origin_to_final.get(&origin).copied().unwrap_or(origin); + attribution.post.entry(final_pc).or_default().insert(cause); } } - _ => {} } } @@ -751,7 +744,7 @@ fn diff_instruction_pcs( } } - for (origin, _) in &after_by_origin { + for origin in after_by_origin.keys() { if !before_by_origin.contains_key(origin) { after_changed.insert(*origin); }