Skip to main content

threatflux_binary_analysis/disasm/
mod.rs

1//! Disassembly module supporting multiple disassembly engines
2//!
3//! This module provides disassembly capabilities using both Capstone and iced-x86 engines.
4//! The choice of engine can be configured based on requirements and availability.
5
6use crate::{
7    AnalysisConfig, BinaryError, BinaryFile, Result,
8    types::{Architecture, Instruction, InstructionCategory, Section},
9};
10
11#[cfg(feature = "disasm-capstone")]
12use crate::types::ControlFlow as FlowType;
13
14#[cfg(feature = "disasm-capstone")]
15mod capstone_engine;
16
17#[cfg(feature = "disasm-iced")]
18mod iced_engine;
19
20/// Disassembly engine selection
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum DisassemblyEngine {
23    /// Use Capstone disassembly engine
24    #[cfg(feature = "disasm-capstone")]
25    Capstone,
26    /// Use iced-x86 disassembly engine  
27    #[cfg(feature = "disasm-iced")]
28    Iced,
29    /// Automatic engine selection
30    Auto,
31}
32
33/// Disassembly configuration
34#[derive(Debug, Clone)]
35pub struct DisassemblyConfig {
36    /// Preferred disassembly engine
37    pub engine: DisassemblyEngine,
38    /// Maximum number of instructions to disassemble
39    pub max_instructions: usize,
40    /// Include instruction details (operands, etc.)
41    pub detailed: bool,
42    /// Enable control flow analysis
43    pub analyze_control_flow: bool,
44    /// Skip invalid instructions
45    pub skip_invalid: bool,
46}
47
48impl Default for DisassemblyConfig {
49    fn default() -> Self {
50        Self {
51            engine: DisassemblyEngine::Auto,
52            max_instructions: 10000,
53            detailed: true,
54            analyze_control_flow: true,
55            skip_invalid: true,
56        }
57    }
58}
59
60/// Disassembler wrapper supporting multiple engines
61pub struct Disassembler {
62    config: DisassemblyConfig,
63    architecture: Architecture,
64}
65
66impl Disassembler {
67    /// Create a new disassembler for the specified architecture
68    pub fn new(architecture: Architecture) -> Result<Self> {
69        Ok(Self {
70            config: DisassemblyConfig::default(),
71            architecture,
72        })
73    }
74
75    /// Create disassembler with custom configuration
76    pub fn with_config(architecture: Architecture, config: DisassemblyConfig) -> Result<Self> {
77        Ok(Self {
78            config,
79            architecture,
80        })
81    }
82
83    /// Disassemble binary code
84    pub fn disassemble(&self, data: &[u8], address: u64) -> Result<Vec<Instruction>> {
85        let byte_length = u64::try_from(data.len()).map_err(|_| {
86            BinaryError::invalid_data("Disassembly byte length does not fit in u64")
87        })?;
88        address.checked_add(byte_length).ok_or_else(|| {
89            BinaryError::invalid_data(format!(
90                "Disassembly address range {address:#x} + {} bytes overflows u64",
91                data.len()
92            ))
93        })?;
94
95        let engine = self.select_engine()?;
96
97        match engine {
98            #[cfg(feature = "disasm-capstone")]
99            DisassemblyEngine::Capstone => {
100                capstone_engine::disassemble(data, address, self.architecture, &self.config)
101            }
102            #[cfg(feature = "disasm-iced")]
103            DisassemblyEngine::Iced => {
104                iced_engine::disassemble(data, address, self.architecture, &self.config)
105            }
106            DisassemblyEngine::Auto => {
107                // Try available engines in order of preference
108                #[cfg(feature = "disasm-capstone")]
109                {
110                    capstone_engine::disassemble(data, address, self.architecture, &self.config)
111                }
112                #[cfg(all(feature = "disasm-iced", not(feature = "disasm-capstone")))]
113                {
114                    iced_engine::disassemble(data, address, self.architecture, &self.config)
115                }
116                #[cfg(not(any(feature = "disasm-capstone", feature = "disasm-iced")))]
117                {
118                    Err(BinaryError::feature_not_available(
119                        "No disassembly engine available. Enable 'disasm-capstone' or 'disasm-iced' feature.",
120                    ))
121                }
122            }
123        }
124    }
125
126    /// Disassemble a specific section of a binary
127    pub fn disassemble_section(
128        &self,
129        binary: &BinaryFile,
130        section_name: &str,
131    ) -> Result<Vec<Instruction>> {
132        for section in binary.sections() {
133            if section.name == section_name {
134                let data = section_file_data(binary, section, 0, None)?;
135                return self.disassemble(data, section.address);
136            }
137        }
138
139        Err(BinaryError::invalid_data(format!(
140            "Section '{}' not found",
141            section_name
142        )))
143    }
144
145    /// Disassemble code at specific address with length
146    pub fn disassemble_at(
147        &self,
148        data: &[u8],
149        address: u64,
150        length: usize,
151    ) -> Result<Vec<Instruction>> {
152        if data.len() < length {
153            return Err(BinaryError::invalid_data(
154                "Insufficient data for disassembly",
155            ));
156        }
157
158        self.disassemble(&data[..length], address)
159    }
160
161    /// Select the appropriate disassembly engine
162    fn select_engine(&self) -> Result<DisassemblyEngine> {
163        match self.config.engine {
164            #[cfg(feature = "disasm-capstone")]
165            DisassemblyEngine::Capstone => Ok(DisassemblyEngine::Capstone),
166            #[cfg(feature = "disasm-iced")]
167            DisassemblyEngine::Iced => Ok(DisassemblyEngine::Iced),
168            DisassemblyEngine::Auto => {
169                // Select best engine for architecture
170                match self.architecture {
171                    Architecture::X86 | Architecture::X86_64 => {
172                        #[cfg(feature = "disasm-iced")]
173                        {
174                            Ok(DisassemblyEngine::Iced)
175                        }
176                        #[cfg(all(feature = "disasm-capstone", not(feature = "disasm-iced")))]
177                        {
178                            Ok(DisassemblyEngine::Capstone)
179                        }
180                        #[cfg(not(any(feature = "disasm-capstone", feature = "disasm-iced")))]
181                        {
182                            Err(BinaryError::feature_not_available(
183                                "No disassembly engine available",
184                            ))
185                        }
186                    }
187                    _ => {
188                        // For non-x86 architectures, prefer Capstone
189                        #[cfg(feature = "disasm-capstone")]
190                        {
191                            Ok(DisassemblyEngine::Capstone)
192                        }
193                        #[cfg(not(feature = "disasm-capstone"))]
194                        {
195                            Err(BinaryError::unsupported_arch(format!(
196                                "Architecture {:?} requires Capstone engine",
197                                self.architecture
198                            )))
199                        }
200                    }
201                }
202            }
203        }
204    }
205}
206
207/// Resolve a checked file-backed range within a section.
208///
209/// `Section::data` is only an inline preview for small sections. Analysis must use
210/// the owning `BinaryFile` so large sections are not silently skipped.
211pub(crate) fn section_file_data<'a>(
212    binary: &'a BinaryFile,
213    section: &Section,
214    relative_offset: u64,
215    requested_length: Option<usize>,
216) -> Result<&'a [u8]> {
217    if relative_offset > section.size || relative_offset > section.file_size {
218        return Err(BinaryError::invalid_data(format!(
219            "Offset {relative_offset} is outside section '{}'",
220            section.name
221        )));
222    }
223
224    let file_offset = section
225        .offset
226        .checked_add(relative_offset)
227        .ok_or_else(|| BinaryError::invalid_data("Section file offset overflows"))?;
228    let start = usize::try_from(file_offset)
229        .map_err(|_| BinaryError::invalid_data("Section file offset does not fit in memory"))?;
230    let data = binary.data();
231    if start > data.len() {
232        return Err(BinaryError::invalid_data(format!(
233            "Section '{}' starts beyond the end of the file",
234            section.name
235        )));
236    }
237
238    let virtual_remaining = section.size - relative_offset;
239    let file_remaining = section.file_size - relative_offset;
240    let section_remaining =
241        usize::try_from(virtual_remaining.min(file_remaining)).unwrap_or(usize::MAX);
242    let available = data.len() - start;
243    let length = requested_length
244        .unwrap_or(section_remaining)
245        .min(section_remaining);
246    if length > available {
247        return Err(BinaryError::invalid_data(format!(
248            "Section '{}' declares {length} file-backed bytes at offset {file_offset}, but only {available} remain",
249            section.name
250        )));
251    }
252    let end = start
253        .checked_add(length)
254        .ok_or_else(|| BinaryError::invalid_data("Section file range overflows"))?;
255    data.get(start..end)
256        .ok_or_else(|| BinaryError::invalid_data("Section file range is out of bounds"))
257}
258
259/// High-level function to disassemble binary data
260pub fn disassemble_binary(
261    binary: &BinaryFile,
262    config: &AnalysisConfig,
263) -> Result<Vec<Instruction>> {
264    let architecture = config.architecture_hint.unwrap_or(binary.architecture());
265    let mut disasm_config = DisassemblyConfig {
266        engine: config.disassembly_engine,
267        max_instructions: config.max_disassembly_instructions,
268        detailed: true,
269        analyze_control_flow: true,
270        skip_invalid: true,
271    };
272
273    let mut all_instructions = Vec::new();
274    let mut remaining_bytes = config.max_analysis_size;
275    let mut remaining_instructions = disasm_config.max_instructions;
276    let mut successful_sections = 0_usize;
277    let mut last_error = None;
278
279    let mut executable_sections: Vec<_> = binary
280        .sections()
281        .iter()
282        .filter(|section| section.permissions.execute)
283        .collect();
284    executable_sections.sort_by(|left, right| {
285        left.address
286            .cmp(&right.address)
287            .then_with(|| left.offset.cmp(&right.offset))
288            .then_with(|| left.name.cmp(&right.name))
289    });
290
291    // Disassemble executable sections
292    for section in executable_sections {
293        if remaining_bytes == 0 || remaining_instructions == 0 {
294            break;
295        }
296
297        let data = match section_file_data(binary, section, 0, Some(remaining_bytes)) {
298            Ok(data) => data,
299            Err(error) => {
300                last_error = Some(error);
301                continue;
302            }
303        };
304        remaining_bytes -= data.len();
305        disasm_config.max_instructions = remaining_instructions;
306        let disassembler = Disassembler::with_config(architecture, disasm_config.clone())?;
307
308        match disassembler.disassemble(data, section.address) {
309            Ok(mut instructions) => {
310                successful_sections += 1;
311                remaining_instructions = remaining_instructions.saturating_sub(instructions.len());
312                all_instructions.append(&mut instructions);
313            }
314            Err(error) => {
315                last_error = Some(error);
316            }
317        }
318    }
319
320    if let (0, Some(error)) = (successful_sections, last_error) {
321        return Err(error);
322    }
323
324    all_instructions.sort_by(|left, right| {
325        left.address
326            .cmp(&right.address)
327            .then_with(|| left.bytes.cmp(&right.bytes))
328            .then_with(|| left.mnemonic.cmp(&right.mnemonic))
329            .then_with(|| left.operands.cmp(&right.operands))
330    });
331    Ok(all_instructions)
332}
333
334/// Determine instruction category from mnemonic
335fn categorize_instruction(mnemonic: &str) -> InstructionCategory {
336    let mnemonic_lower = mnemonic.to_lowercase();
337
338    if mnemonic_lower.starts_with("add")
339        || mnemonic_lower.starts_with("sub")
340        || mnemonic_lower.starts_with("mul")
341        || mnemonic_lower.starts_with("div")
342        || mnemonic_lower.starts_with("inc")
343        || mnemonic_lower.starts_with("dec")
344    {
345        InstructionCategory::Arithmetic
346    } else if mnemonic_lower.starts_with("and")
347        || mnemonic_lower.starts_with("or")
348        || mnemonic_lower.starts_with("xor")
349        || mnemonic_lower.starts_with("not")
350        || mnemonic_lower.starts_with("shl")
351        || mnemonic_lower.starts_with("shr")
352    {
353        InstructionCategory::Logic
354    } else if mnemonic_lower.starts_with("mov")
355        || mnemonic_lower.starts_with("lea")
356        || mnemonic_lower.starts_with("push")
357        || mnemonic_lower.starts_with("pop")
358        || mnemonic_lower.starts_with("load")
359        || mnemonic_lower.starts_with("store")
360    {
361        InstructionCategory::Memory
362    } else if is_control_mnemonic(&mnemonic_lower) && !is_interrupt_mnemonic(&mnemonic_lower) {
363        InstructionCategory::Control
364    } else if mnemonic_lower.starts_with("int")
365        || mnemonic_lower.starts_with("syscall")
366        || mnemonic_lower.starts_with("sysenter")
367        || mnemonic_lower.starts_with("sysexit")
368    {
369        InstructionCategory::System
370    } else if mnemonic_lower.contains("aes")
371        || mnemonic_lower.contains("sha")
372        || mnemonic_lower.contains("crypto")
373    {
374        InstructionCategory::Crypto
375    } else if mnemonic_lower.starts_with("fadd")
376        || mnemonic_lower.starts_with("fsub")
377        || mnemonic_lower.starts_with("fmul")
378        || mnemonic_lower.starts_with("fdiv")
379    {
380        InstructionCategory::Float
381    } else if mnemonic_lower.contains("xmm")
382        || mnemonic_lower.contains("ymm")
383        || mnemonic_lower.contains("zmm")
384        || mnemonic_lower.starts_with("v")
385    {
386        InstructionCategory::Vector
387    } else {
388        InstructionCategory::Unknown
389    }
390}
391
392/// Determine control flow type from instruction
393#[cfg(feature = "disasm-capstone")]
394fn analyze_control_flow(mnemonic: &str, operands: &str) -> FlowType {
395    let mnemonic_lower = mnemonic.to_lowercase();
396
397    if is_return_mnemonic(&mnemonic_lower) {
398        FlowType::Return
399    } else if is_call_mnemonic(&mnemonic_lower) {
400        // Try to extract target address from operands
401        if let Some(addr) = extract_address_from_operands(operands) {
402            FlowType::Call(addr)
403        } else {
404            FlowType::Unknown // Indirect call
405        }
406    } else if is_unconditional_branch_mnemonic(&mnemonic_lower) {
407        if let Some(addr) = extract_address_from_operands(operands) {
408            FlowType::Jump(addr)
409        } else {
410            FlowType::Unknown // Indirect jump
411        }
412    } else if is_conditional_branch_mnemonic(&mnemonic_lower) {
413        // Conditional jumps
414        if let Some(addr) = extract_address_from_operands(operands) {
415            FlowType::ConditionalJump(addr)
416        } else {
417            FlowType::Unknown // Indirect conditional jump
418        }
419    } else if is_interrupt_mnemonic(&mnemonic_lower) {
420        FlowType::Interrupt
421    } else {
422        FlowType::Sequential
423    }
424}
425
426/// Extract address from instruction operands (simplified)
427#[cfg(feature = "disasm-capstone")]
428fn extract_address_from_operands(operands: &str) -> Option<u64> {
429    let operand = operands
430        .trim()
431        .trim_start_matches(['#', '$'])
432        .split_whitespace()
433        .next()?
434        .trim_end_matches(',');
435
436    // Look for hex addresses
437    if let Some(addr) = operand
438        .strip_prefix("0x")
439        .and_then(|hex_part| u64::from_str_radix(hex_part, 16).ok())
440    {
441        return Some(addr);
442    }
443
444    // Look for decimal addresses
445    if let Ok(addr) = operand.parse::<u64>() {
446        return Some(addr);
447    }
448
449    None
450}
451
452fn is_call_mnemonic(mnemonic: &str) -> bool {
453    matches!(
454        mnemonic,
455        "call" | "callq" | "lcall" | "bl" | "blr" | "blx" | "jal" | "jalr" | "bctrl"
456    )
457}
458
459fn is_return_mnemonic(mnemonic: &str) -> bool {
460    matches!(mnemonic, "ret" | "retn" | "retf" | "eret")
461}
462
463fn is_unconditional_branch_mnemonic(mnemonic: &str) -> bool {
464    matches!(
465        mnemonic,
466        "jmp" | "jmpq" | "ljmp" | "j" | "b" | "br" | "bx" | "jr"
467    )
468}
469
470fn is_conditional_branch_mnemonic(mnemonic: &str) -> bool {
471    (mnemonic.starts_with('j') && !is_unconditional_branch_mnemonic(mnemonic))
472        || mnemonic.starts_with("loop")
473        || mnemonic.starts_with("b.")
474        || matches!(
475            mnemonic,
476            "beq"
477                | "beqz"
478                | "bne"
479                | "bnez"
480                | "bgt"
481                | "bgtz"
482                | "bge"
483                | "bgez"
484                | "blt"
485                | "bltz"
486                | "ble"
487                | "blez"
488                | "bhi"
489                | "bhs"
490                | "blo"
491                | "bls"
492                | "bmi"
493                | "bpl"
494                | "bvs"
495                | "bvc"
496                | "cbz"
497                | "cbnz"
498                | "tbz"
499                | "tbnz"
500                | "bc"
501                | "bdnz"
502                | "bdz"
503        )
504}
505
506fn is_interrupt_mnemonic(mnemonic: &str) -> bool {
507    mnemonic.starts_with("int")
508        || matches!(
509            mnemonic,
510            "syscall" | "sysenter" | "sysexit" | "svc" | "swi" | "brk"
511        )
512}
513
514fn is_control_mnemonic(mnemonic: &str) -> bool {
515    is_call_mnemonic(mnemonic)
516        || is_return_mnemonic(mnemonic)
517        || is_unconditional_branch_mnemonic(mnemonic)
518        || is_conditional_branch_mnemonic(mnemonic)
519        || is_interrupt_mnemonic(mnemonic)
520}
521
522#[cfg(test)]
523mod tests {
524    use super::*;
525    use crate::types::*;
526
527    #[cfg(feature = "disasm-iced")]
528    struct TestFormat {
529        architecture: Architecture,
530        sections: Vec<Section>,
531        symbols: Vec<Symbol>,
532        entry_point: Option<u64>,
533        metadata: BinaryMetadata,
534    }
535
536    #[cfg(feature = "disasm-iced")]
537    impl BinaryFormatTrait for TestFormat {
538        fn format_type(&self) -> BinaryFormat {
539            BinaryFormat::Raw
540        }
541
542        fn architecture(&self) -> Architecture {
543            self.architecture
544        }
545
546        fn entry_point(&self) -> Option<u64> {
547            self.entry_point
548        }
549
550        fn sections(&self) -> &[Section] {
551            &self.sections
552        }
553
554        fn symbols(&self) -> &[Symbol] {
555            &self.symbols
556        }
557
558        fn imports(&self) -> &[Import] {
559            &[]
560        }
561
562        fn exports(&self) -> &[Export] {
563            &[]
564        }
565
566        fn metadata(&self) -> &BinaryMetadata {
567            &self.metadata
568        }
569    }
570
571    #[cfg(feature = "disasm-iced")]
572    fn test_binary(architecture: Architecture, addresses: &[u64]) -> BinaryFile {
573        let sections = addresses
574            .iter()
575            .enumerate()
576            .map(|(index, &address)| Section {
577                name: format!(".text{index}"),
578                address,
579                size: 1,
580                offset: index as u64,
581                file_size: 1,
582                permissions: SectionPermissions {
583                    read: true,
584                    write: false,
585                    execute: true,
586                },
587                section_type: SectionType::Code,
588                data: Some(vec![0x90]),
589            })
590            .collect();
591        let metadata = BinaryMetadata {
592            size: addresses.len(),
593            format: BinaryFormat::Raw,
594            architecture,
595            entry_point: None,
596            base_address: None,
597            timestamp: None,
598            compiler_info: None,
599            endian: Endianness::Little,
600            security_features: SecurityFeatures::default(),
601        };
602        BinaryFile {
603            data: vec![0x90; addresses.len()],
604            parsed: Box::new(TestFormat {
605                architecture,
606                sections,
607                symbols: Vec::new(),
608                entry_point: None,
609                metadata,
610            }),
611        }
612    }
613
614    #[cfg(all(feature = "disasm-iced", feature = "control-flow"))]
615    fn large_test_binary() -> (BinaryFile, u64, u64) {
616        const SECTION_OFFSET: usize = 128;
617        const SECTION_ADDRESS: u64 = 0x1000;
618        const SECTION_SIZE: usize = 2048;
619        const CALLER_OFFSET: usize = 1200;
620        const CALLEE_OFFSET: usize = 1500;
621
622        let caller = SECTION_ADDRESS + CALLER_OFFSET as u64;
623        let callee = SECTION_ADDRESS + CALLEE_OFFSET as u64;
624        let displacement = i32::try_from(callee - (caller + 5)).unwrap();
625        let mut data = vec![0x90; SECTION_OFFSET + SECTION_SIZE];
626        data[SECTION_OFFSET + CALLER_OFFSET] = 0xe8;
627        data[SECTION_OFFSET + CALLER_OFFSET + 1..SECTION_OFFSET + CALLER_OFFSET + 5]
628            .copy_from_slice(&displacement.to_le_bytes());
629        data[SECTION_OFFSET + CALLEE_OFFSET] = 0xc3;
630
631        let sections = vec![Section {
632            name: ".text".to_string(),
633            address: SECTION_ADDRESS,
634            size: SECTION_SIZE as u64,
635            offset: SECTION_OFFSET as u64,
636            file_size: SECTION_SIZE as u64,
637            permissions: SectionPermissions {
638                read: true,
639                write: false,
640                execute: true,
641            },
642            section_type: SectionType::Code,
643            data: None,
644        }];
645        let symbols = vec![
646            Symbol {
647                name: "caller".to_string(),
648                demangled_name: None,
649                address: caller,
650                size: 5,
651                symbol_type: SymbolType::Function,
652                binding: SymbolBinding::Global,
653                visibility: SymbolVisibility::Default,
654                section_index: Some(0),
655            },
656            Symbol {
657                name: "callee".to_string(),
658                demangled_name: None,
659                address: callee,
660                size: 1,
661                symbol_type: SymbolType::Function,
662                binding: SymbolBinding::Global,
663                visibility: SymbolVisibility::Default,
664                section_index: Some(0),
665            },
666        ];
667        let metadata = BinaryMetadata {
668            size: data.len(),
669            format: BinaryFormat::Raw,
670            architecture: Architecture::X86_64,
671            entry_point: Some(caller),
672            base_address: Some(SECTION_ADDRESS),
673            timestamp: None,
674            compiler_info: None,
675            endian: Endianness::Little,
676            security_features: SecurityFeatures::default(),
677        };
678        let binary = BinaryFile {
679            data,
680            parsed: Box::new(TestFormat {
681                architecture: Architecture::X86_64,
682                sections,
683                symbols,
684                entry_point: Some(caller),
685                metadata,
686            }),
687        };
688        (binary, caller, callee)
689    }
690
691    #[test]
692    fn test_disassembler_creation() {
693        let result = Disassembler::new(Architecture::X86_64);
694        assert!(result.is_ok());
695    }
696
697    fn assert_address_range_validation(engine: DisassemblyEngine) {
698        let disassembler = Disassembler::with_config(
699            Architecture::X86_64,
700            DisassemblyConfig {
701                engine,
702                ..DisassemblyConfig::default()
703            },
704        )
705        .unwrap();
706
707        let error = disassembler.disassemble(&[0x90], u64::MAX).unwrap_err();
708        assert!(
709            matches!(error, BinaryError::InvalidData(message) if message.contains("address range") && message.contains("overflows u64"))
710        );
711
712        let instructions = disassembler
713            .disassemble(&[0x90], u64::MAX - 1)
714            .expect("the largest representable end-exclusive range must be accepted");
715        assert_eq!(instructions.len(), 1);
716        assert_eq!(instructions[0].address, u64::MAX - 1);
717    }
718
719    #[cfg(feature = "disasm-capstone")]
720    #[test]
721    fn capstone_rejects_wrapping_address_ranges() {
722        assert_address_range_validation(DisassemblyEngine::Capstone);
723    }
724
725    #[cfg(feature = "disasm-iced")]
726    #[test]
727    fn iced_rejects_wrapping_address_ranges() {
728        assert_address_range_validation(DisassemblyEngine::Iced);
729    }
730
731    #[test]
732    fn test_config_default() {
733        let config = DisassemblyConfig::default();
734        assert_eq!(config.engine, DisassemblyEngine::Auto);
735        assert_eq!(config.max_instructions, 10000);
736        assert!(config.detailed);
737        assert!(config.analyze_control_flow);
738    }
739
740    #[test]
741    fn test_instruction_categorization() {
742        assert_eq!(
743            categorize_instruction("add"),
744            InstructionCategory::Arithmetic
745        );
746        assert_eq!(categorize_instruction("mov"), InstructionCategory::Memory);
747        assert_eq!(categorize_instruction("jmp"), InstructionCategory::Control);
748        assert_eq!(categorize_instruction("and"), InstructionCategory::Logic);
749        assert_eq!(
750            categorize_instruction("syscall"),
751            InstructionCategory::System
752        );
753        assert_ne!(
754            categorize_instruction("blendps"),
755            InstructionCategory::Control
756        );
757        assert_eq!(categorize_instruction("bl"), InstructionCategory::Control);
758    }
759
760    #[test]
761    #[cfg(feature = "disasm-capstone")]
762    fn test_control_flow_analysis() {
763        assert_eq!(analyze_control_flow("ret", ""), FlowType::Return);
764        assert_eq!(
765            analyze_control_flow("call", "0x1000"),
766            FlowType::Call(0x1000)
767        );
768        assert_eq!(
769            analyze_control_flow("jmp", "0x2000"),
770            FlowType::Jump(0x2000)
771        );
772        assert_eq!(
773            analyze_control_flow("je", "0x3000"),
774            FlowType::ConditionalJump(0x3000)
775        );
776        assert_eq!(
777            analyze_control_flow("mov", "eax, ebx"),
778            FlowType::Sequential
779        );
780        assert_eq!(
781            analyze_control_flow("bl", "#0x4000"),
782            FlowType::Call(0x4000)
783        );
784        assert_eq!(
785            analyze_control_flow("b.eq", "#0x5000"),
786            FlowType::ConditionalJump(0x5000)
787        );
788    }
789
790    #[test]
791    #[cfg(feature = "disasm-capstone")]
792    fn test_address_extraction() {
793        assert_eq!(extract_address_from_operands("0x1000"), Some(0x1000));
794        assert_eq!(extract_address_from_operands("4096"), Some(4096));
795        assert_eq!(extract_address_from_operands("#0x1000"), Some(0x1000));
796        assert_eq!(extract_address_from_operands("eax"), None);
797    }
798
799    #[cfg(feature = "disasm-iced")]
800    #[test]
801    fn binary_disassembly_returns_error_when_every_section_fails() {
802        let binary = test_binary(Architecture::Arm, &[0x1000]);
803        let config = AnalysisConfig {
804            disassembly_engine: DisassemblyEngine::Iced,
805            max_analysis_size: 16,
806            ..AnalysisConfig::default()
807        };
808
809        assert!(disassemble_binary(&binary, &config).is_err());
810    }
811
812    #[cfg(feature = "disasm-iced")]
813    #[test]
814    fn binary_disassembly_honors_architecture_hint_and_sorts_output() {
815        let binary = test_binary(Architecture::Arm, &[0x2000, 0x1000]);
816        let config = AnalysisConfig {
817            disassembly_engine: DisassemblyEngine::Iced,
818            architecture_hint: Some(Architecture::X86_64),
819            max_analysis_size: 16,
820            ..AnalysisConfig::default()
821        };
822
823        let instructions = disassemble_binary(&binary, &config).unwrap();
824
825        assert_eq!(
826            instructions
827                .iter()
828                .map(|instruction| instruction.address)
829                .collect::<Vec<_>>(),
830            vec![0x1000, 0x2000]
831        );
832    }
833
834    #[cfg(feature = "disasm-iced")]
835    #[test]
836    fn binary_disassembly_honors_global_instruction_limit() {
837        let addresses = (0..20).map(|index| 0x1000 + index).collect::<Vec<_>>();
838        let binary = test_binary(Architecture::X86_64, &addresses);
839        let config = AnalysisConfig {
840            disassembly_engine: DisassemblyEngine::Iced,
841            max_analysis_size: 1024,
842            max_disassembly_instructions: 7,
843            ..AnalysisConfig::default()
844        };
845
846        assert_eq!(disassemble_binary(&binary, &config).unwrap().len(), 7);
847    }
848
849    #[cfg(feature = "disasm-iced")]
850    #[test]
851    fn section_file_data_rejects_truncated_declared_range() {
852        let binary = test_binary(Architecture::X86_64, &[0x1000]);
853        let section = Section {
854            name: ".truncated".to_string(),
855            address: 0x1000,
856            size: 2,
857            offset: 0,
858            file_size: 2,
859            permissions: SectionPermissions {
860                read: true,
861                write: false,
862                execute: true,
863            },
864            section_type: SectionType::Code,
865            data: None,
866        };
867
868        assert!(section_file_data(&binary, &section, 0, None).is_err());
869    }
870
871    #[cfg(all(feature = "disasm-iced", feature = "control-flow"))]
872    #[test]
873    fn large_sections_use_the_binary_file_instead_of_inline_previews() {
874        use crate::analysis::{call_graph::CallGraphAnalyzer, control_flow::ControlFlowAnalyzer};
875
876        let (binary, caller, callee) = large_test_binary();
877        let disassembler = Disassembler::with_config(
878            Architecture::X86_64,
879            DisassemblyConfig {
880                engine: DisassemblyEngine::Iced,
881                max_instructions: 4096,
882                ..DisassemblyConfig::default()
883            },
884        )
885        .unwrap();
886
887        let section_instructions = disassembler.disassemble_section(&binary, ".text").unwrap();
888        assert!(section_instructions.iter().any(|instruction| {
889            instruction.address == caller && instruction.flow == FlowType::Call(callee)
890        }));
891
892        let cfgs = ControlFlowAnalyzer::new(Architecture::X86_64)
893            .analyze_binary(&binary)
894            .unwrap();
895        let caller_cfg = cfgs
896            .iter()
897            .find(|cfg| cfg.function.start_address == caller)
898            .unwrap();
899        assert_eq!(caller_cfg.basic_blocks[0].instructions[0].address, caller);
900
901        let call_graph = CallGraphAnalyzer::new_default()
902            .analyze_binary(&binary)
903            .unwrap();
904        assert!(
905            call_graph
906                .edges
907                .iter()
908                .any(|edge| edge.caller == caller && edge.callee == Some(callee))
909        );
910    }
911}