content: {-|
  
  This module runs each static and symbolic checker pair and outputs the results.
  The more interesting code is in other files:
  - Static/ contains all automatic static checking code
  - Symex/ contains all automatic symbolic checking code
  
  -}
  module Lib ( -- * All final paper checkers, plus new checker heap oob
               uninit -- * For main
             , doUninitCheck -- * For testing
             , concrOOB -- * For main
             , doConcrOOBCheck -- * For testing
             , heapOOB -- * For main
             , doHeapOOBCheck -- * For testing
             , userInput -- * For main
             , doUserInputCheck -- * For testing
             , uaf -- * For main
               -- * Run your own static or symbolic checking
             , checkFiles
             , checkFile
             , symexPath
             )
  where
  import           Checkers.ConcreteOOBStatic
  import           Checkers.HeapOOBStatic
  import           Checkers.StaticConfigs.CheckerConfigs
  import qualified Checkers.SymexConfigs.CheckerConfigs  as SymConfigs
  import           Checkers.UAFStatic
  import           Checkers.UninitStatic
  import           Checkers.UserInputStatic
  import           Control.Concurrent.MVar
  import           Control.Concurrent.STM                (atomically)
  import           Control.Concurrent.STM.TQueue
  import           Control.Concurrent.Thread.Group
  import           Control.Exception                     as E
  import           Control.Log                           as L
  import           Control.Monad
  import           Data.ByteString                       (ByteString, readFile)
  import qualified Data.Map                              as M
  import           Data.Maybe                            (fromJust, isJust)
  import qualified Data.Set                              as S
  import           InternalIR.ModuleInfo
  import           InternalIR.PathExpand
  import           InternalIR.PathInfo
  import           LLVM.AST
  import           LLVM.Context
  import           LLVMAST.Interface
  import           Prelude                               hiding (log, mod, readFile)
  import           Static.Check
  import           Static.CheckerConfigDef
  import qualified Symex.CheckerConfigDef                as SymConfig
  import           Symex.Symex
  import           System.IO.Unsafe                      (unsafePerformIO)
  import           Utils
  
  --
  -- Uninit checker
  --
  
  -- | Run the uninitialized memory static and symbolic checker
  uninit :: FilePath
         -> [String]
         -> IO ()
  uninit fpath extns = do
    candidates <- checkFiles uninitConfig fpath extns
    candidateGroup <- new
  
    forM_ candidates $ \bug ->
      void $ forkIO candidateGroup $ E.handle printHandler $ do
        attackResult <- doUninitCheck bug
        showUninitResult bug attackResult
    wait candidateGroup
  
  doUninitCheck :: UninitBug -> IO SolverResult
  doUninitCheck (UninitBug idx path _ fname fp var ty) = do
    modInfo <- makeModuleInfo fp
    let staticPath = makeIntraproStaticPathStop modInfo fname path idx
    symexPath (SymConfigs.uninitSymexConfig var ty) [fname] [modInfo] staticPath
  
  --
  -- heap out of bounds checker
  --
  
  -- | Run the heap out of bounds static and symbolic checker
  heapOOB :: FilePath
          -> [String]
          -> IO ()
  heapOOB fpath extns = do
    candidates <- checkFiles heapoobConfig fpath extns
    candidateGroup <- new
  
    forM_ candidates $ \bug ->
     void $ forkIO candidateGroup $ E.handle printHandler $ do
        result <- doHeapOOBCheck bug
        showMallocResult bug result
    wait candidateGroup
  
  doHeapOOBCheck :: MOOBBug -> IO SolverResult
  doHeapOOBCheck (MOOBBug (OOBInfo _ asize isize jsize ig _ _ _) fp fun p _) = do
    modInfo <- makeModuleInfo fp
    let staticPath = makeIntraproStaticPath modInfo fun p
    symexPath (SymConfigs.heapoobSymexConfig asize isize jsize ig) [fun] [modInfo] staticPath
  
  --
  -- User input taint checker
  --
  
  -- | Run the user input static and symbolic checker
  userInput :: FilePath
            -> [String]
            -> IO ()
  userInput fpath extns = do
    userInputIndecies <- checkFiles userInputConfig fpath extns
    candidateGroup <- new
    forM_ userInputIndecies $ \bug ->
      void $ forkIO candidateGroup $ E.handle printHandler $ do
      attackResult <- doUserInputCheck bug
      showUserInputResult bug attackResult
    wait candidateGroup
  
  doUserInputCheck :: InputBug -> IO SolverResult
  doUserInputCheck (InputBug fp fname path _ val var) = do
    modInfo <- makeModuleInfo fp
    let staticPath = makeIntraproStaticPath modInfo fname path
    symexPath (SymConfigs.userInputSymexConfig var val) [fname] [modInfo] staticPath
  
  --
  -- Concrete out of bounds checker
  --
  
  -- | Run the concrete out of bounds checker
  concrOOB :: FilePath
           -> [String]
           -> IO ()
  concrOOB fpath extns = do
    candidates <- checkFiles concroobConfig fpath extns
    candidateGroup <- new
    forM_ candidates $ \ bug ->
      void $ forkIO candidateGroup $ E.handle printHandler $ do
      attackResult <- doConcrOOBCheck bug
      showNegBug bug attackResult
    wait candidateGroup
  
  doConcrOOBCheck :: COOBBug -> IO SolverResult
  doConcrOOBCheck (COOBBug _ negIdx negAmt path fname fp) = do
    modInfo <- makeModuleInfo fp
    let staticPath = makeIntraproStaticPath modInfo fname path
    symexPath (SymConfigs.concroobSymexConfig negIdx negAmt) [fname] [modInfo] staticPath
  
  --
  -- UAF checker
  --
  
  -- | Run the UAF static and symbolic checker
  uaf :: FilePath
      -> [String]
      -> IO ()
  uaf fpath extns = do
    candidates <- checkFiles uafConfig fpath extns
    candidateGroup <- new
  
    forM_ candidates $ \bug ->
      void $ forkIO candidateGroup $ E.handle printHandler $ do
        -- L.log L.INFO $ "Symbolically checking static alert: " ++ show bug
        attackResult <- doUAFCheck bug
        showUAFResult bug attackResult
    wait candidateGroup
  
  doUAFCheck :: UAFBug -> IO SolverResult
  doUAFCheck (UAFBug fp fname path _ var ty _) = do
    modInfo <- makeModuleInfo fp
    let staticPath = makeIntraproStaticPath modInfo fname path
    symexPath (SymConfigs.uafSymexConfig var ty) [fname] [modInfo] staticPath
  
  --
  -- Making static paths
  --
  
  makeIntraproStaticPathStop :: ModuleInfo
                             -> Name
                             -> [Name]
                             -> Int
                             -> StaticPath
  makeIntraproStaticPathStop modInfo funName path idx =
    makeSingleStaticPath modInfo funName path (Just idx)
  
  makeIntraproStaticPath :: ModuleInfo
                         -> Name
                         -> [Name]
                         -> StaticPath
  makeIntraproStaticPath modInfo funName path =
    makeSingleStaticPath modInfo funName path Nothing
  
  makeSingleStaticPath :: ModuleInfo
                         -> Name
                         -> [Name]
                         -> Maybe Int
                         -> StaticPath
  makeSingleStaticPath modInfo funName path mTerminalIdx = do
    case bbsFromFunctionName (modAST modInfo) funName of
      Nothing -> error $ "Could not find bbs for this path: " ++ show path
      Just bbs ->
        let bbMap = makeBBMap bbs
        in concatMap  (\(name, pathIndex) ->
                        -- Figure out the previous block
                       let prevBlock = if pathIndex == 0
                                       then Nothing
                                       else Just $ path !! (pathIndex - 1)
                           isLastBlock = pathIndex + 1 >= length path
                        -- Get the individual instruction entries
                           entry idx = StaticPathEntry modInfo funName name idx prevBlock
                           block = bbMap M.! name
                        -- Figure out which index to finish at:
                        -- If we're the last block and have a terminal index, finish there
                        -- Otherwise, go through each line of the block
                           maxEntry = if isLastBlock && isJust mTerminalIdx
                                      then (fromJust mTerminalIdx) - 1
                                      else (length $ getBlockContents block) -1
                           entries = map (\idxNum ->
                                           entry $ SPEInstr idxNum
                                          ) $ [0..maxEntry]
                        -- If there is a next block, add a terminator entry indicating it
                        in if isLastBlock
                           -- USED TO MAKE A TERMINAL FUNCTION CALL (?), NOW JUST TERMINATES
                           then if isJust mTerminalIdx
                                then let newEntry = SPEInstr (fromJust mTerminalIdx)
                                     in entries ++ [entry newEntry]
                                else entries
                           else let nextBlock = path !! (pathIndex + 1)
                                in entries ++ [entry $ SPETerm nextBlock]
                ) $ zip path [0..]
  
  --
  -- Static/symbolic execution:
  -- These three functions (checkFile, checkFile, and symexPath) allow
  -- you to (1) statically check some files for suspicious elements and
  -- then (2) symbolically execute any suspicious paths.
  --
  
  -- | Statically check all files in a directory
  checkFiles :: (Show a, Show b, Ord b)
              => CheckerConfig a b
              -> FilePath
              -> [String]
              -> IO (S.Set b)
  checkFiles ckConfig dirPath extns = do
    fileGroup  <- new
    resultQueue <- newTQueueIO
    sourceFiles <- getSourceFiles dirPath extns
    forM_ sourceFiles $ \file -> void $ forkIO fileGroup $
                                 checkFile resultQueue fileGroup ckConfig file
    wait fileGroup
    S.fromList <$> (atomically $ flushTQueue resultQueue)
  
  -- | Statically check a single file
  checkFile :: (Show a, Show b, Ord b)
            => TQueue b
            -> ThreadGroup
            -> CheckerConfig a b
            -> FilePath
            -> IO ()
  checkFile resultQueue fileGroup ckConfig path = do
    L.log L.INFO $ unwords $ ["Checking file...", path]
    contents <- readFile path
    when (cfgShouldCheckFile ckConfig contents) $ do
      withContext $ \ctx -> E.handle printHandler $ do
        modInfo <- getModInfoCache path ctx contents
        let mod =  modAST modInfo
        when (cfgShouldCheckModule ckConfig $ modAST modInfo) $ do
          L.log L.DEBUG $ unwords $ ["Checking module...", show $ moduleName mod]
          forM_ (getFunctionDefns mod) $ \defn -> do
            when (cfgShouldCheckFunction ckConfig defn) $ void $ forkIO fileGroup $ do
              L.log L.DEBUG $ unwords $ ["Checking function...", show $ getFunctionName' defn]
              staticCheckFunctionAsync resultQueue fileGroup path modInfo defn ckConfig
      where getFunctionDefns :: Module -> [Definition]
            getFunctionDefns = filter isFunction . moduleDefinitions
  
  symexPath :: SymConfig.Config a
            -> [Name] -- ^ Names, with the start function first
            -> [ModuleInfo] -- ^ Module infos, with the start module info first
            -> StaticPath
            -> IO SolverResult
  symexPath config names modInfos staticPath = do
    let path = staticToSimple staticPath
        pl = SymConfig.pathLength config
        ab = SymConfig.arrayBound config
        pathInfos = map (\(mi, name) -> getPathInfo mi pl name path ab) $ zip modInfos names
        pathInfo = foldr mergePathInfos (head pathInfos) (tail pathInfos)
    symex config path pathInfo
  
  moduleInfoMap :: MVar (M.Map FilePath ModuleInfo)
  {-# NOINLINE moduleInfoMap #-}
  moduleInfoMap = unsafePerformIO $ newMVar M.empty
  
  getModInfoCache :: FilePath -> Context -> ByteString -> IO ModuleInfo
  getModInfoCache path ctx contents = do
    map0 <- readMVar moduleInfoMap
    case M.lookup path map0 of
      Just mInfo -> return mInfo
      Nothing -> do mInfo <- getModuleInfo ctx contents
                    modifyMVar_ moduleInfoMap (\map1 -> return $ M.insert path mInfo map1)
                    return mInfo
  
  makeModuleInfo :: FilePath -> IO ModuleInfo
  makeModuleInfo path = do
    map0 <- readMVar moduleInfoMap
    case M.lookup path map0 of
      Just mInfo -> return mInfo
      Nothing -> do mInfo <- withContext $ \ctx -> E.handle handler $ getModuleInfoFromFile ctx path
                    modifyMVar_ moduleInfoMap (\map1 -> return $ M.insert path mInfo map1)
                    return mInfo
   where
     handler :: E.SomeException -> IO ModuleInfo
     handler e = error $ show e
  
  printHandler :: E.SomeException -> IO ()
  printHandler = print
  