diff --git a/reopt/Main_reopt.hs b/reopt/Main_reopt.hs index c5da08ed..b5548acd 100644 --- a/reopt/Main_reopt.hs +++ b/reopt/Main_reopt.hs @@ -189,8 +189,12 @@ data Args = Args -- ^ Path to llvm-mc -- -- Only used when generating object file from assembly generated by llc. + , addEntryPoints :: ![String] + -- ^ List of extra entry points to start discovery from. This does not + -- preclude other entry points from being added through discovery. , includeAddrs :: ![String] - -- ^ List of entry points for translation + -- ^ List of allowed function entry points for discovery. This prevents + -- exploration of entry points beyond the ones explicitly listed. , excludeAddrs :: ![String] -- ^ List of function entry points that we exclude for translation. , loadBaseAddress :: !(Maybe Word64) @@ -249,6 +253,7 @@ defaultArgs = , optLevel = 2 , slashPath = "slash" , llvmMcPath = "llvm-mc" + , addEntryPoints = [] , includeAddrs = [] , excludeAddrs = [] , loadBaseAddress = Nothing @@ -440,7 +445,16 @@ slashPathP = pathP "slash" #slashPath llvmMcPathP :: Parser String llvmMcPathP = pathP "llvm-mc" #llvmMcPath --- | Used to add a new function to ignore translation of. +-- | Used to add extra entry points to seed discovery. +addEntryPointP :: Parser String +addEntryPointP = + strOption + ( long "add-entry-point" + <> metavar "ADDR" + <> help "Extra address to seed discovery frontier with (may be repeated)" + ) + +-- | Used to specify the only addresses we care to explore. includeAddrP :: Parser String includeAddrP = strOption @@ -449,7 +463,7 @@ includeAddrP = <> help "Address of function to include in analysis (may be repeated)" ) --- | Used to add a new function to ignore translation of. +-- | Used to specify addresses we wish **not** to explore. excludeAddrP :: Parser String excludeAddrP = strOption @@ -606,6 +620,7 @@ arguments = <*> optLevelP <*> slashPathP <*> llvmMcPathP + <*> many addEntryPointP <*> many includeAddrP <*> many excludeAddrP <*> optional loadBaseAddressP @@ -674,7 +689,8 @@ argsReoptOptions args = do gdbDebugDirs <- getGdbDebugInfoDirs True pure $ ReoptOptions - { roIncluded = includeAddrs args + { roExtraEntryPoints = addEntryPoints args + , roIncluded = includeAddrs args , roExcluded = excludeAddrs args , roVerboseMode = True , roDiscoveryOptions = args ^. #discOpts diff --git a/src/Reopt.hs b/src/Reopt.hs index f2efe30e..0d4af755 100644 --- a/src/Reopt.hs +++ b/src/Reopt.hs @@ -440,7 +440,9 @@ reoptDefaultDiscoveryOptions = -- | Information from user to control which addresses to include and -- exclude. data ReoptOptions = ReoptOptions - { roIncluded :: [String] + { roExtraEntryPoints :: [String] + -- ^ Additional entry points we want Reopt to consider + , roIncluded :: [String] -- ^ Symbols/addresses user wanted included , roExcluded :: [String] -- ^ Symbols/addresses user wanted exluded. @@ -462,7 +464,8 @@ data ReoptOptions = ReoptOptions defaultReoptOptions :: ReoptOptions defaultReoptOptions = ReoptOptions - { roIncluded = [] + { roExtraEntryPoints = [] + , roIncluded = [] , roExcluded = [] , roVerboseMode = False , roDiscoveryOptions = reoptDefaultDiscoveryOptions @@ -970,7 +973,7 @@ initDiscState mem initPoints regInfo symAddrMap explorePred ainfo reoptOpts = do let resolveEntry qsn | ".cold" `BS.isSuffixOf` qsnBytes qsn = Nothing | otherwise = Just Macaw.MayReturnFun - let entryPoints0 = + let noReturnEntryPoints = Map.mapMaybe resolveEntry (samAddrMap symAddrMap) & addKnownFn symAddrMap "abort" Macaw.NoReturnFun & addKnownFn symAddrMap "exit" Macaw.NoReturnFun @@ -980,7 +983,9 @@ initDiscState mem initPoints regInfo symAddrMap explorePred ainfo reoptOpts = do & addKnownFn symAddrMap "__malloc_assert" Macaw.NoReturnFun & addKnownFn symAddrMap "__stack_chk_fail" Macaw.NoReturnFun & addKnownFn symAddrMap "_ZSt9terminatev" Macaw.NoReturnFun - let entryPoints = foldl (\m a -> Map.insert a Macaw.MayReturnFun m) entryPoints0 initPoints + extraEntryPoints <- mapM (resolveSymAddr mem regInfo symAddrMap) (roExtraEntryPoints reoptOpts) + let mayReturnEntryPoints = initPoints ++ extraEntryPoints + let entryPoints = foldl (\m a -> Map.insert a Macaw.MayReturnFun m) noReturnEntryPoints mayReturnEntryPoints case (roIncluded reoptOpts, roExcluded reoptOpts) of ([], excludeNames) -> do excludeAddrs <- mapM (resolveSymAddr mem regInfo symAddrMap) excludeNames @@ -997,6 +1002,9 @@ initDiscState mem initPoints regInfo symAddrMap explorePred ainfo reoptOpts = do let initState = Macaw.emptyDiscoveryState mem (getAddrSymMap symAddrMap) ainfo & Macaw.trustedFunctionEntryPoints .~ entryPoints + -- NOTE (val) It looks a bit weird that we're not also checking + -- `explorePred a` here. Not sure that's intended, and it's + -- definitely not documented. & Macaw.exploreFnPred .~ (`Set.member` s) & Macaw.markAddrsAsFunction Macaw.InitAddr s pure $! initState