use anyhow::Result;
use smallvec::{smallvec, SmallVec};
use crate::{
analysis::{dis, dis::Target},
module::{Module, Permissions},
VA,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Flow {
Fallthrough(VA),
Call(Target),
UnconditionalJump(Target),
ConditionalJump(VA),
}
impl Flow {
#[must_use]
pub fn swap(&self, va: VA) -> Flow {
match *self {
Flow::Fallthrough(_) => Flow::Fallthrough(va),
Flow::Call(Target::Direct(_)) => Flow::Call(Target::Direct(va)),
Flow::Call(Target::Indirect(_)) => Flow::Call(Target::Indirect(va)),
Flow::UnconditionalJump(Target::Direct(_)) => Flow::UnconditionalJump(Target::Direct(va)),
Flow::UnconditionalJump(Target::Indirect(_)) => Flow::UnconditionalJump(Target::Indirect(va)),
Flow::ConditionalJump(_) => Flow::ConditionalJump(va),
}
}
}
pub type Flows = SmallVec<[Flow; 2]>;
fn is_executable(module: &Module, va: VA) -> bool {
module.probe_va(va, Permissions::X)
}
pub fn get_call_insn_flow(module: &Module, va: VA, insn: &zydis::DecodedInstruction) -> Result<Flows> {
let op = dis::get_first_operand(insn).expect("CALL has no operand");
if let Ok(Some(target)) = dis::get_operand_xref(module, va, insn, op) {
match target {
Target::Direct(va) => {
if is_executable(module, va) {
return Ok(smallvec![Flow::Call(target)]);
} else {
return Ok(smallvec![]);
}
}
Target::Indirect(_) => {
return Ok(smallvec![Flow::Call(target)]);
}
}
}
Ok(smallvec![])
}
pub fn get_jmp_insn_flow(module: &Module, va: VA, insn: &zydis::DecodedInstruction) -> Result<Flows> {
let op = dis::get_first_operand(insn).expect("JMP has no target");
if let Ok(Some(target)) = dis::get_operand_xref(module, va, insn, op) {
match target {
Target::Direct(va) => {
if is_executable(module, va) {
return Ok(smallvec![Flow::UnconditionalJump(target)]);
} else {
return Ok(smallvec![]);
}
}
Target::Indirect(_) => {
return Ok(smallvec![Flow::UnconditionalJump(target)]);
}
}
}
Ok(smallvec![])
}
pub fn get_cjmp_insn_flow(module: &Module, va: VA, insn: &zydis::DecodedInstruction) -> Result<Flows> {
let op = dis::get_first_operand(insn).expect("CJMP has no target");
if let Ok(Some(Target::Direct(dst))) = dis::get_operand_xref(module, va, insn, op) {
if is_executable(module, dst) {
return Ok(smallvec![Flow::ConditionalJump(dst)]);
} else {
return Ok(smallvec![]);
}
}
Ok(smallvec![])
}
pub fn get_insn_flow(module: &Module, va: VA, insn: &zydis::DecodedInstruction) -> Result<Flows> {
let mut flows = match insn.mnemonic {
zydis::Mnemonic::CALL => get_call_insn_flow(module, va, insn)?,
zydis::Mnemonic::JMP => get_jmp_insn_flow(module, va, insn)?,
zydis::Mnemonic::RET | zydis::Mnemonic::IRET | zydis::Mnemonic::IRETD | zydis::Mnemonic::IRETQ => smallvec![],
zydis::Mnemonic::JB
| zydis::Mnemonic::JBE
| zydis::Mnemonic::JCXZ
| zydis::Mnemonic::JECXZ
| zydis::Mnemonic::JKNZD
| zydis::Mnemonic::JKZD
| zydis::Mnemonic::JL
| zydis::Mnemonic::JLE
| zydis::Mnemonic::JNB
| zydis::Mnemonic::JNBE
| zydis::Mnemonic::JNL
| zydis::Mnemonic::JNLE
| zydis::Mnemonic::JNO
| zydis::Mnemonic::JNP
| zydis::Mnemonic::JNS
| zydis::Mnemonic::JNZ
| zydis::Mnemonic::JO
| zydis::Mnemonic::JP
| zydis::Mnemonic::JRCXZ
| zydis::Mnemonic::JS
| zydis::Mnemonic::JZ => get_cjmp_insn_flow(module, va, insn)?,
zydis::Mnemonic::CMOVB
| zydis::Mnemonic::CMOVBE
| zydis::Mnemonic::CMOVL
| zydis::Mnemonic::CMOVLE
| zydis::Mnemonic::CMOVNB
| zydis::Mnemonic::CMOVNBE
| zydis::Mnemonic::CMOVNL
| zydis::Mnemonic::CMOVNLE
| zydis::Mnemonic::CMOVNO
| zydis::Mnemonic::CMOVNP
| zydis::Mnemonic::CMOVNS
| zydis::Mnemonic::CMOVNZ
| zydis::Mnemonic::CMOVO
| zydis::Mnemonic::CMOVP
| zydis::Mnemonic::CMOVS
| zydis::Mnemonic::CMOVZ => smallvec![],
_ => smallvec![],
};
if dis::does_insn_fallthrough(insn) {
flows.push(Flow::Fallthrough(va + insn.length as u64))
}
Ok(flows)
}
#[cfg(test)]
mod tests {
use crate::{analysis::cfg::flow::*, test::*};
#[test]
fn test_get_call_insn_flow() {
let module = load_shellcode32(b"\xE8\x00\x00\x00\x00\x90");
let insn = read_insn(&module, 0x0);
let flows = get_call_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 1);
assert_eq!(flows[0], Flow::Call(Target::Direct(0x5)));
let flows = get_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 2);
assert_eq!(flows[0], Flow::Call(Target::Direct(0x5)));
assert_eq!(flows[1], Flow::Fallthrough(0x5));
}
#[test]
fn test_get_jmp_insn_flow() {
let module = load_shellcode32(b"\xE9\x00\x00\x00\x00\x90");
let insn = read_insn(&module, 0x0);
let flows = get_jmp_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 1);
assert_eq!(flows[0], Flow::UnconditionalJump(Target::Direct(0x5)));
let flows = get_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 1);
assert_eq!(flows[0], Flow::UnconditionalJump(Target::Direct(0x5)));
}
#[test]
fn test_get_cjmp_insn_flow() {
let module = load_shellcode32(b"\x75\x01\xCC\x90");
let insn = read_insn(&module, 0x0);
let flows = get_cjmp_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 1);
assert_eq!(flows[0], Flow::ConditionalJump(0x3));
let flows = get_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 2);
assert_eq!(flows[0], Flow::ConditionalJump(0x3));
assert_eq!(flows[1], Flow::Fallthrough(0x2));
}
#[test]
fn test_get_cmov_insn_flow() {
let module = load_shellcode32(b"\x0F\x44\xC3\x90");
let insn = read_insn(&module, 0x0);
let flows = get_insn_flow(&module, 0x0, &insn).unwrap();
assert_eq!(flows.len(), 1);
assert_eq!(flows[0], Flow::Fallthrough(0x3));
}
}