{-# LANGUAGE PartialTypeSignatures, Strict #-}
module Data.ControlFlow (
ResolvedJumpTarget(..),
post,
fetch_block,
address_has_instruction,
address_has_symbol,
address_is_external,
operand_static_resolve,
resolve_jump_target,
get_internal_addresses,
instruction_jumps_to_external,
show_block,
show_invariants,
function_name_of_entry,
function_name_of_instruction,
isTerminal
)
where
import Data.Binary
import Algorithm.SCC
import Analysis.Context
import Base
import Data.SimplePred
import X86.Conventions
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 qualified X86.Instruction as X86
import qualified X86.Operand as X86
import Typeclasses.HasSize (sizeof)
import Typeclasses.HasAddress (addressof)
import Generic.Address (GenericAddress(..))
import Generic.Operand (GenericOperand(..))
import Generic.Instruction (GenericInstruction(Instruction))
import qualified Generic.Instruction as Instr
post :: CFG -> IS.Key -> IS.IntSet
post :: CFG -> Key -> IntSet
post CFG
g Key
blockId = IntSet -> Maybe IntSet -> IntSet
forall a. a -> Maybe a -> a
fromMaybe IntSet
IS.empty (Key -> IntMap IntSet -> Maybe IntSet
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
blockId (CFG -> IntMap IntSet
cfg_edges CFG
g))
fetch_block ::
CFG
-> Int
-> [X86.Instruction]
fetch_block :: CFG -> Key -> [Instruction]
fetch_block CFG
g Key
blockId =
case Key -> IntMap [Instruction] -> Maybe [Instruction]
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
blockId (IntMap [Instruction] -> Maybe [Instruction])
-> IntMap [Instruction] -> Maybe [Instruction]
forall a b. (a -> b) -> a -> b
$ CFG -> IntMap [Instruction]
cfg_instrs (CFG -> IntMap [Instruction]) -> CFG -> IntMap [Instruction]
forall a b. (a -> b) -> a -> b
$ CFG
g of
Maybe [Instruction]
Nothing -> [Char] -> [Instruction]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [Instruction]) -> [Char] -> [Instruction]
forall a b. (a -> b) -> a -> b
$ [Char]
"Block with ID" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Key -> [Char]
forall a. Show a => a -> [Char]
show Key
blockId [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" not found in cfg."
Just [Instruction]
b -> [Instruction]
b
isTerminal :: CFG -> IS.Key -> Bool
isTerminal :: CFG -> Key -> Bool
isTerminal CFG
cfg Key
b = Maybe IntSet -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe IntSet -> Bool) -> Maybe IntSet -> Bool
forall a b. (a -> b) -> a -> b
$ Key -> IntMap IntSet -> Maybe IntSet
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
b (CFG -> IntMap IntSet
cfg_edges CFG
cfg)
address_has_instruction :: Context -> a -> Bool
address_has_instruction Context
ctxt a
a =
case Context -> Word64 -> Maybe ([Char], [Char], Word64, Word64)
find_section_for_address Context
ctxt (Word64 -> Maybe ([Char], [Char], Word64, Word64))
-> Word64 -> Maybe ([Char], [Char], Word64, Word64)
forall a b. (a -> b) -> a -> b
$ a -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a of
Maybe ([Char], [Char], Word64, Word64)
Nothing -> Bool
False
Just ([Char]
segment,[Char]
section,Word64
_,Word64
_) -> ([Char]
segment,[Char]
section) ([Char], [Char]) -> [([Char], [Char])] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [([Char], [Char])]
sections_with_instructions
address_has_symbol :: Context -> a -> Bool
address_has_symbol Context
ctxt a
a =
case Key -> IntMap [Char] -> Maybe [Char]
forall a. Key -> IntMap a -> Maybe a
IM.lookup (a -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a) (IntMap [Char] -> Maybe [Char]) -> IntMap [Char] -> Maybe [Char]
forall a b. (a -> b) -> a -> b
$ Context -> IntMap [Char]
ctxt_syms Context
ctxt of
Maybe [Char]
Nothing -> Bool
False
Just [Char]
sym -> Bool
True
address_is_external :: Context -> a -> Bool
address_is_external Context
ctxt a
a = Context -> a -> Bool
forall a. Integral a => Context -> a -> Bool
address_has_symbol Context
ctxt a
a Bool -> Bool -> Bool
|| Bool -> Bool
not (Context -> a -> Bool
forall a. Integral a => Context -> a -> Bool
address_has_instruction Context
ctxt a
a)
operand_static_resolve ::
Context
-> X86.Instruction
-> X86.Operand
-> 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 a. HasAddress a => a -> 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 a. HasAddress a => a -> 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
forall a.
(HasAddress a, HasSize a) =>
Context -> a -> (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
forall a.
(HasAddress a, HasSize a) =>
Context -> a -> (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
forall a.
(HasAddress a, HasSize a) =>
Context -> a -> (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 -> a -> (Word64 -> Word64) -> Key -> ResolvedJumpTarget
static_resolve_rip_expr Context
ctxt a
i Word64 -> Word64
f Key
si =
let rip :: Word64
rip = a -> Word64
forall a. HasAddress a => a -> Word64
addressof a
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
$ a -> Key
forall a. HasSize a => a -> Key
sizeof a
i)
a' :: Word64
a' = Word64 -> Word64
f Word64
rip
syms :: IntMap [Char]
syms = Context -> IntMap [Char]
ctxt_syms Context
ctxt in
case (Key -> IntMap [Char] -> Maybe [Char]
forall a. Key -> IntMap a -> Maybe a
IM.lookup (Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a') IntMap [Char]
syms,Context -> Word64 -> Key -> Maybe Word64
read_from_ro_datasection Context
ctxt Word64
a' Key
si) of
(Just [Char]
s, Maybe Word64
a'') ->
[Char] -> ResolvedJumpTarget
External [Char]
s
(Maybe [Char]
Nothing, Just Word64
a'') ->
Word64 -> ResolvedJumpTarget
ImmediateAddress Word64
a''
(Maybe [Char]
Nothing,Maybe Word64
Nothing) ->
ResolvedJumpTarget
Unresolved
resolve_jump_target ::
Context
-> X86.Instruction
-> [ResolvedJumpTarget]
resolve_jump_target :: Context -> Instruction -> [ResolvedJumpTarget]
resolve_jump_target Context
ctxt Instruction
i =
case (Key -> IntMap Indirection -> Maybe Indirection
forall a. Key -> IntMap a -> Maybe a
IM.lookup (Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Key) -> Word64 -> Key
forall a b. (a -> b) -> a -> b
$ Instruction -> Word64
forall a. HasAddress a => a -> Word64
addressof Instruction
i) (IntMap Indirection -> Maybe Indirection)
-> IntMap Indirection -> Maybe Indirection
forall a b. (a -> b) -> a -> b
$ Context -> IntMap Indirection
ctxt_inds Context
ctxt, Instruction -> [Operand]
forall label storage prefix opcode annotation.
GenericInstruction label storage prefix opcode annotation
-> [GenericOperand storage]
Instr.srcs Instruction
i) of
(Just Indirection
ind,[Operand]
_) -> Indirection -> [ResolvedJumpTarget]
jump_targets_of_indirection Indirection
ind
(Maybe Indirection
Nothing,[Operand
op1]) ->
case Context -> Instruction -> Operand -> ResolvedJumpTarget
operand_static_resolve Context
ctxt Instruction
i Operand
op1 of
ResolvedJumpTarget
Unresolved -> [ResolvedJumpTarget
Unresolved]
External [Char]
sym -> [[Char] -> ResolvedJumpTarget
External [Char]
sym]
ImmediateAddress Word64
a ->
case Key -> IntMap [Char] -> Maybe [Char]
forall a. Key -> IntMap a -> Maybe a
IM.lookup (Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a) (IntMap [Char] -> Maybe [Char]) -> IntMap [Char] -> Maybe [Char]
forall a b. (a -> b) -> a -> b
$ Context -> IntMap [Char]
ctxt_syms Context
ctxt of
Just [Char]
sym -> [[Char] -> ResolvedJumpTarget
External [Char]
sym]
Maybe [Char]
Nothing -> if Bool -> Bool
not (Context -> Word64 -> Bool
forall a. Integral a => Context -> a -> Bool
address_has_instruction Context
ctxt Word64
a) then [[Char] -> ResolvedJumpTarget
External ([Char] -> ResolvedJumpTarget) -> [Char] -> ResolvedJumpTarget
forall a b. (a -> b) -> a -> b
$ Word64 -> [Char]
forall a. (Integral a, Show a) => a -> [Char]
showHex Word64
a] else [Word64 -> ResolvedJumpTarget
ImmediateAddress Word64
a]
where
jump_targets_of_indirection :: Indirection -> [ResolvedJumpTarget]
jump_targets_of_indirection (IndirectionResolved Set ResolvedJumpTarget
trgts) = Set ResolvedJumpTarget -> [ResolvedJumpTarget]
forall a. Set a -> [a]
S.toList (Set ResolvedJumpTarget -> [ResolvedJumpTarget])
-> Set ResolvedJumpTarget -> [ResolvedJumpTarget]
forall a b. (a -> b) -> a -> b
$ Set ResolvedJumpTarget
trgts
jump_targets_of_indirection (IndirectionJumpTable (JumpTable Operand
_ Operand
_ [Key]
entries)) = (Key -> ResolvedJumpTarget) -> [Key] -> [ResolvedJumpTarget]
forall a b. (a -> b) -> [a] -> [b]
map (Word64 -> ResolvedJumpTarget
ImmediateAddress (Word64 -> ResolvedJumpTarget)
-> (Key -> Word64) -> Key -> ResolvedJumpTarget
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) [Key]
entries
instruction_jumps_to_external ::
Context
-> X86.Instruction
-> Bool
instruction_jumps_to_external :: Context -> Instruction -> Bool
instruction_jumps_to_external Context
ctxt Instruction
i =
(ResolvedJumpTarget -> Bool) -> [ResolvedJumpTarget] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ResolvedJumpTarget -> Bool
resolve_is_external ([ResolvedJumpTarget] -> Bool) -> [ResolvedJumpTarget] -> Bool
forall a b. (a -> b) -> a -> b
$ Context -> Instruction -> [ResolvedJumpTarget]
resolve_jump_target Context
ctxt Instruction
i
where
resolve_is_external :: ResolvedJumpTarget -> Bool
resolve_is_external (External [Char]
_) = Bool
True
resolve_is_external ResolvedJumpTarget
_ = Bool
False
function_name_of_entry ::
Context
-> Int
-> String
function_name_of_entry :: Context -> Key -> [Char]
function_name_of_entry Context
ctxt Key
a =
case Key -> IntMap [Char] -> Maybe [Char]
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
a (IntMap [Char] -> Maybe [Char]) -> IntMap [Char] -> Maybe [Char]
forall a b. (a -> b) -> a -> b
$ Context -> IntMap [Char]
ctxt_syms Context
ctxt of
Just [Char]
sym -> [Char]
sym
Maybe [Char]
Nothing ->
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
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 [Char]
sym -> [Char]
sym
ResolvedJumpTarget
_ -> [Char]
"0x" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Key -> [Char]
forall a. (Integral a, Show a) => a -> [Char]
showHex Key
a
Maybe Instruction
_ -> [Char]
"0x" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Key -> [Char]
forall a. (Integral a, Show a) => a -> [Char]
showHex Key
a
function_name_of_instruction ::
Context
-> X86.Instruction
-> String
function_name_of_instruction :: Context -> Instruction -> [Char]
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 [Char]
sym -> [Char]
sym
ImmediateAddress Word64
a -> Context -> Key -> [Char]
function_name_of_entry Context
ctxt (Key -> [Char]) -> Key -> [Char]
forall a b. (a -> b) -> a -> b
$ Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a
ResolvedJumpTarget
Unresolved -> [Char]
"indirection@" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Word64 -> [Char]
forall a. (Integral a, Show a) => a -> [Char]
showHex (Instruction -> Word64
forall a. HasAddress a => a -> Word64
addressof Instruction
i)
else
[Char]
""
get_internal_addresses ::
ResolvedJumpTarget
-> [Int]
get_internal_addresses :: ResolvedJumpTarget -> [Key]
get_internal_addresses (External [Char]
_) = []
get_internal_addresses ResolvedJumpTarget
Unresolved = []
get_internal_addresses (ImmediateAddress Word64
a) = [Word64 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a]
show_block ::
CFG
-> Int
-> String
show_block :: CFG -> Key -> [Char]
show_block CFG
g Key
b =
let instrs :: [Key]
instrs = [Char] -> IntMap [Key] -> Key -> [Key]
forall p. [Char] -> IntMap p -> Key -> p
im_lookup ([Char]
"show_block: Block " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Key -> [Char]
forall a. Show a => a -> [Char]
show Key
b [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"in cfg.") (CFG -> IntMap [Key]
cfg_blocks CFG
g) Key
b in
Key -> [Char]
forall a. Show a => a -> [Char]
show Key
b [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" ["
[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Key -> [Char]
forall a. (Integral a, Show a) => a -> [Char]
showHex ([Key] -> Key
forall a. [a] -> a
head [Key]
instrs)
[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
","
[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Key -> [Char]
forall a. (Integral a, Show a) => a -> [Char]
showHex ([Key] -> Key
forall a. [a] -> a
last [Key]
instrs)
[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]"
show_invariants
:: CFG
-> Invariants
-> String
show_invariants :: CFG -> Invariants -> [Char]
show_invariants CFG
g Invariants
invs = [Char] -> [[Char]] -> [Char]
forall a. [a] -> [[a]] -> [a]
intercalate [Char]
"\n\n" ([[Char]] -> [Char]) -> [[Char]] -> [Char]
forall a b. (a -> b) -> a -> b
$ ((Key, Pred) -> [Char]) -> [(Key, Pred)] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (Key, Pred) -> [Char]
forall a. Show a => (Key, a) -> [Char]
show_entry ([(Key, Pred)] -> [[Char]]) -> [(Key, Pred)] -> [[Char]]
forall a b. (a -> b) -> a -> b
$ Invariants -> [(Key, Pred)]
forall a. IntMap a -> [(Key, a)]
IM.toList (Invariants -> [(Key, Pred)]) -> Invariants -> [(Key, Pred)]
forall a b. (a -> b) -> a -> b
$ Invariants
invs
where
show_entry :: (Key, a) -> [Char]
show_entry (Key
blockId, a
p) = [Char]
"Block " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ CFG -> Key -> [Char]
show_block CFG
g Key
blockId [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
":\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ a -> [Char]
forall a. Show a => a -> [Char]
show a
p
instance IntGraph CFG where
intgraph_post :: CFG -> Key -> IntSet
intgraph_post = CFG -> Key -> IntSet
post
intgraph_V :: CFG -> IntSet
intgraph_V = IntMap [Key] -> IntSet
forall a. IntMap a -> IntSet
IM.keysSet (IntMap [Key] -> IntSet) -> (CFG -> IntMap [Key]) -> CFG -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFG -> IntMap [Key]
cfg_blocks