module WithAbstractSymbolicValues.ResolveIndirections where

import Base

import WithAbstractSymbolicValues.Class
import WithAbstractSymbolicValues.FInit
import WithAbstractSymbolicValues.Sstate
import WithAbstractSymbolicValues.SymbolicExecution

import Data.X86.Opcode
import Data.X86.Instruction
import Data.X86.Register
import Data.JumpTarget
import Data.Indirection
import Data.SymbolicExpression (FlagStatus(..)) -- TODO


import qualified Data.Serialize as Cereal
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.IntSet as IS
import qualified Data.IntMap as IM
import qualified Data.Set.NonEmpty as NES

import Data.List (intercalate,partition,intersectBy,find,nub)
import Data.Word 
import Data.Maybe

import Control.Monad.Extra
import Control.Monad.State.Strict hiding (join)
import Control.DeepSeq




stry_resolve_indirection :: WithAbstractSymbolicValues ctxt v p => ctxt -> Sstate v p -> [Instruction] -> Indirections
stry_resolve_indirection :: forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> Sstate v p -> [Instruction] -> Indirections
stry_resolve_indirection ctxt
ctxt p :: Sstate v p
p@(Sstate Map Register v
regs Map (p, Maybe ByteSize) v
mem FlagStatus
flg) [Instruction]
instrs =
  let [Operand
trgt] = Instruction -> [Operand]
srcs (Instruction -> [Operand]) -> Instruction -> [Operand]
forall a b. (a -> b) -> a -> b
$ [Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
last [Instruction]
instrs in
    case Integer -> FlagStatus -> Maybe (Operand, Word64)
forall {a}.
Integral a =>
a -> FlagStatus -> Maybe (Operand, Word64)
flagstatus_to_tries Integer
10000 FlagStatus
flg of -- TODO
      Maybe (Operand, Word64)
Nothing -> Maybe Indirections
try_to_resolve_error_call Maybe Indirections -> Maybe Indirections -> Maybe Indirections
forall a. Eq a => Maybe a -> Maybe a -> Maybe a
`orTry` Operand -> Sstate v p -> Maybe Indirections
try_to_resolve_from_pre Operand
trgt Sstate v p
p Maybe Indirections -> Indirections -> Indirections
forall a. Eq a => Maybe a -> a -> a
`orElse` Indirection -> Indirections
forall a. a -> Set a
S.singleton Indirection
Indirection_Unresolved
      Just (Operand
op1,Word64
n) -> Maybe Indirections
try_to_resolve_error_call Maybe Indirections -> Maybe Indirections -> Maybe Indirections
forall a. Eq a => Maybe a -> Maybe a -> Maybe a
`orTry` Instruction -> Operand -> Word64 -> Operand -> Maybe Indirections
forall {p}.
Integral p =>
Instruction -> Operand -> p -> Operand -> Maybe Indirections
try_to_resolve_using_bound ([Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
head [Instruction]
instrs) Operand
op1 Word64
n Operand
trgt Maybe Indirections -> Maybe Indirections -> Maybe Indirections
forall a. Eq a => Maybe a -> Maybe a -> Maybe a
`orTry` Operand -> Sstate v p -> Maybe Indirections
try_to_resolve_from_pre Operand
trgt Sstate v p
p Maybe Indirections -> Indirections -> Indirections
forall a. Eq a => Maybe a -> a -> a
`orElse` Indirection -> Indirections
forall a. a -> Set a
S.singleton Indirection
Indirection_Unresolved

 where
  try_to_resolve_using_bound :: Instruction -> Operand -> p -> Operand -> Maybe Indirections
try_to_resolve_using_bound Instruction
i Operand
op1 p
n Operand
trgt =
    let values1 :: [Maybe (Set ResolvedJumpTarget)]
values1 = (p -> Maybe (Set ResolvedJumpTarget))
-> [p] -> [Maybe (Set ResolvedJumpTarget)]
forall a b. (a -> b) -> [a] -> [b]
map (\p
n -> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
-> (Sstate v p, VCS v) -> Maybe (Set ResolvedJumpTarget)
forall s a. State s a -> s -> a
evalState (Instruction
-> Operand
-> p
-> Operand
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
forall {i}.
Integral i =>
Instruction
-> Operand
-> i
-> Operand
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
evaluate_target_after_setting_value_and_block Instruction
i Operand
op1 p
n Operand
trgt) (Sstate v p -> Sstate v p
forall {p}. Sstate v p -> Sstate v p
clean_sstate Sstate v p
p, VCS v
forall a. Set a
S.empty)) [p
0..p
n] in
      if [Maybe (Set ResolvedJumpTarget)]
values1 [Maybe (Set ResolvedJumpTarget)]
-> [Maybe (Set ResolvedJumpTarget)] -> Bool
forall a. Eq a => a -> a -> Bool
== [] then
        Maybe Indirections
forall a. Maybe a
Nothing
      else if (Maybe (Set ResolvedJumpTarget) -> Bool)
-> [Maybe (Set ResolvedJumpTarget)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe (Set ResolvedJumpTarget) -> Bool
is_jump_table_entry [Maybe (Set ResolvedJumpTarget)]
values1 then do
        let trgts :: [Word64]
trgts = (Maybe (Set ResolvedJumpTarget) -> Word64)
-> [Maybe (Set ResolvedJumpTarget)] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map (Word64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Word64)
-> (Maybe (Set ResolvedJumpTarget) -> Word64)
-> Maybe (Set ResolvedJumpTarget)
-> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ResolvedJumpTarget -> Word64
get_immediate (ResolvedJumpTarget -> Word64)
-> (Maybe (Set ResolvedJumpTarget) -> ResolvedJumpTarget)
-> Maybe (Set ResolvedJumpTarget)
-> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set ResolvedJumpTarget -> ResolvedJumpTarget
forall a. Set a -> a
S.findMin (Set ResolvedJumpTarget -> ResolvedJumpTarget)
-> (Maybe (Set ResolvedJumpTarget) -> Set ResolvedJumpTarget)
-> Maybe (Set ResolvedJumpTarget)
-> ResolvedJumpTarget
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (Set ResolvedJumpTarget) -> Set ResolvedJumpTarget
forall a. HasCallStack => Maybe a -> a
fromJust) [Maybe (Set ResolvedJumpTarget)]
values1 
            tbl :: JumpTable
tbl = Operand -> Int -> Operand -> IntMap Word64 -> JumpTable
JumpTable Operand
op1 (p -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral p
n) Operand
trgt ([(Int, Word64)] -> IntMap Word64
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, Word64)] -> IntMap Word64)
-> [(Int, Word64)] -> IntMap Word64
forall a b. (a -> b) -> a -> b
$ [Int] -> [Word64] -> [(Int, Word64)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..p -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral p
n] [Word64]
trgts) in
          Indirections -> Maybe Indirections
forall a. a -> Maybe a
Just (Indirections -> Maybe Indirections)
-> Indirections -> Maybe Indirections
forall a b. (a -> b) -> a -> b
$ Indirection -> Indirections
forall a. a -> Set a
S.singleton (Indirection -> Indirections) -> Indirection -> Indirections
forall a b. (a -> b) -> a -> b
$ JumpTable -> Indirection
Indirection_JumpTable JumpTable
tbl
      else
        Maybe Indirections
forall a. Maybe a
Nothing

  try_to_resolve_from_pre :: Operand -> Sstate v p -> Maybe Indirections
try_to_resolve_from_pre Operand
trgt Sstate v p
p =
    let trgts :: Maybe (Set ResolvedJumpTarget)
trgts = State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
-> (Sstate v p, VCS v) -> Maybe (Set ResolvedJumpTarget)
forall s a. State s a -> s -> a
evalState (Operand
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
evaluate_target_after_block Operand
trgt) (Sstate v p
p,VCS v
forall a. Set a
S.empty) in
      (ResolvedJumpTarget -> Indirection)
-> Set ResolvedJumpTarget -> Indirections
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map ResolvedJumpTarget -> Indirection
Indirection_Resolved (Set ResolvedJumpTarget -> Indirections)
-> Maybe (Set ResolvedJumpTarget) -> Maybe Indirections
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Set ResolvedJumpTarget)
trgts


  evaluate_target_after_setting_value_and_block :: Instruction
-> Operand
-> i
-> Operand
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
evaluate_target_after_setting_value_and_block Instruction
i Operand
op1 i
n Operand
trgt = do
    ctxt -> Instruction -> State (Sstate v p, VCS v) ()
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> Instruction -> State (Sstate v p, VCS v) ()
sset_rip ctxt
ctxt ([Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
head [Instruction]
instrs)
    ctxt
-> Instruction
-> Bool
-> Operand
-> v
-> State (Sstate v p, VCS v) ()
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt
-> Instruction
-> Bool
-> Operand
-> v
-> State (Sstate v p, VCS v) ()
swrite_operand ctxt
ctxt Instruction
i Bool
False Operand
op1 (ctxt -> i -> v
forall i. Integral i => ctxt -> i -> v
forall ctxt v p i.
(WithAbstractSymbolicValues ctxt v p, Integral i) =>
ctxt -> i -> v
simmediate ctxt
ctxt i
n)
    Operand
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
evaluate_target_after_block  Operand
trgt

  evaluate_target_after_block :: Operand
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
evaluate_target_after_block Operand
trgt = do
    (Sstate v p
p,VCS v
_) <- StateT (Sstate v p, VCS v) Identity (Sstate v p, VCS v)
forall s (m :: * -> *). MonadState s m => m s
get
    ctxt
-> Bool
-> [Instruction]
-> Maybe [Instruction]
-> State (Sstate v p, VCS v) ()
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt
-> Bool
-> [Instruction]
-> Maybe [Instruction]
-> State (Sstate v p, VCS v) ()
sexec_block ctxt
ctxt Bool
False ([Instruction] -> [Instruction]
forall a. HasCallStack => [a] -> [a]
init [Instruction]
instrs) Maybe [Instruction]
forall a. Maybe a
Nothing
    (Sstate v p
q,VCS v
_) <- StateT (Sstate v p, VCS v) Identity (Sstate v p, VCS v)
forall s (m :: * -> *). MonadState s m => m s
get
    ctxt -> Instruction -> State (Sstate v p, VCS v) ()
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> Instruction -> State (Sstate v p, VCS v) ()
sset_rip ctxt
ctxt ([Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
last [Instruction]
instrs)
    v
val <- ctxt -> String -> Operand -> State (Sstate v p, VCS v) v
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> String -> Operand -> State (Sstate v p, VCS v) v
sread_operand ctxt
ctxt String
"indirection resolving" Operand
trgt
    Maybe (Set ResolvedJumpTarget)
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
forall a. a -> StateT (Sstate v p, VCS v) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Set ResolvedJumpTarget)
 -> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget)))
-> Maybe (Set ResolvedJumpTarget)
-> State (Sstate v p, VCS v) (Maybe (Set ResolvedJumpTarget))
forall a b. (a -> b) -> a -> b
$ ctxt -> v -> Maybe (Set ResolvedJumpTarget)
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> v -> Maybe (Set ResolvedJumpTarget)
stry_jump_targets ctxt
ctxt v
val -- $ trace ("is = " ++ show (instrs) ++ "P:\n" ++ show p ++ "\nQ:\n" ++ show q ++ "\nval,trgt: " ++ show (val,trgt)) $ 

  flagstatus_to_tries :: a -> FlagStatus -> Maybe (Operand, Word64)
flagstatus_to_tries a
max_tries (FS_CMP (Just Bool
True) Operand
op1 (Op_Imm (Immediate BitSize
_ Word64
n))) = if Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= a -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
max_tries then (Operand, Word64) -> Maybe (Operand, Word64)
forall a. a -> Maybe a
Just (Operand
op1,Word64
n) else Maybe (Operand, Word64)
forall a. Maybe a
Nothing
  flagstatus_to_tries a
max_tries FlagStatus
_ = Maybe (Operand, Word64)
forall a. Maybe a
Nothing


  try_to_resolve_error_call :: Maybe Indirections
try_to_resolve_error_call = Indirection -> Indirections
forall a. a -> Set a
S.singleton (Indirection -> Indirections)
-> Maybe Indirection -> Maybe Indirections
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State (Sstate v p, VCS v) (Maybe Indirection)
-> Sstate v p -> Maybe Indirection
forall v p a. State (Sstate v p, VCS v) a -> Sstate v p -> a
evalSstate State (Sstate v p, VCS v) (Maybe Indirection)
evaluate_error_call_after_block Sstate v p
p

  evaluate_error_call_after_block :: State (Sstate v p, VCS v) (Maybe Indirection)
evaluate_error_call_after_block = do
    ctxt
-> Bool
-> [Instruction]
-> Maybe [Instruction]
-> State (Sstate v p, VCS v) ()
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt
-> Bool
-> [Instruction]
-> Maybe [Instruction]
-> State (Sstate v p, VCS v) ()
sexec_block ctxt
ctxt Bool
False ([Instruction] -> [Instruction]
forall a. HasCallStack => [a] -> [a]
init [Instruction]
instrs) Maybe [Instruction]
forall a. Maybe a
Nothing
    ctxt -> Instruction -> State (Sstate v p, VCS v) ()
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> Instruction -> State (Sstate v p, VCS v) ()
sset_rip ctxt
ctxt ([Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
last [Instruction]
instrs)
    v
rdi <- ctxt -> Register -> State (Sstate v p, VCS v) v
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> Register -> State (Sstate v p, VCS v) v
sread_reg ctxt
ctxt (GPR -> Register
Reg64 GPR
RDI)
    Maybe Indirection -> State (Sstate v p, VCS v) (Maybe Indirection)
forall a. a -> StateT (Sstate v p, VCS v) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Indirection
 -> State (Sstate v p, VCS v) (Maybe Indirection))
-> Maybe Indirection
-> State (Sstate v p, VCS v) (Maybe Indirection)
forall a b. (a -> b) -> a -> b
$ ctxt -> Instruction -> v -> Maybe Indirection
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> Instruction -> v -> Maybe Indirection
stry_resolve_error_call ctxt
ctxt ([Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
last [Instruction]
instrs) v
rdi

  is_jump_table_entry :: Maybe (Set ResolvedJumpTarget) -> Bool
is_jump_table_entry Maybe (Set ResolvedJumpTarget)
Nothing      = Bool
False
  is_jump_table_entry (Just Set ResolvedJumpTarget
trgts) = Set ResolvedJumpTarget -> Int
forall a. Set a -> Int
S.size Set ResolvedJumpTarget
trgts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& (ResolvedJumpTarget -> Bool) -> Set ResolvedJumpTarget -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ResolvedJumpTarget -> Bool
is_immediate Set ResolvedJumpTarget
trgts
  is_immediate :: ResolvedJumpTarget -> Bool
is_immediate (ImmediateAddress Word64
_) = Bool
True
  is_immediate ResolvedJumpTarget
_                    = Bool
False
  get_immediate :: ResolvedJumpTarget -> Word64
get_immediate (ImmediateAddress Word64
a) = Word64
a

  clean_sstate :: Sstate v p -> Sstate v p
clean_sstate (Sstate Map Register v
sregs Map (p, Maybe ByteSize) v
smem FlagStatus
fs) = Map Register v
-> Map (p, Maybe ByteSize) v -> FlagStatus -> Sstate v p
forall v p.
Map Register v
-> Map (p, Maybe ByteSize) v -> FlagStatus -> Sstate v p
Sstate Map Register v
sregs (Map (p, Maybe ByteSize) v -> Map (p, Maybe ByteSize) v
forall {a} {b}. Map (a, b) v -> Map (a, b) v
clean_smem Map (p, Maybe ByteSize) v
smem) FlagStatus
fs
  clean_smem :: Map (a, b) v -> Map (a, b) v
clean_smem = ((a, b) -> v -> Bool) -> Map (a, b) v -> Map (a, b) v
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\(a
a,b
si) v
v -> ctxt -> v -> Bool
forall ctxt v p.
WithAbstractSymbolicValues ctxt v p =>
ctxt -> v -> Bool
sis_deterministic ctxt
ctxt v
v)