diff --git a/rust/src/low_level_il.rs b/rust/src/low_level_il.rs index bb4614ae21..78d4ce9c3b 100644 --- a/rust/src/low_level_il.rs +++ b/rust/src/low_level_il.rs @@ -14,7 +14,7 @@ use std::borrow::Cow; use std::fmt; - +use std::fmt::{Debug, Display}; // TODO : provide some way to forbid emitting register reads for certain registers // also writing for certain registers (e.g. zero register must prohibit il.set_reg and il.reg // (replace with nop or const(0) respectively) @@ -83,12 +83,18 @@ impl LowLevelILTempRegister { } } -impl fmt::Debug for LowLevelILTempRegister { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { +impl Debug for LowLevelILTempRegister { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "temp{}", self.temp_id) } } +impl Display for LowLevelILTempRegister { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Debug::fmt(self, f) + } +} + impl TryFrom for LowLevelILTempRegister { type Error = (); @@ -142,11 +148,20 @@ impl LowLevelILRegisterKind { } } -impl fmt::Debug for LowLevelILRegisterKind { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { +impl Debug for LowLevelILRegisterKind { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { match *self { LowLevelILRegisterKind::Arch(ref r) => r.fmt(f), - LowLevelILRegisterKind::Temp(id) => id.fmt(f), + LowLevelILRegisterKind::Temp(ref id) => Debug::fmt(id, f), + } + } +} + +impl Display for LowLevelILRegisterKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match *self { + LowLevelILRegisterKind::Arch(ref r) => write!(f, "{}", r.name()), + LowLevelILRegisterKind::Temp(ref id) => Display::fmt(id, f), } } } @@ -157,36 +172,128 @@ impl From for LowLevelILRegisterKind { } } +#[derive(Copy, Clone, Debug)] +pub struct LowLevelILSSARegister { + pub reg: LowLevelILRegisterKind, + /// The SSA version of the register. + pub version: u32, +} + +impl LowLevelILSSARegister { + pub fn new(reg: LowLevelILRegisterKind, version: u32) -> Self { + Self { reg, version } + } + + pub fn name(&self) -> Cow<'_, str> { + self.reg.name() + } + + pub fn id(&self) -> RegisterId { + self.reg.id() + } +} + +impl Display for LowLevelILSSARegister { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}#{}", self.reg, self.version) + } +} + +/// The kind of SSA register. +/// +/// An SSA register can exist in two states: +/// +/// - Full, e.g. `eax` on x86 +/// - Partial, e.g. `al` on x86 +/// +/// If you intend to query for the ssa uses or definitions you must retrieve the physical register +/// using the function [`LowLevelILSSARegisterKind::physical_reg`] which will give you the actual +/// [`LowLevelILSSARegister`]. #[derive(Copy, Clone, Debug)] pub enum LowLevelILSSARegisterKind { - Full { - kind: LowLevelILRegisterKind, - version: u32, - }, + /// A full register is one that is not aliasing another, such as `eax` on x86 or `rax` on x86_64. + Full(LowLevelILSSARegister), Partial { - full_reg: CoreRegister, + /// This is the non-aliased register. + /// + /// This register is what is used for dataflow, otherwise the backing storage of aliased registers + /// like `al` on x86 would contain separate value information from the physical register `eax`. + /// + /// NOTE: While this is a [`LowLevelILSSARegister`] temporary registers are not allowed in partial + /// assignments, so this will always be an actual architecture register. + full_reg: LowLevelILSSARegister, + /// This is the aliased register. + /// + /// On x86 if the register `al` is used that would be considered a partial register, with the + /// full register `eax` being used as the backing storage. partial_reg: CoreRegister, - version: u32, }, } impl LowLevelILSSARegisterKind { pub fn new_full(kind: LowLevelILRegisterKind, version: u32) -> Self { - Self::Full { kind, version } + Self::Full(LowLevelILSSARegister::new(kind, version)) } - pub fn new_partial(full_reg: CoreRegister, partial_reg: CoreRegister, version: u32) -> Self { + pub fn new_partial( + full_reg: LowLevelILRegisterKind, + version: u32, + partial_reg: CoreRegister, + ) -> Self { Self::Partial { - full_reg, + full_reg: LowLevelILSSARegister::new(full_reg, version), partial_reg, - version, } } - pub fn version(&self) -> u32 { + /// This is the non-aliased register used. This should be called when you intend to actually + /// query for SSA dataflow information, as a partial register is prohibited from being used. + /// + /// # Example + /// + /// On x86 `al` in the LLIL SSA will have a physical register of `eax`. + pub fn physical_reg(&self) -> LowLevelILSSARegister { match *self { - LowLevelILSSARegisterKind::Full { version, .. } - | LowLevelILSSARegisterKind::Partial { version, .. } => version, + LowLevelILSSARegisterKind::Full(reg) => reg, + LowLevelILSSARegisterKind::Partial { full_reg, .. } => full_reg, + } + } + + /// Gets the displayable register, for partial this will be the partial register name. + /// + /// # Example + /// + /// On x86 this will display "al" not "eax". + pub fn name(&self) -> Cow<'_, str> { + match *self { + LowLevelILSSARegisterKind::Full(ref reg) => reg.reg.name(), + LowLevelILSSARegisterKind::Partial { + ref partial_reg, .. + } => partial_reg.name(), + } + } +} + +impl AsRef> for LowLevelILSSARegisterKind { + fn as_ref(&self) -> &LowLevelILSSARegister { + match self { + LowLevelILSSARegisterKind::Full(reg) => reg, + LowLevelILSSARegisterKind::Partial { full_reg, .. } => full_reg, + } + } +} + +impl From> for LowLevelILSSARegisterKind { + fn from(value: LowLevelILSSARegister) -> Self { + LowLevelILSSARegisterKind::Full(value) + } +} + +impl From> for LowLevelILSSARegister { + fn from(value: LowLevelILSSARegisterKind) -> Self { + match value { + LowLevelILSSARegisterKind::Full(reg) => reg, + LowLevelILSSARegisterKind::Partial { full_reg, .. } => full_reg, } } } diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs index 15a83c19e7..b5661e32b6 100644 --- a/rust/src/low_level_il/function.rs +++ b/rust/src/low_level_il/function.rs @@ -264,19 +264,16 @@ impl LowLevelILFunction { #[must_use] pub fn get_ssa_register_uses( &self, - reg: LowLevelILSSARegisterKind, + reg: impl AsRef>, ) -> Vec> { use binaryninjacore_sys::BNGetLowLevelILSSARegisterUses; - let register_id = match reg { - LowLevelILSSARegisterKind::Full { kind, .. } => kind.id(), - LowLevelILSSARegisterKind::Partial { partial_reg, .. } => partial_reg.id(), - }; + let reg = reg.as_ref(); let mut count = 0; let instrs = unsafe { BNGetLowLevelILSSARegisterUses( self.handle, - register_id.into(), - reg.version() as usize, + reg.id().into(), + reg.version as usize, &mut count, ) }; @@ -292,19 +289,12 @@ impl LowLevelILFunction { #[must_use] pub fn get_ssa_register_definition( &self, - reg: &LowLevelILSSARegisterKind, + reg: impl AsRef>, ) -> Option> { use binaryninjacore_sys::BNGetLowLevelILSSARegisterDefinition; - let register_id = match reg { - LowLevelILSSARegisterKind::Full { kind, .. } => kind.id(), - LowLevelILSSARegisterKind::Partial { partial_reg, .. } => partial_reg.id(), - }; + let reg = reg.as_ref(); let instr_idx = unsafe { - BNGetLowLevelILSSARegisterDefinition( - self.handle, - register_id.into(), - reg.version() as usize, - ) + BNGetLowLevelILSSARegisterDefinition(self.handle, reg.id().into(), reg.version as usize) }; self.instruction_from_index(LowLevelInstructionIndex(instr_idx)) } @@ -313,14 +303,11 @@ impl LowLevelILFunction { #[must_use] pub fn get_ssa_register_value( &self, - reg: &LowLevelILSSARegisterKind, + reg: impl AsRef>, ) -> Option { - let register_id = match reg { - LowLevelILSSARegisterKind::Full { kind, .. } => kind.id(), - LowLevelILSSARegisterKind::Partial { partial_reg, .. } => partial_reg.id(), - }; + let reg = reg.as_ref(); let value = unsafe { - BNGetLowLevelILSSARegisterValue(self.handle, register_id.into(), reg.version() as usize) + BNGetLowLevelILSSARegisterValue(self.handle, reg.id().into(), reg.version as usize) }; if value.state == BNRegisterValueType::UndeterminedValue { return None; diff --git a/rust/src/low_level_il/operation.rs b/rust/src/low_level_il/operation.rs index fa320060fa..2f081b5e2b 100644 --- a/rust/src/low_level_il/operation.rs +++ b/rust/src/low_level_il/operation.rs @@ -413,11 +413,11 @@ where let full_raw_id = RegisterId(self.op.operands[0] as u32); let version = self.op.operands[1] as u32; let partial_raw_id = RegisterId(self.op.operands[2] as u32); - let full_reg = - CoreRegister::new(self.function.arch(), full_raw_id).expect("Bad register ID"); + let full_reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), full_raw_id) + .expect("Bad register ID"); let partial_reg = CoreRegister::new(self.function.arch(), partial_raw_id).expect("Bad register ID"); - LowLevelILSSARegisterKind::new_partial(full_reg, partial_reg, version) + LowLevelILSSARegisterKind::new_partial(full_reg_kind, version, partial_reg) } pub fn source_expr(&self) -> LowLevelILExpression<'func, M, F, ValueExpr> { @@ -868,11 +868,11 @@ where let full_raw_id = RegisterId(self.op.operands[0] as u32); let version = self.op.operands[1] as u32; let partial_raw_id = RegisterId(self.op.operands[2] as u32); - let full_reg = - CoreRegister::new(self.function.arch(), full_raw_id).expect("Bad register ID"); + let full_reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), full_raw_id) + .expect("Bad register ID"); let partial_reg = CoreRegister::new(self.function.arch(), partial_raw_id).expect("Bad register ID"); - LowLevelILSSARegisterKind::new_partial(full_reg, partial_reg, version) + LowLevelILSSARegisterKind::new_partial(full_reg_kind, version, partial_reg) } } diff --git a/rust/tests/low_level_il.rs b/rust/tests/low_level_il.rs index 79a7996487..e6f8d8ff3c 100644 --- a/rust/tests/low_level_il.rs +++ b/rust/tests/low_level_il.rs @@ -259,13 +259,19 @@ fn test_llil_ssa() { LowLevelILInstructionKind::SetRegSsa(op) => { assert_eq!(op.size(), 4); match op.dest_reg() { - LowLevelILSSARegisterKind::Full { kind, version } => { - assert_eq!(kind.name(), "edi"); - assert_eq!(version, 1); + LowLevelILSSARegisterKind::Full(reg) => { + assert_eq!(reg.name(), "edi"); + assert_eq!(reg.version, 1); } _ => panic!("Expected LowLevelILSSARegisterKind::Full"), } assert_eq!(op.source_expr().index, LowLevelExpressionIndex(0)); + + // Verify dest_reg does not have a use, so let's verify the ssa register definition. + let dest_reg_def = llil_ssa_function + .get_ssa_register_definition(op.dest_reg()) + .expect("Valid ssa reg def"); + assert_eq!(dest_reg_def.address(), ssa_instr_0.address()); } _ => panic!("Expected SetRegSsa"), } @@ -287,6 +293,23 @@ fn test_llil_ssa() { let dest_memory_version = op.dest_memory_version(); assert_eq!(dest_memory_version, 1); assert_eq!(dest_expr.index, LowLevelExpressionIndex(4)); + + // Grab the SP register so we can verify its use. + let dest_expr_kind = dest_expr.kind(); + let sub_expr = dest_expr_kind.as_binary_op().unwrap(); + match sub_expr.left().kind() { + LowLevelILExpressionKind::RegSsa(reg) => { + // Verify esp#0 has a single use in the next instruction (same address however). + let sp_0_uses = llil_ssa_function.get_ssa_register_uses(reg.source_reg()); + println!("{:?}", sp_0_uses); + assert_eq!(sp_0_uses.len(), 2); + let _next_instr_use = sp_0_uses + .iter() + .find(|inst| inst.index != ssa_instr_1.index) + .expect("Failed to get next instructions use of sp"); + } + _ => panic!("Expected RegSsa"), + } } _ => panic!("Expected StoreSsa"), }