From b00b8aa123cdf7a0094674f87990a932920b5190 Mon Sep 17 00:00:00 2001 From: g4titanx Date: Fri, 23 Jan 2026 04:49:21 +0100 Subject: [PATCH] feat(transform): impl. cluster_shuffle transform --- crates/transforms/src/cluster_shuffle.rs | 210 ++++++++++++++++++++++- crates/transforms/src/obfuscator.rs | 7 +- 2 files changed, 210 insertions(+), 7 deletions(-) diff --git a/crates/transforms/src/cluster_shuffle.rs b/crates/transforms/src/cluster_shuffle.rs index 7e95611..38a0869 100644 --- a/crates/transforms/src/cluster_shuffle.rs +++ b/crates/transforms/src/cluster_shuffle.rs @@ -13,9 +13,11 @@ //! [storage_gate][dispatcher_tier_0][decoy_stub][controller_real] //! ``` -use crate::{Result, Transform}; -use azoth_core::cfg_ir::CfgIrBundle; -use rand::rngs::StdRng; +use crate::{Error, Result, Transform}; +use azoth_core::cfg_ir::{Block, CfgIrBundle}; +use petgraph::graph::NodeIndex; +use rand::{rngs::StdRng, seq::SliceRandom}; +use std::collections::{HashMap, HashSet}; use tracing::debug; /// Cluster-level shuffle wrapper. @@ -33,8 +35,204 @@ impl Transform for ClusterShuffle { "ClusterShuffle" } - fn apply(&self, _ir: &mut CfgIrBundle, _rng: &mut StdRng) -> Result { - debug!("ClusterShuffle: placeholder apply (no-op)"); - Ok(false) + fn apply(&self, ir: &mut CfgIrBundle, rng: &mut StdRng) -> Result { + let mut body_nodes: Vec<(usize, NodeIndex)> = ir + .cfg + .node_indices() + .filter_map(|n| match &ir.cfg[n] { + Block::Body(body) => Some((body.start_pc, n)), + _ => None, + }) + .collect(); + + if body_nodes.len() <= 1 { + debug!("ClusterShuffle: not enough blocks to shuffle"); + return Ok(false); + } + + body_nodes.sort_by_key(|(pc, _)| *pc); + let original_order: Vec = body_nodes.iter().map(|(pc, _)| *pc).collect(); + + let mut clusters: Vec> = Vec::new(); + let mut assigned: HashSet = HashSet::new(); + let start_pc_map: HashMap = + body_nodes.iter().map(|(pc, n)| (*n, *pc)).collect(); + + // Cluster 1: dispatcher blocks (kept adjacent as a group). + if !ir.dispatcher_blocks.is_empty() { + let mut dispatcher_cluster: Vec = body_nodes + .iter() + .filter_map(|(_, n)| { + if ir.dispatcher_blocks.contains(&n.index()) { + Some(*n) + } else { + None + } + }) + .collect(); + if !dispatcher_cluster.is_empty() { + dispatcher_cluster.sort_by_key(|n| start_pc_map[n]); + assigned.extend(dispatcher_cluster.iter().copied()); + clusters.push(dispatcher_cluster); + } + } + + // Cluster 2: stub+decoy pairs. + if let Some(stub_patches) = &ir.stub_patches { + for (stub_node, _, _, decoy_node) in stub_patches { + let mut cluster = Vec::new(); + for node in [*stub_node, *decoy_node] { + if assigned.contains(&node) { + continue; + } + if matches!(ir.cfg.node_weight(node), Some(Block::Body(_))) { + cluster.push(node); + } + } + if !cluster.is_empty() { + cluster.sort_by_key(|n| start_pc_map[n]); + assigned.extend(cluster.iter().copied()); + clusters.push(cluster); + } + } + } + + // Remaining nodes: singleton clusters. + for (_, node) in &body_nodes { + if assigned.contains(node) { + continue; + } + clusters.push(vec![*node]); + } + + let mut cluster_order: Vec> = clusters.clone(); + cluster_order.shuffle(rng); + + let mut new_order: Vec = Vec::with_capacity(body_nodes.len()); + for cluster in &cluster_order { + let mut ordered = cluster.clone(); + ordered.sort_by_key(|n| start_pc_map[n]); + new_order.extend(ordered.iter().map(|n| start_pc_map[n])); + } + + if original_order == new_order && cluster_order.len() > 1 { + debug!("ClusterShuffle: shuffle produced no change; rotating clusters"); + cluster_order.rotate_left(1); + } + + debug!( + "ClusterShuffle: original order (by start_pc): {:?}", + original_order + ); + debug!( + "ClusterShuffle: cluster order (sizes): {:?}", + cluster_order.iter().map(|c| c.len()).collect::>() + ); + + // Assign temporary PCs to enforce cluster order. + let mut pos = 0usize; + for cluster in &cluster_order { + let mut ordered = cluster.clone(); + ordered.sort_by_key(|n| start_pc_map[n]); + for node in ordered { + if let Some(Block::Body(body)) = ir.cfg.node_weight_mut(node) { + let temp_pc = pos * 1_000_000; + body.start_pc = temp_pc; + pos += 1; + } + } + } + + ir.reindex_pcs() + .map_err(|e| Error::CoreError(e.to_string()))?; + + Ok(true) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use azoth_core::cfg_ir::Block; + use azoth_core::process_bytecode_to_cfg; + use rand::SeedableRng; + + const STORAGE_BYTECODE: &str = include_str!("../../../tests/bytecode/storage.hex"); + + fn contiguous_positions(nodes: &[NodeIndex], order: &[NodeIndex]) -> bool { + let mut positions: Vec = nodes + .iter() + .filter_map(|n| order.iter().position(|o| o == n)) + .collect(); + positions.sort_unstable(); + if positions.len() <= 1 { + return true; + } + let min = positions[0]; + let max = positions[positions.len() - 1]; + max - min + 1 == positions.len() + } + + #[tokio::test] + async fn keeps_clusters_adjacent() { + let bytecode = STORAGE_BYTECODE.trim(); + let (mut cfg_ir, _, _, _) = process_bytecode_to_cfg(bytecode, false, bytecode, false) + .await + .unwrap(); + + let mut bodies: Vec<_> = cfg_ir + .cfg + .node_indices() + .filter(|n| matches!(cfg_ir.cfg[*n], Block::Body(_))) + .collect(); + bodies.sort_by_key(|n| { + if let Block::Body(ref b) = cfg_ir.cfg[*n] { + b.start_pc + } else { + 0 + } + }); + + assert!(bodies.len() >= 5, "storage fixture should have body blocks"); + + // Mark first three blocks as dispatcher cluster. + for node in bodies.iter().take(3) { + cfg_ir.dispatcher_blocks.insert(node.index()); + } + + // Create one stub+decoy pair from the next two blocks. + let stub = bodies[3]; + let decoy = bodies[4]; + cfg_ir.stub_patches = Some(vec![(stub, 0, 1, decoy)]); + + let mut rng = StdRng::seed_from_u64(0x5157_u64); + let shuffle = ClusterShuffle::new(); + let changed = shuffle.apply(&mut cfg_ir, &mut rng).unwrap(); + assert!(changed, "cluster shuffle should report changes"); + + let mut ordered: Vec = cfg_ir + .cfg + .node_indices() + .filter(|n| matches!(cfg_ir.cfg[*n], Block::Body(_))) + .collect(); + ordered.sort_by_key(|n| { + if let Block::Body(ref b) = cfg_ir.cfg[*n] { + b.start_pc + } else { + 0 + } + }); + + let dispatcher_nodes: Vec = bodies.iter().take(3).copied().collect(); + assert!( + contiguous_positions(&dispatcher_nodes, &ordered), + "dispatcher cluster should stay adjacent" + ); + + let pair = vec![stub, decoy]; + assert!( + contiguous_positions(&pair, &ordered), + "stub+decoy cluster should stay adjacent" + ); } } diff --git a/crates/transforms/src/obfuscator.rs b/crates/transforms/src/obfuscator.rs index c741d98..a3408a3 100644 --- a/crates/transforms/src/obfuscator.rs +++ b/crates/transforms/src/obfuscator.rs @@ -1,4 +1,5 @@ use crate::arithmetic_chain::ArithmeticChain; +use crate::cluster_shuffle::ClusterShuffle; use crate::function_dispatcher::FunctionDispatcher; use crate::push_split::PushSplit; use crate::Transform; @@ -62,7 +63,11 @@ impl Default for ObfuscationConfig { fn default() -> Self { Self { seed: Seed::generate(), - transforms: vec![Box::new(ArithmeticChain::new()), Box::new(PushSplit::new())], + transforms: vec![ + Box::new(ArithmeticChain::new()), + Box::new(PushSplit::new()), + Box::new(ClusterShuffle::new()), + ], preserve_unknown_opcodes: true, } }