{-# LANGUAGE PartialTypeSignatures, Strict #-}

{-|
Module      : FunctionNames
Description : Provides functions to 

Contains function relating to control flow, including functions for
resolving the targets of jumps and calls.
-}



module Analysis.FunctionNames where

import Analysis.Context
import Base
import Data.JumpTarget
import X86.Conventions
import Generic.Binary

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Set as S
import Data.Maybe (fromJust, fromMaybe, isNothing)
import Data.List
import Data.List.Split (chunksOf)
import Data.Word (Word64)
import Control.Monad ((>=>))
import Debug.Trace
import Numeric (readHex)
import System.IO.Unsafe (unsafePerformIO)
import X86.Register (Register(..))
import X86.Opcode (Opcode(JMP), isCall, isJump)
import X86.Instruction
import qualified X86.Operand as X86
import Generic.HasSize (sizeof)
import Generic.Address (GenericAddress(..))
import Generic.Operand (GenericOperand(..))
import Generic.Instruction (GenericInstruction(Instruction))
import Generic.SymbolicConstituents
import qualified Generic.Instruction as Instr
import Control.Applicative ((<|>))


-- | Returns true iff a symbol is associated with the address.
address_has_external_symbol :: Context -> a -> Bool
address_has_external_symbol Context
ctxt a
a =
  case Key -> IntMap Symbol -> Maybe Symbol
forall a. Key -> IntMap a -> Maybe a
IM.lookup (a -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a) (IntMap Symbol -> Maybe Symbol) -> IntMap Symbol -> Maybe Symbol
forall a b. (a -> b) -> a -> b
$ Context -> IntMap Symbol
ctxt_symbol_table Context
ctxt of
    Just (Relocated_Function String
_) -> Bool
True
    Just (Relocated_Label String
_) -> Bool
True
    Maybe Symbol
_ -> Bool
False


-- | many operands can statically be resolved, even though technically they are indirect (relative to RIP).
-- Examples:
--
-- @10005464e: call RIP + 1751660@ resolves to an immediate jump target by resolving the RIP-relative addressing.
--
-- @10005464e: call qword ptr [RIP + 1751660]@ read from address 1002000c0, but address has a symbol associated to it. This function call will resolve to an external function.
operand_static_resolve ::
  Context               -- ^ The context
  -> Instruction    -- ^ The instruction
  -> X86.Operand        -- ^ The operand of the instruction to be resolved
  -> ResolvedJumpTarget
operand_static_resolve :: Context -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve Context
ctxt Instruction
i (Immediate Word64
a')                                                         = Word64 -> ResolvedJumpTarget
ImmediateAddress Word64
a'
operand_static_resolve Context
ctxt Instruction
i (EffectiveAddress (AddressPlus (AddressStorage Register
RIP) (AddressImm Word64
imm))) = Word64 -> ResolvedJumpTarget
ImmediateAddress (Word64 -> ResolvedJumpTarget) -> Word64 -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ Instruction -> Word64
forall storage prefix opcode annotation.
GenericInstruction AddressWord64 storage prefix opcode annotation
-> Word64
addressof Instruction
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Key -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Instruction -> Key
forall a. HasSize a => a -> Key
sizeof Instruction
i) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
imm
operand_static_resolve Context
ctxt Instruction
i (EffectiveAddress (AddressPlus (AddressImm Word64
imm) (AddressStorage Register
RIP))) = Word64 -> ResolvedJumpTarget
ImmediateAddress (Word64 -> ResolvedJumpTarget) -> Word64 -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ Instruction -> Word64
forall storage prefix opcode annotation.
GenericInstruction AddressWord64 storage prefix opcode annotation
-> Word64
addressof Instruction
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Key -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Instruction -> Key
forall a. HasSize a => a -> Key
sizeof Instruction
i) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
imm
operand_static_resolve Context
ctxt Instruction
i (Memory (AddressPlus  (AddressStorage Register
RIP) (AddressImm Word64
imm)) Key
si)       = Context
-> Instruction -> (Word64 -> Word64) -> Key -> ResolvedJumpTarget
static_resolve_rip_expr Context
ctxt Instruction
i (\Word64
rip -> Word64
rip Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
imm) Key
si
operand_static_resolve Context
ctxt Instruction
i (Memory (AddressPlus  (AddressImm Word64
imm) (AddressStorage Register
RIP)) Key
si)       = Context
-> Instruction -> (Word64 -> Word64) -> Key -> ResolvedJumpTarget
static_resolve_rip_expr Context
ctxt Instruction
i (\Word64
rip -> Word64
rip Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
imm) Key
si
operand_static_resolve Context
ctxt Instruction
i (Memory (AddressMinus (AddressStorage Register
RIP) (AddressImm Word64
imm)) Key
si)       = Context
-> Instruction -> (Word64 -> Word64) -> Key -> ResolvedJumpTarget
static_resolve_rip_expr Context
ctxt Instruction
i (\Word64
rip -> Word64
rip Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
imm) Key
si
operand_static_resolve Context
ctxt Instruction
i Operand
_                                                                      = ResolvedJumpTarget
Unresolved

static_resolve_rip_expr :: Context -> Instruction -> (Word64 -> Word64) -> Int -> ResolvedJumpTarget
static_resolve_rip_expr :: Context
-> Instruction -> (Word64 -> Word64) -> Key -> ResolvedJumpTarget
static_resolve_rip_expr Context
ctxt Instruction
i Word64 -> Word64
f Key
si =
  let rip :: Word64
rip     = Instruction -> Word64
forall storage prefix opcode annotation.
GenericInstruction AddressWord64 storage prefix opcode annotation
-> Word64
addressof Instruction
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ (Key -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Key -> Word64) -> Key -> Word64
forall a b. (a -> b) -> a -> b
$ Instruction -> Key
forall a. HasSize a => a -> Key
sizeof Instruction
i)
      a' :: Word64
a'      = Word64 -> Word64
f Word64
rip in
    (Word64 -> Maybe ResolvedJumpTarget
forall a. Integral a => a -> Maybe ResolvedJumpTarget
try_relocated_function Word64
a' Maybe ResolvedJumpTarget
-> Maybe ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Word64 -> Maybe ResolvedJumpTarget
try_relocation Word64
a') Maybe ResolvedJumpTarget
-> ResolvedJumpTarget -> ResolvedJumpTarget
forall a. Eq a => Maybe a -> a -> a
`orElse` ResolvedJumpTarget
Unresolved
 where
  try_relocated_function :: a -> Maybe ResolvedJumpTarget
try_relocated_function a
a' =
    case Key -> IntMap Symbol -> Maybe Symbol
forall a. Key -> IntMap a -> Maybe a
IM.lookup (a -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a') (IntMap Symbol -> Maybe Symbol) -> IntMap Symbol -> Maybe Symbol
forall a b. (a -> b) -> a -> b
$ Context -> IntMap Symbol
ctxt_symbol_table Context
ctxt of
      -- Example:
      --   Instruction 10005464e: CALL 64 ptr [RIP + 1751660] 6 read from address 1002000c0 which has symbol _objc_msgSend producing address 0
      --   Address *[1002000c0,8] is treated as an external function call       
      Just (Relocated_Function String
s) -> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. a -> Maybe a
Just (ResolvedJumpTarget -> Maybe ResolvedJumpTarget)
-> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> ResolvedJumpTarget
External String
s
      Maybe Symbol
_ -> Maybe ResolvedJumpTarget
forall a. Maybe a
Nothing

  try_relocation :: Word64 -> Maybe ResolvedJumpTarget
try_relocation Word64
a' =
    case (Relocation -> Bool) -> Set Relocation -> Maybe Relocation
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Word64 -> Relocation -> Bool
is_reloc_for Word64
a') (Set Relocation -> Maybe Relocation)
-> Set Relocation -> Maybe Relocation
forall a b. (a -> b) -> a -> b
$ Context -> Set Relocation
ctxt_relocs Context
ctxt of
      Just (Relocation Word64
_ Word64
a1) -> Word64 -> Maybe ResolvedJumpTarget
forall a. Integral a => a -> Maybe ResolvedJumpTarget
try_relocated_function Word64
a1 Maybe ResolvedJumpTarget
-> Maybe ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. a -> Maybe a
Just (ResolvedJumpTarget -> Maybe ResolvedJumpTarget)
-> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ Word64 -> ResolvedJumpTarget
ImmediateAddress Word64
a1)
      Maybe Relocation
_ -> Maybe ResolvedJumpTarget
forall a. Maybe a
Nothing
  is_reloc_for :: Word64 -> Relocation -> Bool
is_reloc_for Word64
a' (Relocation Word64
a0 Word64
_) = Word64
a' Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
a0

{--
      (Nothing, Just a'') ->
        -- Example:
        --   Instruction 10011e093: CALL 64 ptr [RIP + 1098831] 6 read from address 10022a4e8 producing address 100131d63
        --   Address 100131d63 is returned as that is the function pointer to be called
        -- trace ("Instruction " ++ show i ++ " read from address " ++ showHex a' ++ " producing address " ++ showHex a'') $
        ImmediateAddress a''
--}


-- | Tries to retrieve a function name with an entry address.
-- If the entry matches a known symbol, return that.
-- Otherwise, simply return the entry address itself in hexadecimal notation.
-- However, there is one exception: 
-- 	if the first instruction at the entry address immediately jumps to an external function,
-- 	return the name of that external function instead. This happens in a @.got@ section.
function_name_of_entry ::
  Context  -- ^ The context
  -> Int   -- ^ The entry address
  -> String
function_name_of_entry :: Context -> Key -> String
function_name_of_entry Context
ctxt Key
a =
  case Key -> IntMap Symbol -> Maybe Symbol
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
a (IntMap Symbol -> Maybe Symbol) -> IntMap Symbol -> Maybe Symbol
forall a b. (a -> b) -> a -> b
$ Context -> IntMap Symbol
ctxt_symbol_table Context
ctxt of
    Just (Relocated_Function String
sym) -> String
sym
    Maybe Symbol
_ ->
      case IO (Maybe Instruction) -> Maybe Instruction
forall a. IO a -> a
unsafePerformIO (IO (Maybe Instruction) -> Maybe Instruction)
-> IO (Maybe Instruction) -> Maybe Instruction
forall a b. (a -> b) -> a -> b
$ Context -> Word64 -> IO (Maybe Instruction)
fetch_instruction Context
ctxt (Key -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
a) of -- TODO. However, should be safe as result is immutable.
        Just i :: Instruction
i@(Instruction AddressWord64
_ Maybe Prefix
_ Opcode
JMP Maybe Operand
Nothing [Operand
op1] Maybe Key
_)  ->
          case Context -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve Context
ctxt Instruction
i Operand
op1 of
            External String
sym -> String
sym
            ResolvedJumpTarget
_ -> String
"0x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Key -> String
forall a. (Integral a, Show a) => a -> String
showHex Key
a
        Maybe Instruction
_ -> String
"0x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Key -> String
forall a. (Integral a, Show a) => a -> String
showHex Key
a

-- | Tries to retrieve a function name for a @call@-instruction (see @`function_name_of_entry`@).
--
-- Returns the empty string if the given instruction is not a call or a jump.
function_name_of_instruction ::
  Context            -- ^ The context
  -> Instruction -- ^ The instruction
  -> String
function_name_of_instruction :: Context -> Instruction -> String
function_name_of_instruction Context
ctxt i :: Instruction
i@(Instruction AddressWord64
_ Maybe Prefix
_ Opcode
_ Maybe Operand
_ [Operand]
ops Maybe Key
_) =
  if Opcode -> Bool
isCall (Instruction -> Opcode
forall label storage prefix opcode annotation.
GenericInstruction label storage prefix opcode annotation -> opcode
Instr.opcode Instruction
i) Bool -> Bool -> Bool
|| Opcode -> Bool
isJump (Instruction -> Opcode
forall label storage prefix opcode annotation.
GenericInstruction label storage prefix opcode annotation -> opcode
Instr.opcode Instruction
i) then
    case Context -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve Context
ctxt Instruction
i ([Operand] -> Operand
forall a. [a] -> a
head [Operand]
ops) of
      External String
sym       -> String
sym
      ImmediateAddress Word64
a -> Context -> Key -> String
function_name_of_entry Context
ctxt (Key -> String) -> Key -> String
forall a b. (a -> b) -> a -> b
$ Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a
      ResolvedJumpTarget
Unresolved         -> String
"indirection@" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall a. (Integral a, Show a) => a -> String
showHex (Instruction -> Word64
forall storage prefix opcode annotation.
GenericInstruction AddressWord64 storage prefix opcode annotation
-> Word64
addressof Instruction
i)
  else
    String
""