[ObjC] Set return type on objc_msgSend calls that send an init message to a known class

Author: bdash

plugins/workflow_objc/src/activities/objc_msg_send_calls.rs

@@ -107,7 +107,7 @@ fn process_instruction(
     };
 
     let mut function_changed = false;
-    if adjust_call_type::process_call(bv, func, insn, &selector, message_send_type).is_ok() {
+    if adjust_call_type::process_call(bv, func, ssa, insn, &selector, message_send_type).is_ok() {
         function_changed = true;
     }
 
@@ -166,7 +166,7 @@ fn selector_from_call(
         return None;
     };
 
-    let raw_selector = ssa.get_ssa_register_value(&reg.source_reg())?.value as u64;
+    let raw_selector = ssa.get_ssa_register_value(reg.source_reg())?.value as u64;
     if raw_selector == 0 {
         return None;
     }

plugins/workflow_objc/src/activities/objc_msg_send_calls/adjust_call_type.rs

@@ -1,26 +1,150 @@
 use binaryninja::{
+    architecture::CoreRegister,
     binary_view::{BinaryView, BinaryViewBase as _, BinaryViewExt},
     confidence::Conf,
     function::Function,
     low_level_il::{
-        function::{Mutable, SSA},
-        instruction::LowLevelILInstruction,
+        expression::{
+            ExpressionHandler as _, LowLevelILExpression, LowLevelILExpressionKind, ValueExpr,
+        },
+        function::{LowLevelILFunction, Mutable, SSA},
+        instruction::{InstructionHandler as _, LowLevelILInstruction, LowLevelILInstructionKind},
+        operation::{CallSsa, Operation},
+        LowLevelILSSARegisterKind,
     },
     rc::Ref,
     types::{FunctionParameter, Type},
+    variable::PossibleValueSet,
 };
+use bstr::ByteSlice as _;
 
 use super::MessageSendType;
-use crate::{metadata::Selector, workflow::Confidence, Error};
+use crate::{activities::util, metadata::Selector, workflow::Confidence, Error};
 
 fn named_type(bv: &BinaryView, name: &str) -> Option<Ref<Type>> {
     bv.type_by_name(name)
         .map(|t| Type::named_type_from_type(name, &t))
 }
 
+// j_ prefixes are for stub functions in the dyld shared cache.
+const ALLOC_FUNCTIONS: &[&str] = &[
+    "_objc_alloc_init",
+    "_objc_alloc_initWithZone",
+    "_objc_alloc",
+    "_objc_allocWithZone",
+    "_objc_opt_new",
+    "j__objc_alloc_init",
+    "j__objc_alloc_initWithZone",
+    "j__objc_alloc",
+    "j__objc_allocWithZone",
+    "j__objc_opt_new",
+];
+
+/// Extract parameter expressions from a call, handling the SeparateParamListSsa wrapper.
+fn call_param_exprs<'a>(
+    call_op: &'a Operation<'a, Mutable, SSA, CallSsa>,
+) -> Option<Vec<LowLevelILExpression<'a, Mutable, SSA, ValueExpr>>> {
+    let LowLevelILExpressionKind::CallParamSsa(params) = &call_op.param_expr().kind() else {
+        return None;
+    };
+
+    let param_exprs = params.param_exprs();
+    Some(
+        if let Some(LowLevelILExpressionKind::SeparateParamListSsa(inner)) =
+            param_exprs.first().map(|e| e.kind())
+        {
+            inner.param_exprs()
+        } else {
+            param_exprs
+        },
+    )
+}
+
+/// Follow an SSA register back through register-to-register copies to find the
+/// instruction that originally defined its value.
+fn source_def_for_register<'a>(
+    ssa: &'a LowLevelILFunction<Mutable, SSA>,
+    reg: LowLevelILSSARegisterKind<CoreRegister>,
+) -> Option<LowLevelILInstruction<'a, Mutable, SSA>> {
+    let mut def = ssa.get_ssa_register_definition(reg)?;
+    while let LowLevelILInstructionKind::SetRegSsa(set_reg) = def.kind() {
+        let LowLevelILExpressionKind::RegSsa(src_reg) = set_reg.source_expr().kind() else {
+            break;
+        };
+        def = ssa.get_ssa_register_definition(src_reg.source_reg())?;
+    }
+    Some(def)
+}
+
+/// For init-family selectors on a normal message send, try to determine the return type
+/// by tracing the receiver register back to an alloc call and resolving the class.
+fn return_type_for_init_receiver(
+    bv: &BinaryView,
+    func: &Function,
+    ssa: &LowLevelILFunction<Mutable, SSA>,
+    insn: &LowLevelILInstruction<Mutable, SSA>,
+    selector: &Selector,
+    message_send_type: MessageSendType,
+) -> Option<Ref<Type>> {
+    if message_send_type != MessageSendType::Normal || !selector.name.starts_with("init") {
+        return None;
+    }
+
+    let call_op = match insn.kind() {
+        LowLevelILInstructionKind::CallSsa(op) | LowLevelILInstructionKind::TailCallSsa(op) => op,
+        _ => return None,
+    };
+
+    let param_exprs = call_param_exprs(&call_op)?;
+    let LowLevelILExpressionKind::RegSsa(receiver_reg) = param_exprs.first()?.kind() else {
+        return None;
+    };
+
+    let def = source_def_for_register(ssa, receiver_reg.source_reg())?;
+    let def_call_op = match def.kind() {
+        LowLevelILInstructionKind::CallSsa(op) | LowLevelILInstructionKind::TailCallSsa(op) => op,
+        _ => return None,
+    };
+
+    // Check if the defining call is to an alloc function.
+    let target_values = def_call_op.target().possible_values();
+    let call_target = match target_values {
+        PossibleValueSet::ConstantValue { value }
+        | PossibleValueSet::ConstantPointerValue { value }
+        | PossibleValueSet::ImportedAddressValue { value } => value as u64,
+        _ => return None,
+    };
+
+    let target_name = bv
+        .symbol_by_address(call_target)?
+        .raw_name()
+        .to_string_lossy()
+        .into_owned();
+    if !ALLOC_FUNCTIONS.contains(&target_name.as_str()) {
+        return None;
+    }
+
+    // Get the class from the alloc call's first parameter.
+    let alloc_params = call_param_exprs(&def_call_op)?;
+    let LowLevelILExpressionKind::RegSsa(class_reg) = alloc_params.first()?.kind() else {
+        return None;
+    };
+
+    let class_addr = ssa.get_ssa_register_value(class_reg.source_reg())?.value as u64;
+    if class_addr == 0 {
+        return None;
+    }
+
+    let class_symbol_name = bv.symbol_by_address(class_addr)?.full_name();
+    let class_name = util::class_name_from_symbol_name(class_symbol_name.to_bytes().as_bstr())?;
+    let class_type = bv.type_by_name(class_name.to_str_lossy())?;
+    Some(Type::pointer(&func.arch(), &class_type))
+}
+
 pub fn process_call(
     bv: &BinaryView,
     func: &Function,
+    ssa: &LowLevelILFunction<Mutable, SSA>,
     insn: &LowLevelILInstruction<Mutable, SSA>,
     selector: &Selector,
     message_send_type: MessageSendType,
@@ -39,8 +163,9 @@ pub fn process_call(
     };
     let sel = named_type(bv, "SEL").unwrap_or_else(|| Type::pointer(&arch, &Type::char()));
 
-    // TODO: Infer return type based on receiver type / selector.
-    let return_type = id.clone();
+    let return_type =
+        return_type_for_init_receiver(bv, func, ssa, insn, selector, message_send_type)
+            .unwrap_or_else(|| id.clone());
 
     let mut params = vec![
         FunctionParameter::new(receiver_type, receiver_name.to_string(), None),