use crate::{arch_info, Error, Result};
#[cfg(feature = "serde")]
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::{convert::TryInto, fmt};
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ArchitectureIdentifier {
Amd64,
Arm64,
Virtual,
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug)]
pub struct Header {
pub arch_id: ArchitectureIdentifier,
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone, Copy)]
#[repr(transparent)]
pub struct Vip(pub u64);
bitflags! {
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct RegisterFlags: u64 {
const VIRTUAL = 0;
const PHYSICAL = 1 << 0;
const LOCAL = 1 << 1;
const FLAGS = 1 << 2;
const STACK_POINTER = 1 << 3;
const IMAGE_BASE = 1 << 4;
const VOLATILE = 1 << 5;
const READONLY = 1 << 6;
const UNDEFINED = 1 << 7;
const INTERNAL = 1 << 8;
const SPECIAL = Self::FLAGS.bits | Self::STACK_POINTER.bits | Self::IMAGE_BASE.bits | Self::UNDEFINED.bits;
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone, Copy)]
pub struct Reg {
pub flags: RegisterFlags,
pub combined_id: u64,
pub bit_count: i32,
pub bit_offset: i32,
}
impl Reg {
pub fn local_id(&self) -> u64 {
self.combined_id & !(0xff << 56)
}
pub fn arch_id(&self) -> ArchitectureIdentifier {
match self.combined_id & (0xff << 56) {
0 => ArchitectureIdentifier::Amd64,
1 => ArchitectureIdentifier::Arm64,
2 => ArchitectureIdentifier::Virtual,
_ => unreachable!(),
}
}
}
impl fmt::Display for Reg {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let mut prefix = String::new();
if self.flags.contains(RegisterFlags::VOLATILE) {
prefix = "?".to_string();
}
if self.flags.contains(RegisterFlags::READONLY) {
prefix += "&&";
}
let mut suffix = String::new();
if self.bit_offset != 0 {
suffix = format!("@{}", self.bit_offset);
}
if self.bit_count != 64 {
suffix.push_str(&format!(":{}", self.bit_count));
}
if self.flags.contains(RegisterFlags::INTERNAL) {
write!(f, "{}sr{}{}", prefix, self.local_id(), suffix)?;
return Ok(());
} else if self.flags.contains(RegisterFlags::UNDEFINED) {
write!(f, "{}UD{}", prefix, suffix)?;
return Ok(());
} else if self.flags.contains(RegisterFlags::FLAGS) {
write!(f, "{}$flags{}", prefix, suffix)?;
return Ok(());
} else if self.flags.contains(RegisterFlags::STACK_POINTER) {
write!(f, "{}$sp{}", prefix, suffix)?;
return Ok(());
} else if self.flags.contains(RegisterFlags::IMAGE_BASE) {
write!(f, "{}base{}", prefix, suffix)?;
return Ok(());
} else if self.flags.contains(RegisterFlags::LOCAL) {
write!(f, "{}t{}{}", prefix, self.local_id(), suffix)?;
return Ok(());
}
if self.flags.contains(RegisterFlags::PHYSICAL) {
match self.arch_id() {
ArchitectureIdentifier::Amd64 => {
write!(
f,
"{}{}{}",
prefix,
arch_info::X86_REGISTER_NAME_MAPPING[self.local_id() as usize],
suffix
)?;
return Ok(());
}
ArchitectureIdentifier::Arm64 => {
write!(
f,
"{}{}{}",
prefix,
arch_info::AARCH64_REGISTER_NAME_MAPPING[self.local_id() as usize],
suffix
)?;
return Ok(());
}
_ => {}
}
}
write!(f, "{}vr{}{}", prefix, self.local_id(), suffix)?;
Ok(())
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug)]
pub struct RoutineConvention {
pub volatile_registers: Vec<Reg>,
pub param_registers: Vec<Reg>,
pub retval_registers: Vec<Reg>,
pub frame_register: Reg,
pub shadow_space: u64,
pub purge_stack: bool,
}
#[derive(Clone, Copy)]
pub(crate) union Immediate {
pub(crate) u64: u64,
pub(crate) i64: i64,
}
impl Immediate {
pub(crate) fn u64(&self) -> u64 {
unsafe { self.u64 }
}
pub(crate) fn set_u64(&mut self, imm: u64) {
self.u64 = imm;
}
pub(crate) fn i64(&self) -> i64 {
unsafe { self.i64 }
}
pub(crate) fn set_i64(&mut self, imm: i64) {
self.i64 = imm;
}
}
#[cfg(feature = "serde")]
impl Serialize for Immediate {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_i64(self.i64())
}
}
#[cfg(feature = "serde")]
impl<'de> Deserialize<'de> for Immediate {
fn deserialize<D>(deserializer: D) -> std::result::Result<Immediate, D::Error>
where
D: Deserializer<'de>,
{
Ok(Immediate {
i64: i64::deserialize(deserializer)?,
})
}
}
impl fmt::Debug for Immediate {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Immediate")
.field("u64", &self.u64())
.field("i64", &self.i64())
.finish()
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone, Copy)]
pub struct Imm {
pub(crate) value: Immediate,
pub bit_count: u32,
}
impl Imm {
pub fn new<T: Into<u64>>(value: T, bit_count: u32) -> Imm {
assert!(bit_count % 8 == 0);
Imm {
value: Immediate { u64: value.into() },
bit_count,
}
}
pub fn new_signed<T: Into<i64>>(value: T, bit_count: u32) -> Imm {
assert!(bit_count % 8 == 0);
Imm {
value: Immediate { i64: value.into() },
bit_count,
}
}
pub fn u64(&self) -> u64 {
self.value.u64()
}
pub fn set_u64(&mut self, imm: u64) {
self.value.set_u64(imm);
}
pub fn i64(&self) -> i64 {
self.value.i64()
}
pub fn set_i64(&mut self, imm: i64) {
self.value.set_i64(imm);
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone, Copy)]
pub enum Operand {
Imm(Imm),
Reg(Reg),
}
impl<'a, 'b> TryInto<&'b Imm> for &'a Operand
where
'a: 'b,
{
type Error = Error;
fn try_into(self) -> Result<&'a Imm> {
match self {
Operand::Imm(ref i) => Ok(i),
_ => Err(Error::OperandTypeMismatch),
}
}
}
impl<'a, 'b> TryInto<&'b Reg> for &'a Operand
where
'a: 'b,
{
type Error = Error;
fn try_into(self) -> Result<&'a Reg> {
match self {
Operand::Reg(r) => Ok(r),
_ => Err(Error::OperandTypeMismatch),
}
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug)]
pub struct Instruction {
pub op: Op,
pub vip: Vip,
pub sp_offset: i64,
pub sp_index: u32,
pub sp_reset: bool,
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug)]
pub enum Op {
Mov(Operand, Operand),
Movsx(Operand, Operand),
Str(Operand, Operand, Operand),
Ldd(Operand, Operand, Operand),
Neg(Operand),
Add(Operand, Operand),
Sub(Operand, Operand),
Mul(Operand, Operand),
Mulhi(Operand, Operand),
Imul(Operand, Operand),
Imulhi(Operand, Operand),
Div(Operand, Operand, Operand),
Rem(Operand, Operand, Operand),
Idiv(Operand, Operand, Operand),
Irem(Operand, Operand, Operand),
Popcnt(Operand),
Bsf(Operand),
Bsr(Operand),
Not(Operand),
Shr(Operand, Operand),
Shl(Operand, Operand),
Xor(Operand, Operand),
Or(Operand, Operand),
And(Operand, Operand),
Ror(Operand, Operand),
Rol(Operand, Operand),
Tg(Operand, Operand, Operand),
Tge(Operand, Operand, Operand),
Te(Operand, Operand, Operand),
Tne(Operand, Operand, Operand),
Tl(Operand, Operand, Operand),
Tle(Operand, Operand, Operand),
Tug(Operand, Operand, Operand),
Tuge(Operand, Operand, Operand),
Tul(Operand, Operand, Operand),
Tule(Operand, Operand, Operand),
Ifs(Operand, Operand, Operand),
Js(Operand, Operand, Operand),
Jmp(Operand),
Vexit(Operand),
Vxcall(Operand),
Nop,
Sfence,
Lfence,
Vemit(Operand),
Vpinr(Operand),
Vpinw(Operand),
Vpinrm(Operand, Operand, Operand),
Vpinwm(Operand, Operand, Operand),
}
impl Op {
pub fn name(&self) -> &'static str {
match self {
Op::Mov(_, _) => "mov",
Op::Movsx(_, _) => "movsx",
Op::Str(_, _, _) => "str",
Op::Ldd(_, _, _) => "ldd",
Op::Neg(_) => "neg",
Op::Add(_, _) => "add",
Op::Sub(_, _) => "sub",
Op::Mul(_, _) => "mul",
Op::Mulhi(_, _) => "mulhi",
Op::Imul(_, _) => "imul",
Op::Imulhi(_, _) => "imulhi",
Op::Div(_, _, _) => "div",
Op::Rem(_, _, _) => "rem",
Op::Idiv(_, _, _) => "idiv",
Op::Irem(_, _, _) => "irem",
Op::Popcnt(_) => "popcnt",
Op::Bsf(_) => "bsf",
Op::Bsr(_) => "bsr",
Op::Not(_) => "not",
Op::Shr(_, _) => "shr",
Op::Shl(_, _) => "shl",
Op::Xor(_, _) => "xor",
Op::Or(_, _) => "or",
Op::And(_, _) => "and",
Op::Ror(_, _) => "ror",
Op::Rol(_, _) => "rol",
Op::Tg(_, _, _) => "tg",
Op::Tge(_, _, _) => "tge",
Op::Te(_, _, _) => "te",
Op::Tne(_, _, _) => "tne",
Op::Tl(_, _, _) => "tl",
Op::Tle(_, _, _) => "tle",
Op::Tug(_, _, _) => "tug",
Op::Tuge(_, _, _) => "tuge",
Op::Tul(_, _, _) => "tul",
Op::Tule(_, _, _) => "tule",
Op::Ifs(_, _, _) => "ifs",
Op::Js(_, _, _) => "js",
Op::Jmp(_) => "jmp",
Op::Vexit(_) => "vexit",
Op::Vxcall(_) => "vxcall",
Op::Nop => "nop",
Op::Sfence => "sfence",
Op::Lfence => "lfence",
Op::Vemit(_) => "vemit",
Op::Vpinr(_) => "vpinr",
Op::Vpinw(_) => "vpinw",
Op::Vpinrm(_, _, _) => "vpinrm",
Op::Vpinwm(_, _, _) => "vpinwm",
}
}
pub fn operands(&self) -> Vec<&Operand> {
match *self {
Op::Nop | Op::Sfence | Op::Lfence => vec![],
Op::Neg(ref op1)
| Op::Popcnt(ref op1)
| Op::Bsf(ref op1)
| Op::Bsr(ref op1)
| Op::Not(ref op1)
| Op::Jmp(ref op1)
| Op::Vexit(ref op1)
| Op::Vxcall(ref op1)
| Op::Vemit(ref op1)
| Op::Vpinr(ref op1)
| Op::Vpinw(ref op1) => vec![op1],
Op::Mov(ref op1, ref op2)
| Op::Movsx(ref op1, ref op2)
| Op::Add(ref op1, ref op2)
| Op::Sub(ref op1, ref op2)
| Op::Mul(ref op1, ref op2)
| Op::Mulhi(ref op1, ref op2)
| Op::Imul(ref op1, ref op2)
| Op::Imulhi(ref op1, ref op2)
| Op::Shr(ref op1, ref op2)
| Op::Shl(ref op1, ref op2)
| Op::Xor(ref op1, ref op2)
| Op::Or(ref op1, ref op2)
| Op::And(ref op1, ref op2)
| Op::Ror(ref op1, ref op2)
| Op::Rol(ref op1, ref op2) => vec![op1, op2],
Op::Str(ref op1, ref op2, ref op3)
| Op::Ldd(ref op1, ref op2, ref op3)
| Op::Div(ref op1, ref op2, ref op3)
| Op::Rem(ref op1, ref op2, ref op3)
| Op::Idiv(ref op1, ref op2, ref op3)
| Op::Irem(ref op1, ref op2, ref op3)
| Op::Tg(ref op1, ref op2, ref op3)
| Op::Tge(ref op1, ref op2, ref op3)
| Op::Te(ref op1, ref op2, ref op3)
| Op::Tne(ref op1, ref op2, ref op3)
| Op::Tl(ref op1, ref op2, ref op3)
| Op::Tle(ref op1, ref op2, ref op3)
| Op::Tug(ref op1, ref op2, ref op3)
| Op::Tuge(ref op1, ref op2, ref op3)
| Op::Tul(ref op1, ref op2, ref op3)
| Op::Tule(ref op1, ref op2, ref op3)
| Op::Ifs(ref op1, ref op2, ref op3)
| Op::Js(ref op1, ref op2, ref op3)
| Op::Vpinrm(ref op1, ref op2, ref op3)
| Op::Vpinwm(ref op1, ref op2, ref op3) => vec![op1, op2, op3],
}
}
pub fn operands_mut(&mut self) -> Vec<&mut Operand> {
match *self {
Op::Nop | Op::Sfence | Op::Lfence => vec![],
Op::Neg(ref mut op1)
| Op::Popcnt(ref mut op1)
| Op::Bsf(ref mut op1)
| Op::Bsr(ref mut op1)
| Op::Not(ref mut op1)
| Op::Jmp(ref mut op1)
| Op::Vexit(ref mut op1)
| Op::Vxcall(ref mut op1)
| Op::Vemit(ref mut op1)
| Op::Vpinr(ref mut op1)
| Op::Vpinw(ref mut op1) => vec![op1],
Op::Mov(ref mut op1, ref mut op2)
| Op::Movsx(ref mut op1, ref mut op2)
| Op::Add(ref mut op1, ref mut op2)
| Op::Sub(ref mut op1, ref mut op2)
| Op::Mul(ref mut op1, ref mut op2)
| Op::Mulhi(ref mut op1, ref mut op2)
| Op::Imul(ref mut op1, ref mut op2)
| Op::Imulhi(ref mut op1, ref mut op2)
| Op::Shr(ref mut op1, ref mut op2)
| Op::Shl(ref mut op1, ref mut op2)
| Op::Xor(ref mut op1, ref mut op2)
| Op::Or(ref mut op1, ref mut op2)
| Op::And(ref mut op1, ref mut op2)
| Op::Ror(ref mut op1, ref mut op2)
| Op::Rol(ref mut op1, ref mut op2) => vec![op1, op2],
Op::Str(ref mut op1, ref mut op2, ref mut op3)
| Op::Ldd(ref mut op1, ref mut op2, ref mut op3)
| Op::Div(ref mut op1, ref mut op2, ref mut op3)
| Op::Rem(ref mut op1, ref mut op2, ref mut op3)
| Op::Idiv(ref mut op1, ref mut op2, ref mut op3)
| Op::Irem(ref mut op1, ref mut op2, ref mut op3)
| Op::Tg(ref mut op1, ref mut op2, ref mut op3)
| Op::Tge(ref mut op1, ref mut op2, ref mut op3)
| Op::Te(ref mut op1, ref mut op2, ref mut op3)
| Op::Tne(ref mut op1, ref mut op2, ref mut op3)
| Op::Tl(ref mut op1, ref mut op2, ref mut op3)
| Op::Tle(ref mut op1, ref mut op2, ref mut op3)
| Op::Tug(ref mut op1, ref mut op2, ref mut op3)
| Op::Tuge(ref mut op1, ref mut op2, ref mut op3)
| Op::Tul(ref mut op1, ref mut op2, ref mut op3)
| Op::Tule(ref mut op1, ref mut op2, ref mut op3)
| Op::Ifs(ref mut op1, ref mut op2, ref mut op3)
| Op::Js(ref mut op1, ref mut op2, ref mut op3)
| Op::Vpinrm(ref mut op1, ref mut op2, ref mut op3)
| Op::Vpinwm(ref mut op1, ref mut op2, ref mut op3) => vec![op1, op2, op3],
}
}
pub fn is_volatile(&self) -> bool {
matches!(
self,
Op::Sfence
| Op::Lfence
| Op::Vemit(_)
| Op::Vpinr(_)
| Op::Vpinw(_)
| Op::Vpinrm(_, _, _)
| Op::Vpinwm(_, _, _)
)
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug)]
pub struct BasicBlock {
pub vip: Vip,
pub sp_offset: i64,
pub sp_index: u32,
pub last_temporary_index: u32,
pub instructions: Vec<Instruction>,
pub prev_vip: Vec<Vip>,
pub next_vip: Vec<Vip>,
}
pub type SubroutineConvention = RoutineConvention;
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug)]
pub struct VTIL {
pub header: Header,
pub vip: Vip,
pub routine_convention: RoutineConvention,
pub subroutine_convention: SubroutineConvention,
pub spec_subroutine_conventions: Vec<SubroutineConvention>,
pub explored_blocks: Vec<BasicBlock>,
}