From 4da8f6720483dd812f68503119daefe2ea261d09 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Mon, 24 Nov 2025 14:25:46 +0100 Subject: [PATCH 1/3] fix: use substring matching to calc. precise offset where runtime starts --- crates/core/src/detection/sections.rs | 151 ++++++-------------------- crates/core/src/lib.rs | 43 +++++--- 2 files changed, 60 insertions(+), 134 deletions(-) diff --git a/crates/core/src/detection/sections.rs b/crates/core/src/detection/sections.rs index 3fb8e02..7d6c801 100644 --- a/crates/core/src/detection/sections.rs +++ b/crates/core/src/detection/sections.rs @@ -41,10 +41,14 @@ impl Section { } } -/// Locates all non-overlapping, offset-ordered sections in the bytecode. -pub fn locate_sections(bytes: &[u8], instructions: &[Instruction]) -> Result, Error> { +/// Locates all non-overlapping, offset-ordered sections in the deployment bytecode. +pub fn locate_sections( + deployment_bytes: &[u8], + instructions: &[Instruction], + runtime_bytes: &[u8], +) -> Result, Error> { let mut sections = Vec::new(); - let total_len = bytes.len(); + let total_len = deployment_bytes.len(); tracing::debug!( "Processing bytecode: {} bytes, {} instructions", @@ -53,7 +57,7 @@ pub fn locate_sections(bytes: &[u8], instructions: &[Instruction]) -> Result Result Runtime split using dispatcher pattern let (mut init_end, mut runtime_start, mut runtime_len) = - detect_init_runtime_split(instructions).unwrap_or((0, 0, aux_offset)); + detect_init_runtime_split(deployment_bytes, runtime_bytes).unwrap_or((0, 0, aux_offset)); tracing::debug!( "Initial detection: init_end={}, runtime_start={}, runtime_len={}", @@ -489,128 +493,43 @@ fn detect_padding(instructions: &[Instruction], aux_offset: usize) -> Option<(us }) } -/// Detects the Init to Runtime split using the dispatcher pattern. +/// Detects the Init to Runtime split by finding runtime bytecode within deployment bytecode. +/// +/// Uses a simple sliding window search to find where the runtime bytecode appears within +/// the deployment bytecode. This is more reliable than heuristic-based pattern matching. /// /// # Arguments -/// * `instructions` - Decoded instructions. +/// * `deployment_bytes` - The full deployment bytecode. +/// * `runtime_bytes` - The runtime bytecode to locate within deployment. /// /// # Returns -/// Optional tuple of (init_end, runtime_start, runtime_len) if pattern is found, None otherwise. -fn detect_init_runtime_split(instructions: &[Instruction]) -> Option<(usize, usize, usize)> { - // Try the strict pattern first (for backwards compatibility) - if let Some(result) = detect_strict_deployment_pattern(instructions) { - return Some(result); - } - - // Fallback: Look for any CODECOPY + RETURN pattern - if let Some(result) = detect_codecopy_return_pattern(instructions) { - return Some(result); +/// Optional tuple of (init_end, runtime_start, runtime_len) if found, None otherwise. +fn detect_init_runtime_split( + deployment_bytes: &[u8], + runtime_bytes: &[u8], +) -> Option<(usize, usize, usize)> { + if runtime_bytes.len() < 20 { + return None; } - None -} - -/// Detects the strict deployment pattern (original heuristic) -fn detect_strict_deployment_pattern(instructions: &[Instruction]) -> Option<(usize, usize, usize)> { - for i in 0..instructions.len().saturating_sub(6) { - if matches!(instructions[i].op, Opcode::PUSH(_) | Opcode::PUSH0) - && matches!(instructions[i + 1].op, Opcode::PUSH(_) | Opcode::PUSH0) - && matches!(instructions[i + 2].op, Opcode::PUSH0 | Opcode::PUSH(1)) - && instructions[i + 2].imm.as_deref() == Some("00") - && instructions[i + 3].op == Opcode::CODECOPY - && matches!(instructions[i + 4].op, Opcode::PUSH(_) | Opcode::PUSH0) - && instructions[i + 5].op == Opcode::RETURN - { - let runtime_len = instructions[i] - .imm - .as_ref() - .and_then(|s| usize::from_str_radix(s, 16).ok())?; - let runtime_ofs = instructions[i + 1] - .imm - .as_ref() - .and_then(|s| usize::from_str_radix(s, 16).ok())?; - let init_end = instructions[i + 5].pc + 1; + let needle = &runtime_bytes[..20]; + let runtime_offset = deployment_bytes + .windows(needle.len()) + .position(|window| window == needle)?; - tracing::debug!( - "Found strict deployment pattern at {}: init_end={}, runtime_start={}, runtime_len={}", - i, - init_end, - runtime_ofs, - runtime_len - ); + let remaining_runtime = &runtime_bytes[20..]; + let remaining_deployment = &deployment_bytes[runtime_offset + 20..]; - return Some((init_end, runtime_ofs, runtime_len)); + if remaining_deployment.starts_with(remaining_runtime) { + // if runtime starts at offset 0, we treat it as runtime-only bytecode with no init. + if runtime_offset == 0 { + return Some((0, 0, runtime_bytes.len())); } - } - None -} - -/// Fallback: Look for CODECOPY + RETURN pattern with more flexibility -fn detect_codecopy_return_pattern(instructions: &[Instruction]) -> Option<(usize, usize, usize)> { - // Find CODECOPY instruction - let codecopy_idx = instructions - .iter() - .position(|instruction| instruction.op == Opcode::CODECOPY)?; - - // Look for RETURN after CODECOPY (within reasonable distance) - let return_idx = instructions[codecopy_idx + 1..] - .iter() - .take(10) // Look within next 10 instructions - .position(|instruction| instruction.op == Opcode::RETURN) - .map(|pos| codecopy_idx + 1 + pos)?; - // Try to extract runtime parameters from PUSH instructions before CODECOPY - let mut runtime_len = None; - let mut runtime_start = None; - - // Look backwards from CODECOPY for PUSH instructions - // CODECOPY stack layout: [destOffset, offset, size] where offset is where runtime starts, - // and size is how many bytes to copy. Scanning backwards, we encounter them in reverse order. - for instruction in (0..codecopy_idx).rev().take(10) { - if matches!( - instructions[instruction].op, - Opcode::PUSH(_) | Opcode::PUSH0 - ) && let Some(immediate) = &instructions[instruction].imm - && let Ok(value) = usize::from_str_radix(immediate, 16) - { - if runtime_start.is_none() && value > 0 && value < 100000 { - // First reasonable value (scanning backwards) is the offset where runtime starts - runtime_start = Some(value); - } else if runtime_len.is_none() && value > 0 && value < 100000 { - // Second reasonable value is the size of the runtime code - runtime_len = Some(value); - } - - if runtime_len.is_some() && runtime_start.is_some() { - break; - } - } + return Some((runtime_offset - 1, runtime_offset, runtime_bytes.len())); } - // If we found CODECOPY + RETURN but can't extract parameters, - // make reasonable assumptions - let runtime_len = runtime_len.unwrap_or_else(|| { - // Estimate runtime length from instruction count after return - instructions.len().saturating_sub(return_idx + 1) * 2 // rough estimate - }); - - let runtime_start = runtime_start.unwrap_or_else(|| { - // Assume runtime starts right after the RETURN instruction - instructions[return_idx].pc + 1 - }); - - let init_end = instructions[return_idx].pc + 1; - - tracing::debug!( - "Found fallback deployment pattern: CODECOPY at {}, RETURN at {}, init_end={}, runtime_start={}, runtime_len={}", - codecopy_idx, - return_idx, - init_end, - runtime_start, - runtime_len - ); - - Some((init_end, runtime_start, runtime_len)) + None } /// Detects ConstructorArgs section between Init end and Runtime start. @@ -658,7 +577,7 @@ mod tests { use super::*; - const STORAGE_BYTECODE: &str = "6080604052348015600e575f5ffd5b50603e80601a5f395ff3fe60806040525f5ffdfea2646970667358221220e8c66682f723c073c8c5ec2c0de0795c9b8b64e310482b13bc56a554d057842b64736f6c634300081e0033"; + const STORAGE_BYTECODE: &str = include_str!("../../../../tests/bytecode/storage.hex"); #[test] fn detect_auxdata_works() { diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 477009f..9c9a72b 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -93,8 +93,10 @@ pub fn input_to_bytes(input: &str, is_file: bool) -> Result> { /// This function handles the complete pipeline from raw bytecode to CFG-IR, performing /// all necessary preprocessing steps that `cfg_ir::build_cfg_ir` expects to be done already. pub async fn process_bytecode_to_cfg( - bytecode: &str, - is_file: bool, + deployment_bytecode: &str, + deployment_is_file: bool, + runtime_bytecode: &str, + runtime_is_file: bool, ) -> std::result::Result< ( cfg_ir::CfgIrBundle, @@ -104,8 +106,10 @@ pub async fn process_bytecode_to_cfg( ), Box, > { - let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode, is_file).await?; - let sections = detection::locate_sections(&bytes, &instructions)?; + let (instructions, _, _, bytes) = + decoder::decode_bytecode(deployment_bytecode, deployment_is_file).await?; + let runtime_bytes = input_to_bytes(runtime_bytecode, runtime_is_file)?; + let sections = detection::locate_sections(&bytes, &instructions, &runtime_bytes)?; let (_, report) = strip::strip_bytecode(&bytes, §ions)?; // Filter instructions to only runtime section before building CFG @@ -148,6 +152,11 @@ mod tests { use crate::detection::SectionKind; use std::fs; + const COUNTER_DEPLOYMENT_BYTECODE: &str = + include_str!("../../../tests/bytecode/counter/counter_deployment.hex"); + const COUNTER_RUNTIME_BYTECODE: &str = + include_str!("../../../tests/bytecode/counter/counter_runtime.hex"); + #[test] fn normalize_hex_handles_whitespace_and_padding() { let raw = " 0xabc\n"; @@ -170,27 +179,25 @@ mod tests { #[tokio::test] async fn process_bytecode_builds_runtime_cfg() { - let bytecode = include_str!("../../../tests/bytecode/storage.hex"); - - let (bundle, instructions, sections, bytes) = process_bytecode_to_cfg(bytecode, false) - .await - .expect("cfg build"); + let (bundle, instructions, sections, bytes) = process_bytecode_to_cfg( + COUNTER_DEPLOYMENT_BYTECODE, + false, + COUNTER_RUNTIME_BYTECODE, + false, + ) + .await + .expect("cfg build"); assert!(!instructions.is_empty()); let runtime_section = sections .iter() .find(|s| s.kind == SectionKind::Runtime) .expect("runtime section present"); - assert_eq!(runtime_section.offset, 0x1a); - assert_eq!(runtime_section.len, 0x9); - - assert_eq!(bundle.runtime_bounds(), Some((0x1a, 0x23))); - assert_eq!(bundle.clean_report.clean_len, 0x9); - assert_eq!(bundle.clean_report.runtime_layout.len(), 1); - assert_eq!(bundle.clean_report.runtime_layout[0].offset, 0x1a); - assert_eq!(bundle.clean_report.runtime_layout[0].len, 0x9); - assert_eq!(bundle.original_bytecode, bytes); + // Counter contract has runtime at offset 28 (0x1c) + assert!(runtime_section.offset > 0); + assert!(runtime_section.len > 0); assert!(bundle.cfg.node_count() > 0, "cfg contains blocks"); + assert_eq!(bundle.original_bytecode, bytes); } } From ad1414202668012d8322d8a1b6bc03c6f070626a Mon Sep 17 00:00:00 2001 From: g4titanx Date: Mon, 24 Nov 2025 14:28:38 +0100 Subject: [PATCH 2/3] chore: update callers --- crates/analysis/src/obfuscation.rs | 13 ++++- crates/cli/src/commands/analyze.rs | 14 ++++-- crates/cli/src/commands/cfg.rs | 10 +++- crates/cli/src/commands/obfuscate.rs | 9 +++- crates/cli/src/commands/strip.rs | 10 +++- crates/transforms/src/obfuscator.rs | 5 +- crates/verification/src/lib.rs | 7 ++- crates/verification/src/semantics.rs | 3 +- examples/src/main.rs | 71 ++++++++++++++++++++-------- 9 files changed, 104 insertions(+), 38 deletions(-) diff --git a/crates/analysis/src/obfuscation.rs b/crates/analysis/src/obfuscation.rs index 18e77cc..9940f1a 100644 --- a/crates/analysis/src/obfuscation.rs +++ b/crates/analysis/src/obfuscation.rs @@ -26,6 +26,8 @@ pub struct AnalysisConfig<'a> { pub iterations: usize, /// Original bytecode (hex string with or without `0x` prefix). pub original_bytecode: &'a str, + /// Runtime bytecode (hex string with or without `0x` prefix). + pub runtime_bytecode: &'a str, /// Path to write the markdown report. pub report_path: PathBuf, /// Maximum attempts per iteration before giving up on a seed. @@ -34,10 +36,11 @@ pub struct AnalysisConfig<'a> { impl<'a> AnalysisConfig<'a> { /// Create config with sensible defaults. - pub fn new(original_bytecode: &'a str, iterations: usize) -> Self { + pub fn new(original_bytecode: &'a str, runtime_bytecode: &'a str, iterations: usize) -> Self { Self { iterations, original_bytecode, + runtime_bytecode, report_path: PathBuf::from("obfuscation_analysis_report.md"), max_attempts: 5, } @@ -370,7 +373,13 @@ pub async fn analyze_obfuscation( obfuscation_config.preserve_unknown_opcodes = true; obfuscation_config.transforms = passes.iter().map(|p| p.build()).collect(); - match obfuscate_bytecode(config.original_bytecode, obfuscation_config).await { + match obfuscate_bytecode( + config.original_bytecode, + config.runtime_bytecode, + obfuscation_config, + ) + .await + { Ok(result) => { let transforms_applied = result.metadata.transforms_applied.clone(); let obfuscated_bytes = hex_to_bytes(&result.obfuscated_bytecode)?; diff --git a/crates/cli/src/commands/analyze.rs b/crates/cli/src/commands/analyze.rs index 22d92b7..3a98775 100644 --- a/crates/cli/src/commands/analyze.rs +++ b/crates/cli/src/commands/analyze.rs @@ -3,16 +3,20 @@ use async_trait::async_trait; use azoth_analysis::obfuscation::{analyze_obfuscation, AnalysisConfig, AnalysisError}; use clap::Args; use std::{error::Error, path::PathBuf}; -const DEFAULT_BYTECODE_PATH: &str = "examples/escrow-bytecode/artifacts/deployment_bytecode.hex"; +const DEFAULT_DEPLOYMENT_PATH: &str = "examples/escrow-bytecode/artifacts/deployment_bytecode.hex"; +const DEFAULT_RUNTIME_PATH: &str = "examples/escrow-bytecode/artifacts/runtime_bytecode.hex"; /// Analyze how much bytecode survives obfuscation across multiple seeds. #[derive(Args)] pub struct AnalyzeArgs { /// Number of obfuscated samples to generate. pub iterations: usize, - /// Input bytecode as hex, .hex file, or binary file. - #[arg(value_name = "BYTECODE", default_value = DEFAULT_BYTECODE_PATH)] + /// Input deployment bytecode as hex, .hex file, or binary file. + #[arg(value_name = "BYTECODE", default_value = DEFAULT_DEPLOYMENT_PATH)] pub input: String, + /// Input runtime bytecode as hex, .hex file, or binary file. + #[arg(long, value_name = "RUNTIME", default_value = DEFAULT_RUNTIME_PATH)] + pub runtime: String, /// Where to write the markdown report (default: ./obfuscation_analysis_report.md). #[arg(long, value_name = "PATH")] output: Option, @@ -27,13 +31,15 @@ impl super::Command for AnalyzeArgs { let AnalyzeArgs { iterations, input, + runtime, output, max_attempts, } = self; let input_hex = read_input(&input)?; + let runtime_hex = read_input(&runtime)?; - let mut config = AnalysisConfig::new(&input_hex, iterations); + let mut config = AnalysisConfig::new(&input_hex, &runtime_hex, iterations); config.max_attempts = max_attempts; if let Some(path) = output { config.report_path = path; diff --git a/crates/cli/src/commands/cfg.rs b/crates/cli/src/commands/cfg.rs index 9af6494..c147919 100644 --- a/crates/cli/src/commands/cfg.rs +++ b/crates/cli/src/commands/cfg.rs @@ -6,6 +6,7 @@ use async_trait::async_trait; use azoth_core::cfg_ir::{build_cfg_ir, Block, CfgIrBundle, EdgeType}; use azoth_core::decoder::decode_bytecode; use azoth_core::detection::locate_sections; +use azoth_core::input_to_bytes; use azoth_core::strip::strip_bytecode; use clap::Args; use std::error::Error; @@ -15,8 +16,11 @@ use std::path::Path; /// Arguments for the `cfg` subcommand. #[derive(Args)] pub struct CfgArgs { - /// Input bytecode as a hex string (0x...) or file path containing EVM bytecode. + /// Input deployment bytecode as a hex string (0x...) or file path containing EVM bytecode. pub input: String, + /// Input runtime bytecode as a hex string (0x...) or file path containing EVM bytecode. + #[arg(long)] + pub runtime: String, /// Output file for Graphviz .dot (default: stdout) #[arg(short, long)] output: Option, @@ -27,8 +31,10 @@ pub struct CfgArgs { impl super::Command for CfgArgs { async fn execute(self) -> Result<(), Box> { let is_file = !self.input.starts_with("0x") && Path::new(&self.input).is_file(); + let runtime_is_file = !self.runtime.starts_with("0x") && Path::new(&self.runtime).is_file(); let (instructions, _, _, bytes) = decode_bytecode(&self.input, is_file).await?; - let sections = locate_sections(&bytes, &instructions)?; + let runtime_bytes = input_to_bytes(&self.runtime, runtime_is_file)?; + let sections = locate_sections(&bytes, &instructions, &runtime_bytes)?; let (_clean_runtime, clean_report) = strip_bytecode(&bytes, §ions)?; let cfg_ir = build_cfg_ir(&instructions, §ions, clean_report, &bytes)?; diff --git a/crates/cli/src/commands/obfuscate.rs b/crates/cli/src/commands/obfuscate.rs index 04cf347..c35e6b9 100644 --- a/crates/cli/src/commands/obfuscate.rs +++ b/crates/cli/src/commands/obfuscate.rs @@ -19,8 +19,11 @@ use std::path::Path; /// Arguments for the `obfuscate` subcommand. #[derive(Args)] pub struct ObfuscateArgs { - /// Input bytecode as a hex string, .hex file, or binary file containing EVM bytecode. + /// Input deployment bytecode as a hex string, .hex file, or binary file containing EVM bytecode. pub input: String, + /// Input runtime bytecode as a hex string, .hex file, or binary file containing EVM bytecode. + #[arg(long)] + pub runtime: String, /// Cryptographic seed for deterministic obfuscation #[arg(long)] seed: Option, @@ -42,6 +45,7 @@ impl super::Command for ObfuscateArgs { async fn execute(self) -> Result<(), Box> { let ObfuscateArgs { input, + runtime, seed, passes, emit, @@ -50,6 +54,7 @@ impl super::Command for ObfuscateArgs { // Step 1: Read and normalize input let input_bytecode = read_input(&input)?; + let runtime_bytecode = read_input(&runtime)?; // Step 2: Build transforms from CLI args let transforms = build_passes(&passes)?; @@ -68,7 +73,7 @@ impl super::Command for ObfuscateArgs { config.preserve_unknown_opcodes = true; // Step 4: Run obfuscation pipeline - let result = match obfuscate_bytecode(&input_bytecode, config).await { + let result = match obfuscate_bytecode(&input_bytecode, &runtime_bytecode, config).await { Ok(result) => result, Err(e) => return Err(format!("{e}").into()), }; diff --git a/crates/cli/src/commands/strip.rs b/crates/cli/src/commands/strip.rs index f3d0091..e616ad8 100644 --- a/crates/cli/src/commands/strip.rs +++ b/crates/cli/src/commands/strip.rs @@ -5,6 +5,7 @@ use async_trait::async_trait; use azoth_core::decoder::decode_bytecode; use azoth_core::detection::locate_sections; +use azoth_core::input_to_bytes; use azoth_core::strip::strip_bytecode; use clap::Args; use serde_json; @@ -14,8 +15,11 @@ use std::path::Path; /// Arguments for the `strip` subcommand. #[derive(Args)] pub struct StripArgs { - /// Input bytecode as a hex string (0x...) or file path containing EVM bytecode. + /// Input deployment bytecode as a hex string (0x...) or file path containing EVM bytecode. pub input: String, + /// Input runtime bytecode as a hex string (0x...) or file path containing EVM bytecode. + #[arg(long)] + pub runtime: String, /// Output raw cleaned runtime hex instead of JSON report #[arg(long)] raw: bool, @@ -26,8 +30,10 @@ pub struct StripArgs { impl super::Command for StripArgs { async fn execute(self) -> Result<(), Box> { let is_file = !self.input.starts_with("0x") && Path::new(&self.input).is_file(); + let runtime_is_file = !self.runtime.starts_with("0x") && Path::new(&self.runtime).is_file(); let (instructions, _, _, bytes) = decode_bytecode(&self.input, is_file).await?; - let sections = locate_sections(&bytes, &instructions)?; + let runtime_bytes = input_to_bytes(&self.runtime, runtime_is_file)?; + let sections = locate_sections(&bytes, &instructions, &runtime_bytes)?; let (clean_runtime, report) = strip_bytecode(&bytes, §ions)?; if self.raw { diff --git a/crates/transforms/src/obfuscator.rs b/crates/transforms/src/obfuscator.rs index 92b730e..313cedd 100644 --- a/crates/transforms/src/obfuscator.rs +++ b/crates/transforms/src/obfuscator.rs @@ -94,7 +94,8 @@ pub struct ObfuscationMetadata { /// Main obfuscation pipeline pub async fn obfuscate_bytecode( - input_bytecode: &str, + deployment_bytecode: &str, + runtime_bytecode: &str, config: ObfuscationConfig, ) -> Result> { tracing::debug!("Starting obfuscation pipeline:"); @@ -102,7 +103,7 @@ pub async fn obfuscate_bytecode( // Step 1: Process bytecode to CFG-IR let (mut cfg_ir, instructions, sections, bytes) = - process_bytecode_to_cfg(input_bytecode, false).await?; + process_bytecode_to_cfg(deployment_bytecode, false, runtime_bytecode, false).await?; let original_size = bytes.len(); tracing::debug!(" Input size: {} bytes", original_size); diff --git a/crates/verification/src/lib.rs b/crates/verification/src/lib.rs index fe14a9d..e077230 100644 --- a/crates/verification/src/lib.rs +++ b/crates/verification/src/lib.rs @@ -37,7 +37,9 @@ impl FormalVerifier { pub async fn prove_equivalence( &mut self, original_bytecode: &[u8], + original_runtime: &[u8], obfuscated_bytecode: &[u8], + obfuscated_runtime: &[u8], security_properties: &[SecurityProperty], ) -> VerificationResult { let start_time = Instant::now(); @@ -46,9 +48,10 @@ impl FormalVerifier { // Parse both contracts into semantic representations let original_semantics = - semantics::extract_semantics_from_bytecode(original_bytecode).await?; + semantics::extract_semantics_from_bytecode(original_bytecode, original_runtime).await?; let obfuscated_semantics = - semantics::extract_semantics_from_bytecode(obfuscated_bytecode).await?; + semantics::extract_semantics_from_bytecode(obfuscated_bytecode, obfuscated_runtime) + .await?; tracing::debug!("Extracted semantics for both contracts"); diff --git a/crates/verification/src/semantics.rs b/crates/verification/src/semantics.rs index 1aa1367..1e4b9fc 100644 --- a/crates/verification/src/semantics.rs +++ b/crates/verification/src/semantics.rs @@ -277,6 +277,7 @@ pub fn extract_semantics(cfg_bundle: &CfgIrBundle) -> VerificationResult VerificationResult { tracing::debug!( "Extracting semantics from bytecode ({} bytes)", @@ -288,7 +289,7 @@ pub async fn extract_semantics_from_bytecode( .await .map_err(|e| Error::BytecodeAnalysis(format!("Failed to decode bytecode: {e}")))?; - let sections = detection::locate_sections(bytecode, &instructions) + let sections = detection::locate_sections(bytecode, &instructions, runtime_bytes) .map_err(|e| Error::BytecodeAnalysis(format!("Failed to detect sections: {e}")))?; let (_clean_runtime, clean_report) = strip::strip_bytecode(bytecode, §ions) diff --git a/examples/src/main.rs b/examples/src/main.rs index cfd3c7d..ad6889a 100644 --- a/examples/src/main.rs +++ b/examples/src/main.rs @@ -5,23 +5,29 @@ use azoth_transform::obfuscator::{obfuscate_bytecode, ObfuscationConfig, Obfusca use serde_json::json; use std::fs; -const MIRAGE_ESCROW_PATH: &str = "escrow-bytecode/artifacts/bytecode.hex"; +const MIRAGE_ESCROW_DEPLOYMENT_PATH: &str = "escrow-bytecode/artifacts/deployment_bytecode.hex"; +const MIRAGE_ESCROW_RUNTIME_PATH: &str = "escrow-bytecode/artifacts/runtime_bytecode.hex"; #[tokio::main] async fn main() -> Result<(), Box> { println!("Mirage Privacy Protocol - Obfuscation Workflow"); println!("================================================="); - // Load contract bytecode - let original_bytecode = load_mirage_contract()?; + // Load contract bytecode (both deployment and runtime) + let (original_bytecode, runtime_bytecode) = load_mirage_contract()?; let seed_k2 = Seed::from_hex("0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef")?; - println!("Loaded Escrow bytecode: {} bytes", original_bytecode.len()); + println!( + "Loaded Escrow bytecode: {} bytes (deployment), {} bytes (runtime)", + original_bytecode.len(), + runtime_bytecode.len() + ); // SENDER: Compile with obfuscation O(S, K2) println!("\nSENDER: Compiling contract with obfuscation..."); - let obfuscation_result = apply_mirage_obfuscation(&original_bytecode, &seed_k2).await?; + let obfuscation_result = + apply_mirage_obfuscation(&original_bytecode, &runtime_bytecode, &seed_k2).await?; let obfuscated_bytecode = hex::decode( obfuscation_result .obfuscated_bytecode @@ -51,7 +57,8 @@ async fn main() -> Result<(), Box> { // VERIFIER: Verify bytecode integrity println!("\nVERIFIER: Verifying deterministic compilation with K2..."); - let recompilation_result = apply_mirage_obfuscation(&original_bytecode, &seed_k2).await?; + let recompilation_result = + apply_mirage_obfuscation(&original_bytecode, &runtime_bytecode, &seed_k2).await?; let recompiled_bytecode = hex::decode( recompilation_result .obfuscated_bytecode @@ -94,7 +101,7 @@ async fn main() -> Result<(), Box> { // Deterministic compilation verification println!("\nDETERMINISTIC COMPILATION TEST:"); - verify_deterministic_compilation_test(&original_bytecode, &seed_k2).await?; + verify_deterministic_compilation_test(&original_bytecode, &runtime_bytecode, &seed_k2).await?; // Generate comprehensive report let report = generate_workflow_report( @@ -120,32 +127,53 @@ async fn main() -> Result<(), Box> { Ok(()) } -/// Load Escrow contract bytecode from submodule artifact -fn load_mirage_contract() -> Result, Box> { - let content = fs::read_to_string(MIRAGE_ESCROW_PATH) - .map_err(|_| format!("Failed to load bytecode from {MIRAGE_ESCROW_PATH}\nRun './run_escrow.sh' to update submodule"))?; - - let clean_bytecode = content.trim().strip_prefix("0x").unwrap_or(content.trim()); - - if clean_bytecode.is_empty() || clean_bytecode.len() < 20 { - return Err("Invalid or empty bytecode in artifact".into()); +/// Load Escrow contract bytecode from submodule artifact (both deployment and runtime) +fn load_mirage_contract() -> Result<(Vec, Vec), Box> { + let deployment_content = fs::read_to_string(MIRAGE_ESCROW_DEPLOYMENT_PATH).map_err(|_| { + format!("Failed to load deployment bytecode from {MIRAGE_ESCROW_DEPLOYMENT_PATH}") + })?; + let runtime_content = fs::read_to_string(MIRAGE_ESCROW_RUNTIME_PATH).map_err(|_| { + format!("Failed to load runtime bytecode from {MIRAGE_ESCROW_RUNTIME_PATH}") + })?; + + let clean_deployment = deployment_content + .trim() + .strip_prefix("0x") + .unwrap_or(deployment_content.trim()); + let clean_runtime = runtime_content + .trim() + .strip_prefix("0x") + .unwrap_or(runtime_content.trim()); + + if clean_deployment.is_empty() || clean_deployment.len() < 20 { + return Err("Invalid or empty deployment bytecode in artifact".into()); + } + if clean_runtime.is_empty() || clean_runtime.len() < 20 { + return Err("Invalid or empty runtime bytecode in artifact".into()); } - hex::decode(clean_bytecode).map_err(|e| format!("Hex decode error: {e}").into()) + let deployment = + hex::decode(clean_deployment).map_err(|e| format!("Deployment hex decode error: {e}"))?; + let runtime = + hex::decode(clean_runtime).map_err(|e| format!("Runtime hex decode error: {e}"))?; + + Ok((deployment, runtime)) } /// Apply Mirage obfuscation transforms using the unified pipeline async fn apply_mirage_obfuscation( bytecode: &[u8], + runtime_bytecode: &[u8], seed_k2: &Seed, ) -> Result> { let hex_input = format!("0x{}", hex::encode(bytecode)); + let runtime_hex = format!("0x{}", hex::encode(runtime_bytecode)); // Create Mirage-specific transform configuration let config = create_mirage_config(seed_k2); // Use the unified obfuscation pipeline - obfuscate_bytecode(&hex_input, config).await + obfuscate_bytecode(&hex_input, &runtime_hex, config).await } /// Create Mirage-specific obfuscation configuration @@ -231,10 +259,11 @@ fn calculate_gas_percentage_increase(original: u64, new: u64) -> f64 { /// Verify deterministic compilation produces identical results async fn verify_deterministic_compilation_test( bytecode: &[u8], + runtime_bytecode: &[u8], seed: &Seed, ) -> Result<(), Box> { - let result1 = apply_mirage_obfuscation(bytecode, seed).await?; - let result2 = apply_mirage_obfuscation(bytecode, seed).await?; + let result1 = apply_mirage_obfuscation(bytecode, runtime_bytecode, seed).await?; + let result2 = apply_mirage_obfuscation(bytecode, runtime_bytecode, seed).await?; if result1.obfuscated_bytecode != result2.obfuscated_bytecode { return Err("Same seed produced different bytecode - not deterministic!".into()); @@ -243,7 +272,7 @@ async fn verify_deterministic_compilation_test( // Test different seeds produce different results let different_seed = Seed::generate(); - let diff_result = apply_mirage_obfuscation(bytecode, &different_seed).await?; + let diff_result = apply_mirage_obfuscation(bytecode, runtime_bytecode, &different_seed).await?; if result1.obfuscated_bytecode == diff_result.obfuscated_bytecode { return Err("Different seeds produced identical bytecode!".into()); } From 0d44544947913efea99561dbe5bf3e85128fa95d Mon Sep 17 00:00:00 2001 From: g4titanx Date: Mon, 24 Nov 2025 14:29:07 +0100 Subject: [PATCH 3/3] chore: update tests --- tests/src/analysis/metrics.rs | 18 ++-- tests/src/core/cfg_ir.rs | 2 +- tests/src/core/detection/dispatcher.rs | 2 +- tests/src/core/detection/sections.rs | 2 +- tests/src/core/strip.rs | 4 +- tests/src/e2e/deploy.rs | 103 +++++++++++++------- tests/src/e2e/escrow.rs | 12 ++- tests/src/e2e/test_counter.rs | 11 ++- tests/src/transforms/function_dispatcher.rs | 25 +++-- tests/src/transforms/jump_address.rs | 4 +- tests/src/transforms/opaque_predicate.rs | 4 +- tests/src/transforms/shuffle.rs | 8 +- 12 files changed, 128 insertions(+), 67 deletions(-) diff --git a/tests/src/analysis/metrics.rs b/tests/src/analysis/metrics.rs index c95d0bd..2be4c0a 100644 --- a/tests/src/analysis/metrics.rs +++ b/tests/src/analysis/metrics.rs @@ -16,7 +16,7 @@ async fn test_collect_metrics_simple() { let bytecode = "0x600160015601"; // PUSH1 0x01, PUSH1 0x01, ADD let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode, false).await.unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); @@ -45,7 +45,7 @@ async fn test_collect_metrics_single_block() { let bytecode = "0x600050"; // PUSH1 0x00, STOP let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode, false).await.unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); @@ -71,7 +71,7 @@ async fn test_collect_metrics_branching() { let bytecode = "0x6000600157600256"; // PUSH1 0x00, JUMPI, JUMPDEST, STOP let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode, false).await.unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); @@ -115,7 +115,7 @@ async fn test_collect_metrics_no_body_blocks() { let bytecode = "0x00"; // STOP let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode, false).await.unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); @@ -137,7 +137,7 @@ async fn test_compare_metrics() { .await .unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); let metrics_before = collect_metrics(&cfg_ir, &report).unwrap(); @@ -147,7 +147,7 @@ async fn test_compare_metrics() { .await .unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); let metrics_after = collect_metrics(&cfg_ir, &report).unwrap(); @@ -168,7 +168,7 @@ async fn test_potency_edge_increase() { let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode_simple, false) .await .unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); let metrics_simple = collect_metrics(&cfg_ir, &report).unwrap(); @@ -177,7 +177,7 @@ async fn test_potency_edge_increase() { let (instructions, _, _, bytes) = decoder::decode_bytecode(bytecode_complex, false) .await .unwrap(); - let sections = detection::locate_sections(&bytes, &instructions).unwrap(); + let sections = detection::locate_sections(&bytes, &instructions, &[]).unwrap(); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).unwrap(); let cfg_ir = cfg_ir::build_cfg_ir(&instructions, §ions, report.clone(), &bytes).unwrap(); let metrics_complex = collect_metrics(&cfg_ir, &report).unwrap(); @@ -197,7 +197,7 @@ async fn test_dominator_computation() { .without_time() .init(); let bytecode = "0x6000600157600256"; // PUSH1 0x00, PUSH1 0x01, JUMPI, PUSH1 0x02, JUMP - let (cfg_ir, _, _, _) = azoth_core::process_bytecode_to_cfg(bytecode, false) + let (cfg_ir, _, _, _) = azoth_core::process_bytecode_to_cfg(bytecode, false, bytecode, false) .await .unwrap(); diff --git a/tests/src/core/cfg_ir.rs b/tests/src/core/cfg_ir.rs index 1e58517..478c322 100644 --- a/tests/src/core/cfg_ir.rs +++ b/tests/src/core/cfg_ir.rs @@ -10,7 +10,7 @@ async fn build_storage_cfg() -> CfgIrBundle { let (instructions, _info, _asm, bytes) = decoder::decode_bytecode(STORAGE_BYTECODE, false) .await .expect("failed to decode storage bytecode"); - let sections = detection::locate_sections(&bytes, &instructions) + let sections = detection::locate_sections(&bytes, &instructions, &[]) .expect("failed to locate sections for storage bytecode"); let (_clean_runtime, report) = strip::strip_bytecode(&bytes, §ions).expect("failed to strip storage bytecode"); diff --git a/tests/src/core/detection/dispatcher.rs b/tests/src/core/detection/dispatcher.rs index 55a9aeb..8d8730d 100644 --- a/tests/src/core/detection/dispatcher.rs +++ b/tests/src/core/detection/dispatcher.rs @@ -20,7 +20,7 @@ async fn test_dispatcher_detection() { tracing::debug!("DecodeInfo: {:?}", info); // Step 1: Locate sections to find runtime - let sections = locate_sections(&bytes, &instructions).unwrap(); + let sections = locate_sections(&bytes, &instructions, &[]).unwrap(); tracing::debug!("Detected sections: {:?}", sections); // Step 2: Extract runtime instructions diff --git a/tests/src/core/detection/sections.rs b/tests/src/core/detection/sections.rs index 922a652..45301ba 100644 --- a/tests/src/core/detection/sections.rs +++ b/tests/src/core/detection/sections.rs @@ -19,7 +19,7 @@ async fn test_full_deploy_payload_properties() { tracing::debug!("Instructions: {:?}", instructions); tracing::debug!("DecodeInfo: {:?}", info); - let sections = locate_sections(&bytes, &instructions).unwrap(); + let sections = locate_sections(&bytes, &instructions, &[]).unwrap(); tracing::debug!("Detected sections: {:?}", sections); for (i, section) in sections.iter().enumerate() { diff --git a/tests/src/core/strip.rs b/tests/src/core/strip.rs index e43e5f4..1d99427 100644 --- a/tests/src/core/strip.rs +++ b/tests/src/core/strip.rs @@ -18,7 +18,7 @@ async fn test_round_trip() { let (instructions, _, _, bytecode) = decode_bytecode(COUNTER_DEPLOYMENT_BYTECODE, false) .await .unwrap(); - let sections = detection::locate_sections(&bytecode, &instructions).unwrap(); + let sections = detection::locate_sections(&bytecode, &instructions, &[]).unwrap(); let (clean_runtime, mut report) = strip_bytecode(&bytecode, §ions).unwrap(); let rebuilt = report.reassemble(&clean_runtime); @@ -42,7 +42,7 @@ async fn test_runtime_only() { let (instructions, _, _, bytecode) = decode_bytecode(COUNTER_RUNTIME_BYTECODE, false) .await .unwrap(); - let sections = detection::locate_sections(&bytecode, &instructions).unwrap(); + let sections = detection::locate_sections(&bytecode, &instructions, &[]).unwrap(); let (clean_runtime, mut report) = strip_bytecode(&bytecode, §ions).unwrap(); let rebuilt = report.reassemble(&clean_runtime); diff --git a/tests/src/e2e/deploy.rs b/tests/src/e2e/deploy.rs index 10ef7ae..ae31c5f 100644 --- a/tests/src/e2e/deploy.rs +++ b/tests/src/e2e/deploy.rs @@ -1,5 +1,6 @@ use super::{ - mock_token_bytecode, prepare_bytecode, ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, MOCK_TOKEN_ADDR, + mock_token_bytecode, prepare_bytecode, ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, MOCK_TOKEN_ADDR, }; use azoth_core::seed::Seed; use azoth_transform::jump_address_transformer::JumpAddressTransformer; @@ -109,9 +110,13 @@ async fn test_function_dispatch_only() -> Result<()> { println!("Testing FunctionDispatcher only (no additional transforms)"); let config = create_config_with_transforms(vec![], seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with function dispatcher: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with function dispatcher: {}", e))?; assert!(result .metadata @@ -137,9 +142,13 @@ async fn test_shuffle_transform() -> Result<()> { let transforms: Vec> = vec![Box::new(Shuffle)]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with shuffle: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with shuffle: {}", e))?; assert!(result .metadata @@ -170,9 +179,13 @@ async fn test_jump_address_transform() -> Result<()> { let transforms: Vec> = vec![Box::new(JumpAddressTransformer::new())]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with jump address transformer: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with jump address transformer: {}", e))?; assert!(result .metadata @@ -199,9 +212,13 @@ async fn test_opaque_predicate_transform() -> Result<()> { let transforms: Vec> = vec![Box::new(OpaquePredicate::new())]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with opaque predicate: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with opaque predicate: {}", e))?; assert!(result .metadata @@ -224,9 +241,13 @@ async fn test_shuffle_and_jump_address() -> Result<()> { let transforms: Vec> = vec![Box::new(Shuffle), Box::new(JumpAddressTransformer::new())]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with shuffle + jump address: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with shuffle + jump address: {}", e))?; assert!(result .metadata @@ -255,9 +276,13 @@ async fn test_shuffle_and_opaque_predicate() -> Result<()> { let transforms: Vec> = vec![Box::new(Shuffle), Box::new(OpaquePredicate::new())]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with shuffle + opaque predicate: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with shuffle + opaque predicate: {}", e))?; assert!(result .metadata @@ -289,14 +314,18 @@ async fn test_jump_address_and_opaque_predicate() -> Result<()> { Box::new(OpaquePredicate::new()), ]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| { - eyre!( - "Failed to obfuscate with jump address + opaque predicate: {}", - e - ) - })?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| { + eyre!( + "Failed to obfuscate with jump address + opaque predicate: {}", + e + ) + })?; assert!(result .metadata @@ -331,9 +360,13 @@ async fn test_all_transforms_enabled() -> Result<()> { Box::new(OpaquePredicate::new()), ]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate with all transforms: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate with all transforms: {}", e))?; // Verify all transforms were applied assert!(result @@ -379,9 +412,13 @@ async fn test_gas_consumption_analysis() -> Result<()> { let transforms: Vec> = vec![Box::new(Shuffle)]; let config = create_config_with_transforms(transforms, seed); - let result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate: {}", e))?; + let result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate: {}", e))?; let (_, obfuscated_gas) = deploy_and_verify_contract_revm(&result.obfuscated_bytecode, "Obfuscated")?; diff --git a/tests/src/e2e/escrow.rs b/tests/src/e2e/escrow.rs index dfce6a8..c776809 100644 --- a/tests/src/e2e/escrow.rs +++ b/tests/src/e2e/escrow.rs @@ -6,7 +6,7 @@ use super::{ mock_token_bytecode, prepare_bytecode, EscrowMappings, ObfuscatedCaller, - ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, MOCK_TOKEN_ADDR, + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, ESCROW_CONTRACT_RUNTIME_BYTECODE, MOCK_TOKEN_ADDR, }; use azoth_transform::obfuscator::{obfuscate_bytecode, ObfuscationConfig}; use azoth_transform::shuffle::Shuffle; @@ -33,9 +33,13 @@ async fn test_obfuscated_function_calls() -> Result<()> { let mut config = ObfuscationConfig::default(); config.transforms.push(Box::new(Shuffle)); - let obfuscation_result = obfuscate_bytecode(ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, config) - .await - .map_err(|e| eyre!("Failed to obfuscate bytecode: {:?}", e))?; + let obfuscation_result = obfuscate_bytecode( + ESCROW_CONTRACT_DEPLOYMENT_BYTECODE, + ESCROW_CONTRACT_RUNTIME_BYTECODE, + config, + ) + .await + .map_err(|e| eyre!("Failed to obfuscate bytecode: {:?}", e))?; println!( "✓ Contract obfuscated ({} -> {} bytes, {:+.1}%)", diff --git a/tests/src/e2e/test_counter.rs b/tests/src/e2e/test_counter.rs index b32c286..3a98047 100644 --- a/tests/src/e2e/test_counter.rs +++ b/tests/src/e2e/test_counter.rs @@ -58,10 +58,13 @@ async fn test_obfuscated_counter_deploys_and_counts() -> Result<()> { .without_time() .try_init(); - let obfuscation_result = - obfuscate_bytecode(COUNTER_DEPLOYMENT_BYTECODE, ObfuscationConfig::default()) - .await - .map_err(|e| eyre!("Bytecode transformation failed: {}", e))?; + let obfuscation_result = obfuscate_bytecode( + COUNTER_DEPLOYMENT_BYTECODE, + COUNTER_RUNTIME_BYTECODE, + ObfuscationConfig::default(), + ) + .await + .map_err(|e| eyre!("Bytecode transformation failed: {}", e))?; assert!( obfuscation_result diff --git a/tests/src/transforms/function_dispatcher.rs b/tests/src/transforms/function_dispatcher.rs index d25acb7..d1a0f5f 100644 --- a/tests/src/transforms/function_dispatcher.rs +++ b/tests/src/transforms/function_dispatcher.rs @@ -62,8 +62,12 @@ async fn test_dispatcher_transformation_and_determinism() { let config1 = ObfuscationConfig::with_seed(seed.clone()); let config2 = ObfuscationConfig::with_seed(seed.clone()); - let result1 = obfuscate_bytecode(SIMPLE_BYTECODE, config1).await.unwrap(); - let result2 = obfuscate_bytecode(SIMPLE_BYTECODE, config2).await.unwrap(); + let result1 = obfuscate_bytecode(SIMPLE_BYTECODE, SIMPLE_BYTECODE, config1) + .await + .unwrap(); + let result2 = obfuscate_bytecode(SIMPLE_BYTECODE, SIMPLE_BYTECODE, config2) + .await + .unwrap(); println!("\nTransformation result:"); println!(" Original: {} bytes", result1.original_size); @@ -164,9 +168,10 @@ async fn test_counter_dispatcher_detection() { .without_time() .try_init(); - let (_, instructions, sections, _) = process_bytecode_to_cfg(COUNTER_BYTECODE, false) - .await - .unwrap(); + let (_, instructions, sections, _) = + process_bytecode_to_cfg(COUNTER_BYTECODE, false, COUNTER_BYTECODE, false) + .await + .unwrap(); let runtime_section = sections .iter() @@ -190,9 +195,13 @@ async fn test_counter_dispatcher_detection() { "Expected four selectors in Counter dispatcher" ); - let result = obfuscate_bytecode(COUNTER_BYTECODE, ObfuscationConfig::default()) - .await - .expect("obfuscator should succeed"); + let result = obfuscate_bytecode( + COUNTER_BYTECODE, + COUNTER_BYTECODE, + ObfuscationConfig::default(), + ) + .await + .expect("obfuscator should succeed"); assert!( result diff --git a/tests/src/transforms/jump_address.rs b/tests/src/transforms/jump_address.rs index 278458c..47fe48c 100644 --- a/tests/src/transforms/jump_address.rs +++ b/tests/src/transforms/jump_address.rs @@ -13,7 +13,9 @@ async fn test_jump_address_transformer() { // Simple bytecode with a conditional jump let bytecode = "0x60085760015b00"; // PUSH1 0x08, JUMPI, PUSH1 0x01, JUMPDEST, STOP - let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false).await.unwrap(); + let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false, bytecode, false) + .await + .unwrap(); // Count instructions before transformation let mut instruction_count_before = 0; diff --git a/tests/src/transforms/opaque_predicate.rs b/tests/src/transforms/opaque_predicate.rs index 9786e60..17b0ce2 100644 --- a/tests/src/transforms/opaque_predicate.rs +++ b/tests/src/transforms/opaque_predicate.rs @@ -12,7 +12,9 @@ async fn test_opaque_predicate_adds_blocks() { .without_time() .init(); let bytecode = "0x6001600260016003"; // PUSH1 0x01, PUSH1 0x02, PUSH1 0x01, PUSH1 0x03 - let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false).await.unwrap(); + let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false, bytecode, false) + .await + .unwrap(); let before = collect_metrics(&cfg_ir, &cfg_ir.clean_report).unwrap(); let seed = Seed::from_hex("0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef") diff --git a/tests/src/transforms/shuffle.rs b/tests/src/transforms/shuffle.rs index 872fd5f..2c25e98 100644 --- a/tests/src/transforms/shuffle.rs +++ b/tests/src/transforms/shuffle.rs @@ -12,7 +12,9 @@ async fn test_shuffle_reorders_blocks() { .without_time() .init(); let bytecode = "0x6004565b60016000555b60026000555b6003600055"; - let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false).await.unwrap(); + let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false, bytecode, false) + .await + .unwrap(); let before = collect_metrics(&cfg_ir, &cfg_ir.clean_report).unwrap(); let seed = Seed::generate(); @@ -36,7 +38,9 @@ async fn test_shuffle_storage_bytecode() { .try_init(); let bytecode = include_str!("../../bytecode/storage.hex").trim(); - let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false).await.unwrap(); + let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false, bytecode, false) + .await + .unwrap(); // Collect block start PCs before shuffle let before_pcs: Vec = cfg_ir