{-# 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 Binary.FunctionNames where

import Binary.Generic


import Base
import Data.JumpTarget
import Conventions
import Data.X86.Opcode
import Data.X86.Instruction
import Data.X86.Register
import Data.Symbol



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 Control.Applicative ((<|>))






-- | 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.
try_plt_target_for_entry ::
  BinaryClass bin => 
     bin
  -> Word64 -- ^ The entry address
  -> Maybe ResolvedJumpTarget
try_plt_target_for_entry :: forall bin.
BinaryClass bin =>
bin -> Word64 -> Maybe ResolvedJumpTarget
try_plt_target_for_entry bin
bin Word64
a = Word64 -> Maybe ResolvedJumpTarget
jmps_to_external_function Word64
a Maybe ResolvedJumpTarget
-> Maybe ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. Eq a => Maybe a -> Maybe a -> Maybe a
`orTry` Word64 -> Maybe ResolvedJumpTarget
endbr64_jmps_to_external_function Word64
a
  --firstJusts [ {--is_external_function a,--} jmps_to_external_function a, endbr64_jmps_to_external_function a ]
 where
{--
  is_external_function a = do
    let SymbolTable syms _ = binary_get_symbols $ fst bin
    let sym_a = IM.lookup (fromIntegral a) syms
    case sym_a of
      Just (PointerToLabel sym True) -> Just $ External $ strip_GLIBC sym
      _                              -> Nothing 
--}
  jmps_to_external_function :: Word64 -> Maybe ResolvedJumpTarget
jmps_to_external_function Word64
a =
    case bin -> Word64 -> Maybe Instruction
forall bin. BinaryClass bin => bin -> Word64 -> Maybe Instruction
fetch_instruction bin
bin Word64
a of
      Just i :: Instruction
i@(Instruction Word64
_ [Prefix]
_  Opcode
JMP Maybe Operand
Nothing [Operand
op1] Int
_) -> Instruction -> Operand -> Maybe ResolvedJumpTarget
resolve_operand Instruction
i Operand
op1
      Maybe Instruction
_                                             -> Maybe ResolvedJumpTarget
forall a. Maybe a
Nothing

  endbr64_jmps_to_external_function :: Word64 -> Maybe ResolvedJumpTarget
endbr64_jmps_to_external_function Word64
a =
    case bin -> Word64 -> Maybe Instruction
forall bin. BinaryClass bin => bin -> Word64 -> Maybe Instruction
fetch_instruction bin
bin Word64
a of
      Just (Instruction Word64
_ [Prefix]
_ Opcode
ENDBR64 Maybe Operand
Nothing [Operand]
_ Int
si) -> Word64 -> Maybe ResolvedJumpTarget
jmps_to_external_function (Word64
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
si)
      Maybe Instruction
_                                           -> Maybe ResolvedJumpTarget
forall a. Maybe a
Nothing

  resolve_operand :: Instruction -> Operand -> Maybe ResolvedJumpTarget
resolve_operand Instruction
i Operand
op1 =
    case bin -> Instruction -> Operand -> ResolvedJumpTarget
forall bin.
BinaryClass bin =>
bin -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve bin
bin Instruction
i Operand
op1 of
      External String
sym -> 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
sym
      ExternalDeref String
sym -> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. a -> Maybe a
Just (ResolvedJumpTarget -> Maybe ResolvedJumpTarget)
-> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> ResolvedJumpTarget
ExternalDeref String
sym
      ImmediateAddress Word64
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
a'
      ResolvedJumpTarget
_ -> Maybe ResolvedJumpTarget
forall a. Maybe a
Nothing


-- | Given an instruction that calls or jumps, try to find a jump target
jump_target_for_instruction ::
  BinaryClass bin =>
     bin
  -> Instruction -- ^ The instruction
  -> ResolvedJumpTarget
jump_target_for_instruction :: forall bin.
BinaryClass bin =>
bin -> Instruction -> ResolvedJumpTarget
jump_target_for_instruction bin
bin i :: Instruction
i@(Instruction Word64
_ [Prefix]
_ Opcode
_ Maybe Operand
_ [Operand]
ops Int
_) =
  case bin -> Instruction -> Operand -> ResolvedJumpTarget
forall bin.
BinaryClass bin =>
bin -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve bin
bin Instruction
i ([Operand] -> Operand
forall a. HasCallStack => [a] -> a
head [Operand]
ops) of
    External String
sym       -> String -> ResolvedJumpTarget
External String
sym
    ExternalDeref String
sym  -> String -> ResolvedJumpTarget
ExternalDeref String
sym
    ResolvedJumpTarget
Unresolved         -> ResolvedJumpTarget
Unresolved
    ImmediateAddress Word64
a ->
      case bin -> Word64 -> Maybe ResolvedJumpTarget
forall bin.
BinaryClass bin =>
bin -> Word64 -> Maybe ResolvedJumpTarget
try_plt_target_for_entry bin
bin Word64
a of
        Just ResolvedJumpTarget
trgt -> ResolvedJumpTarget
trgt
        Maybe ResolvedJumpTarget
Nothing   -> Word64 -> ResolvedJumpTarget
ImmediateAddress Word64
a



-- | 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 ::
  BinaryClass bin =>
     bin
  -> Instruction    -- ^ The instruction
  -> Operand        -- ^ The operand of the instruction to be resolved
  -> ResolvedJumpTarget
operand_static_resolve :: forall bin.
BinaryClass bin =>
bin -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve bin
bin Instruction
i (Op_Imm Immediate
_)   = String -> ResolvedJumpTarget
forall a. HasCallStack => String -> a
error String
"todo"
operand_static_resolve bin
bin Instruction
i (Op_Jmp (Immediate BitSize
_ Word64
v)) = Word64 -> ResolvedJumpTarget
ImmediateAddress Word64
v
operand_static_resolve bin
bin Instruction
i (Op_Const Int
_) = String -> ResolvedJumpTarget
forall a. HasCallStack => String -> a
error String
"todo"
operand_static_resolve bin
bin Instruction
i (Op_Near Operand
op) = bin -> Instruction -> Operand -> ResolvedJumpTarget
forall bin.
BinaryClass bin =>
bin -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve bin
bin Instruction
i Operand
op
operand_static_resolve bin
bin Instruction
i (Op_Far  Operand
op) = String -> ResolvedJumpTarget
forall a. HasCallStack => String -> a
error String
"todo"
operand_static_resolve bin
bin Instruction
i (Op_Mem BitSize
_ BitSize
_ (Reg64 GPR
RIP) Register
RegNone Word8
0 Int
displ Maybe SReg
Nothing) = bin -> Instruction -> Word64 -> ResolvedJumpTarget
forall {p} {p}.
(Show p, BinaryClass p) =>
p -> p -> Word64 -> ResolvedJumpTarget
try_read_function_pointer bin
bin Instruction
i (Word64 -> ResolvedJumpTarget) -> Word64 -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Instruction -> Word64
inAddress Instruction
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Instruction -> Int
inSize Instruction
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
displ)
operand_static_resolve bin
bin Instruction
i (Op_Mem BitSize
_ BitSize
_ (Reg64 GPR
RIP) Register
RegNone Word8
scale Int
displ Maybe SReg
Nothing) = String -> ResolvedJumpTarget
forall a. HasCallStack => String -> a
error (String -> ResolvedJumpTarget) -> String -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ (Instruction, Integer) -> String
forall a. Show a => a -> String
show (Instruction
i,Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Instruction -> Word64
inAddress Instruction
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Instruction -> Int
inSize Instruction
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
displ))
operand_static_resolve bin
bin Instruction
i Operand
_            = ResolvedJumpTarget
Unresolved



-- If *[a',8] = fptr, then return that fptr
try_read_function_pointer :: p -> p -> Word64 -> ResolvedJumpTarget
try_read_function_pointer p
bin p
i Word64
a' = Word64 -> Maybe ResolvedJumpTarget
forall {a}. Integral a => a -> Maybe ResolvedJumpTarget
try_symbol Word64
a' Maybe ResolvedJumpTarget
-> Maybe ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. Eq a => Maybe a -> Maybe a -> Maybe a
`orTry` Word64 -> Maybe ResolvedJumpTarget
try_relocation Word64
a' Maybe ResolvedJumpTarget
-> ResolvedJumpTarget -> ResolvedJumpTarget
forall a. Eq a => Maybe a -> a -> a
`orElse` Word64 -> ResolvedJumpTarget
immediate_address Word64
a'
 where
  try_symbol :: a -> Maybe ResolvedJumpTarget
try_symbol a
a' = do
    let SymbolTable IntMap Symbol
syms Set String
_ = p -> SymbolTable
forall a. BinaryClass a => a -> SymbolTable
binary_get_symbols p
bin
    case Int -> IntMap Symbol -> Maybe Symbol
forall a. Int -> IntMap a -> Maybe a
IM.lookup (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a') IntMap Symbol
syms 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 (PointerToLabel String
s Bool
True)  -> 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 -> ResolvedJumpTarget) -> String -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> String
strip_GLIBC String
s
      Just (PointerToLabel String
s Bool
False) -> IntMap Symbol -> String -> Maybe ResolvedJumpTarget
find_address_of_label IntMap Symbol
syms String
s
      Just (AddressOfLabel String
s Bool
True)  -> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. a -> Maybe a
Just (ResolvedJumpTarget -> Maybe ResolvedJumpTarget)
-> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> ResolvedJumpTarget
ExternalDeref (String -> ResolvedJumpTarget) -> String -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> String
strip_GLIBC String
s
      Just (AddressOfObject String
s Bool
True) -> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a. a -> Maybe a
Just (ResolvedJumpTarget -> Maybe ResolvedJumpTarget)
-> ResolvedJumpTarget -> Maybe ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> ResolvedJumpTarget
ExternalDeref (String -> ResolvedJumpTarget) -> String -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String -> String
strip_GLIBC String
s


      Just Symbol
s -> String -> Maybe ResolvedJumpTarget
forall a. HasCallStack => String -> a
error (String -> Maybe ResolvedJumpTarget)
-> String -> Maybe ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ (p, String, Symbol) -> String
forall a. Show a => a -> String
show (p
i,a -> String
forall {a}. Integral a => a -> String
showHex a
a',Symbol
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
$ p -> Set Relocation
forall a. BinaryClass a => a -> Set Relocation
binary_get_relocations p
bin of
      Maybe Relocation
Nothing -> Maybe ResolvedJumpTarget
forall a. Maybe a
Nothing
      Just (Relocation Word64
_ Word64
a1) -> 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 --TODO what if a1 is a symbol?
        --case try_symbol a1 of
        --  Nothing -> Just $ ImmediateAddress a1
        --  Just sym -> Just sym -- error $ show (a',a1,sym)
    

  is_reloc_for :: Word64 -> Relocation -> Bool
is_reloc_for Word64
a' (Relocation Word64
a Word64
_) = Word64
a Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
a'

  find_address_of_label :: IntMap Symbol -> String -> Maybe ResolvedJumpTarget
find_address_of_label IntMap Symbol
syms String
l = (Word64 -> ResolvedJumpTarget
ImmediateAddress (Word64 -> ResolvedJumpTarget)
-> ((Int, Symbol) -> Word64) -> (Int, Symbol) -> ResolvedJumpTarget
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64)
-> ((Int, Symbol) -> Int) -> (Int, Symbol) -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Symbol) -> Int
forall a b. (a, b) -> a
fst) ((Int, Symbol) -> ResolvedJumpTarget)
-> Maybe (Int, Symbol) -> Maybe ResolvedJumpTarget
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (((Int, Symbol) -> Bool) -> [(Int, Symbol)] -> Maybe (Int, Symbol)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(Int
a,Symbol
symbol) -> Symbol
symbol Symbol -> Symbol -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Bool -> Symbol
AddressOfLabel String
l Bool
False) ([(Int, Symbol)] -> Maybe (Int, Symbol))
-> [(Int, Symbol)] -> Maybe (Int, Symbol)
forall a b. (a -> b) -> a -> b
$ IntMap Symbol -> [(Int, Symbol)]
forall a. IntMap a -> [(Int, a)]
IM.assocs IntMap Symbol
syms)

  immediate_address :: Word64 -> ResolvedJumpTarget
immediate_address Word64
a
    | p -> Word64 -> Bool
forall bin. BinaryClass bin => bin -> Word64 -> Bool
address_has_instruction p
bin Word64
a = String -> ResolvedJumpTarget
forall a. HasCallStack => String -> a
error (String -> ResolvedJumpTarget) -> String -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ String
"Does this happen?" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall {a}. Integral a => a -> String
showHex Word64
a -- Should NOT return this! ImmediateAddress a
    | Bool
otherwise = ResolvedJumpTarget
Unresolved



function_name_of_entry ::
  BinaryClass bin =>
     bin
  -> Word64 -- ^ The entry address
  -> String
function_name_of_entry :: forall bin. BinaryClass bin => bin -> Word64 -> String
function_name_of_entry bin
bin Word64
a =
  case bin -> Word64 -> Maybe ResolvedJumpTarget
forall bin.
BinaryClass bin =>
bin -> Word64 -> Maybe ResolvedJumpTarget
try_plt_target_for_entry bin
bin Word64
a of
    Maybe ResolvedJumpTarget
Nothing                    -> String
"0x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall {a}. Integral a => a -> String
showHex Word64
a
    Just (External String
sym)        -> String
sym
    Just (ExternalDeref  String
sym)  -> String
"*" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
sym
    Just (ImmediateAddress Word64
a') -> String
"0x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall {a}. Integral a => a -> String
showHex Word64
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 ::
  BinaryClass bin =>
     bin
  -> Instruction -- ^ The instruction
  -> String
function_name_of_instruction :: forall bin. BinaryClass bin => bin -> Instruction -> String
function_name_of_instruction bin
bin i :: Instruction
i@(Instruction Word64
_ [Prefix]
_ Opcode
_ Maybe Operand
_ [Operand]
ops Int
_)
  |  Opcode -> Bool
isCall (Instruction -> Opcode
inOperation Instruction
i) Bool -> Bool -> Bool
|| Opcode -> Bool
isJump (Instruction -> Opcode
inOperation Instruction
i) =
    case bin -> Instruction -> ResolvedJumpTarget
forall bin.
BinaryClass bin =>
bin -> Instruction -> ResolvedJumpTarget
jump_target_for_instruction bin
bin Instruction
i of
      External String
sym       -> String
sym
      ExternalDeref String
sym  -> String
"*" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
sym
      ImmediateAddress Word64
a -> String
"0x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall {a}. Integral a => a -> String
showHex Word64
a
      ResolvedJumpTarget
Unresolved         -> String
"indirection@" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall {a}. Integral a => a -> String
showHex (Instruction -> Word64
inAddress Instruction
i)
  | Bool
otherwise = String
""