thalir-core 0.1.0

Core IR types and builders for ThalIR smart contract analysis
Documentation
#![allow(unused_imports)]
#![allow(unused_variables)]

use thalir_core::{
    analysis::{Pass, PassManager},
    block::{BasicBlock, BlockId, Terminator},
    contract::{Contract, ContractMetadata, StorageLayout, StorageSlot},
    function::{Function, FunctionBody, FunctionSignature, Mutability, Visibility},
    instructions::Instruction,
    types::Type,
    values::{Constant, TempId, Value},
    ObfuscationConfig, ObfuscationLevel, ObfuscationPass, VulnerabilityMapper,
};

use indexmap::IndexMap;
use num_bigint::BigUint;

fn create_test_contract_with_identifiable_names() -> Contract {
    let mut contract = Contract {
        name: "NovelBondingCurveAMM".to_string(),
        functions: IndexMap::new(),
        storage_layout: StorageLayout {
            slots: vec![
                StorageSlot {
                    slot: BigUint::from(0u32),
                    offset: 0,
                    var_type: Type::Uint(256),
                    name: "liquidityPoolReserves".to_string(),
                    packed_with: Vec::new(),
                },
                StorageSlot {
                    slot: BigUint::from(1u32),
                    offset: 0,
                    var_type: Type::Address,
                    name: "contractOwner".to_string(),
                    packed_with: Vec::new(),
                },
            ],
            mappings: Vec::new(),
            arrays: Vec::new(),
            structs: Vec::new(),
        },
        events: Vec::new(),
        modifiers: Vec::new(),
        constants: Vec::new(),
        metadata: ContractMetadata {
            source_file: Some("NovelBondingCurveAMM.sol".to_string()),
            source_code: Some("contract NovelBondingCurveAMM { ... }".to_string()),
            ..Default::default()
        },
        source_files: thalir_core::SourceFiles::new(),
    };

    let mut func_body = FunctionBody::new();
    func_body.blocks.insert(
        func_body.entry_block,
        BasicBlock {
            id: func_body.entry_block,
            params: Vec::new(),
            instructions: vec![Instruction::Require {
                condition: Value::Constant(Constant::Bool(true)),
                message: "Invalid bonding curve parameters".to_string(),
            }],
            terminator: Terminator::Return(None),
            metadata: Default::default(),
        },
    );

    let function = Function {
        signature: FunctionSignature {
            name: "calculateBondingCurve".to_string(),
            params: vec![],
            returns: vec![],
            is_payable: false,
        },
        visibility: Visibility::Public,
        mutability: Mutability::NonPayable,
        modifiers: Vec::new(),
        body: func_body,
        metadata: Default::default(),
    };

    contract
        .functions
        .insert("calculateBondingCurve".to_string(), function);
    contract
}

#[test]
fn test_end_to_end_obfuscation_workflow() {
    let mut contract = create_test_contract_with_identifiable_names();

    assert_eq!(contract.name, "NovelBondingCurveAMM");
    assert!(contract.functions.contains_key("calculateBondingCurve"));
    assert_eq!(
        contract.storage_layout.slots[0].name,
        "liquidityPoolReserves"
    );
    assert_eq!(contract.storage_layout.slots[1].name, "contractOwner");

    let mut manager = PassManager::new();
    let obf_config = ObfuscationConfig {
        level: ObfuscationLevel::Minimal,
        retain_mapping: true,
        strip_string_constants: true,
        strip_error_messages: true,
        strip_metadata: true,
        ..Default::default()
    };

    manager.register_pass(ObfuscationPass::new(obf_config));
    manager.run_all(&mut contract).unwrap();

    assert_eq!(contract.name, "contract_0");
    assert!(!contract.functions.contains_key("calculateBondingCurve"));
    assert!(contract.functions.contains_key("fn_0"));
    assert_eq!(contract.storage_layout.slots[0].name, "var_0");
    assert_eq!(contract.storage_layout.slots[1].name, "var_1");

    assert!(contract.metadata.source_file.is_none());
    assert!(contract.metadata.source_code.is_none());

    let func = contract.functions.get("fn_0").unwrap();
    if let Some(block) = func.body.blocks.get(&func.body.entry_block) {
        if let Some(Instruction::Require { message, .. }) = block.instructions.first() {
            assert_eq!(message, "error_0");
        }
    }

    let pass = manager.get_pass::<ObfuscationPass>().unwrap();
    let mapping = pass.export_mapping();

    assert_eq!(
        mapping.deobfuscate("contract_0"),
        Some("NovelBondingCurveAMM")
    );
    assert_eq!(mapping.deobfuscate("fn_0"), Some("calculateBondingCurve"));
    assert_eq!(mapping.deobfuscate("var_0"), Some("liquidityPoolReserves"));
    assert_eq!(mapping.deobfuscate("var_1"), Some("contractOwner"));

    let obfuscated_report = "VULNERABILITY FOUND:\n\
        Reentrancy risk in contract_0::fn_0 at position [1]\n\
        Access to storage var_0 after external call\n\
        Consider using checks-effects-interactions pattern";

    let mapper = VulnerabilityMapper::from_mapping(mapping);
    let original_report = mapper.deobfuscate_report(obfuscated_report);

    assert!(original_report.contains("NovelBondingCurveAMM"));
    assert!(original_report.contains("calculateBondingCurve"));
    assert!(original_report.contains("liquidityPoolReserves"));
    assert!(!original_report.contains("contract_0"));
    assert!(!original_report.contains("fn_0"));
    assert!(!original_report.contains("var_0"));
}

#[test]
fn test_obfuscation_with_standard_hashing() {
    let mut contract = create_test_contract_with_identifiable_names();

    let mut manager = PassManager::new();
    let obf_config = ObfuscationConfig {
        level: ObfuscationLevel::Standard,
        retain_mapping: true,
        hash_salt: Some("test-salt-12345".to_string()),
        ..Default::default()
    };

    manager.register_pass(ObfuscationPass::new(obf_config));
    manager.run_all(&mut contract).unwrap();

    assert!(contract.name.starts_with("c_"));
    assert_eq!(contract.name.len(), 8);

    let func_names: Vec<_> = contract.functions.keys().collect();
    assert!(func_names[0].starts_with("f_"));
    assert_eq!(func_names[0].len(), 8);

    assert!(contract.storage_layout.slots[0].name.starts_with("v_"));
    assert_eq!(contract.storage_layout.slots[0].name.len(), 8);
}

#[test]
fn test_obfuscation_determinism() {
    let mut contract1 = create_test_contract_with_identifiable_names();
    let mut contract2 = create_test_contract_with_identifiable_names();

    let salt = "deterministic-salt";

    let config = ObfuscationConfig {
        level: ObfuscationLevel::Standard,
        hash_salt: Some(salt.to_string()),
        retain_mapping: true,
        ..Default::default()
    };

    let mut manager1 = PassManager::new();
    manager1.register_pass(ObfuscationPass::new(config.clone()));
    manager1.run_all(&mut contract1).unwrap();

    let mut manager2 = PassManager::new();
    manager2.register_pass(ObfuscationPass::new(config));
    manager2.run_all(&mut contract2).unwrap();

    assert_eq!(contract1.name, contract2.name);

    let func1_names: Vec<_> = contract1.functions.keys().collect();
    let func2_names: Vec<_> = contract2.functions.keys().collect();
    assert_eq!(func1_names, func2_names);

    assert_eq!(
        contract1.storage_layout.slots[0].name,
        contract2.storage_layout.slots[0].name
    );
}

#[test]
fn test_obfuscation_none_level_is_noop() {
    let mut contract = create_test_contract_with_identifiable_names();
    let original_name = contract.name.clone();
    let original_func_name = contract.functions.keys().next().unwrap().clone();

    let mut manager = PassManager::new();
    let obf_config = ObfuscationConfig {
        level: ObfuscationLevel::None,
        ..Default::default()
    };

    manager.register_pass(ObfuscationPass::new(obf_config));
    manager.run_all(&mut contract).unwrap();

    assert_eq!(contract.name, original_name);
    assert!(contract.functions.contains_key(&original_func_name));
}

#[test]
fn test_pass_preserves_structural_analyses() {
    let obf_config = ObfuscationConfig {
        level: ObfuscationLevel::Standard,
        ..Default::default()
    };

    let pass = ObfuscationPass::new(obf_config);
    let preserved = pass.preserved_analyses();

    use thalir_core::analysis::AnalysisID;

    assert!(preserved.contains(&AnalysisID::ControlFlow));
    assert!(preserved.contains(&AnalysisID::Dominator));
    assert!(preserved.contains(&AnalysisID::DefUse));
    assert!(preserved.contains(&AnalysisID::AliasAnalysis));
}