express recovery intent to reoptPrepareForRecovery

GaloisInc · Sep 27, 2023 · a6e9bf7 · a6e9bf7
1 parent 58b9920
commit a6e9bf7
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 9 deletions.
diff --git a/reopt/Main_reopt.hs b/reopt/Main_reopt.hs
@@ -692,7 +692,13 @@ showCFG args elfPath = do
     runReoptM printLogEvent $ do
       hdrAnn <- resolveHeader (headerPath args) (clangPath args)
       (_, _, _, discState) <-
-        reoptPrepareForRecovery (loadOptions args) reoptOpts hdrAnn (unnamedFunPrefix args) hdrInfo
+        reoptPrepareForRecovery
+          (loadOptions args)
+          reoptOpts
+          hdrAnn
+          (unnamedFunPrefix args)
+          hdrInfo
+          dontIntendToRecover
       pure $ show $ Macaw.ppDiscoveryStateBlocks discState
   handleEitherWithExit mr
 
@@ -714,7 +720,7 @@ showConstraints args elfPath = do
       funPrefix = unnamedFunPrefix args
 
     (os, symAddrMap, debugTypeMap, discState) <-
-      reoptPrepareForRecovery (loadOptions args) rOpts hdrAnn funPrefix origElf
+      reoptPrepareForRecovery (loadOptions args) rOpts hdrAnn funPrefix origElf dontIntendToRecover
 
     let sysp = osPersonality os
     recoverX86Output <-
@@ -818,10 +824,13 @@ performReopt args elfPath = do
       funPrefix = unnamedFunPrefix args
 
     (os, symAddrMap, debugTypeMap, discState) <-
-      reoptPrepareForRecovery (loadOptions args) rOpts hdrAnn funPrefix origElf
-
-    when (shouldRecover args) $
-      checkNoSymbolUsesReservedPrefix funPrefix symAddrMap
+      reoptPrepareForRecovery
+        (loadOptions args)
+        rOpts
+        hdrAnn
+        funPrefix
+        origElf
+        (IntendToRecover $ shouldRecover args)
 
     case cfgExportPath args of
       Nothing -> pure ()

diff --git a/src/Reopt.hs b/src/Reopt.hs
@@ -53,6 +53,9 @@ module Reopt (
   -- * Function recovery
   Reopt.TypeInference.HeaderTypes.AnnDeclarations,
   Reopt.TypeInference.HeaderTypes.emptyAnnDeclarations,
+  IntendToRecover (IntendToRecover),
+  intendToRecover,
+  dontIntendToRecover,
   RecoveredModule,
   recoveredDefs,
   reoptPrepareForRecovery,
@@ -2699,6 +2702,11 @@ reoptRecoveryLoop symAddrMap rOpts funPrefix sysp debugTypeMap = go
       then traceM "NOLOOP" >> return (newDiscState, recoverX86Output, recMod, moduleConstraints)
       else traceM "LOOP" >> go newDiscState
 
+newtype IntendToRecover = IntendToRecover {getIntendToRecover :: Bool}
+intendToRecover, dontIntendToRecover :: IntendToRecover
+intendToRecover = IntendToRecover True
+dontIntendToRecover = IntendToRecover False
+
 -- | Factors out some pre-computation we frequently do prior to running the
 -- Reopt recovery.
 reoptPrepareForRecovery ::
@@ -2708,11 +2716,13 @@ reoptPrepareForRecovery ::
   -- | Prefix to use when encountering unnamed functions
   BSC.ByteString ->
   Elf.ElfHeaderInfo 64 ->
+  -- | Whether we should perform checks relative to recovery
+  IntendToRecover ->
   ReoptM X86_64 r (X86OS, SymAddrMap 64, FunTypeMaps 64, Macaw.DiscoveryState X86_64)
-reoptPrepareForRecovery loadOpts reoptOpts hdrAnn funPrefix hdrInfo = do
+reoptPrepareForRecovery loadOpts reoptOpts hdrAnn funPrefix hdrInfo recIntent = do
   (os, initState) <- reoptX86Init loadOpts reoptOpts hdrInfo
   let symAddrMap = initDiscSymAddrMap initState
-  checkNoSymbolUsesReservedPrefix funPrefix symAddrMap
+  when (getIntendToRecover recIntent) $ checkNoSymbolUsesReservedPrefix funPrefix symAddrMap
   let ainfo = osArchitectureInfo os
   (debugTypeMap, discState) <- doDiscovery hdrAnn hdrInfo ainfo initState reoptOpts
   return (os, symAddrMap, debugTypeMap, discState)
@@ -2737,7 +2747,8 @@ recoverX86Elf ::
     , ModuleConstraints X86_64
     )
 recoverX86Elf loadOpts reoptOpts hdrAnn unnamedFunPrefix hdrInfo = do
-  (os, symAddrMap, debugTypeMap, discState) <- reoptPrepareForRecovery loadOpts reoptOpts hdrAnn unnamedFunPrefix hdrInfo
+  (os, symAddrMap, debugTypeMap, discState) <-
+    reoptPrepareForRecovery loadOpts reoptOpts hdrAnn unnamedFunPrefix hdrInfo intendToRecover
   let sysp = osPersonality os
   (finalDiscState, recoverX86Output, recMod, moduleConstraints) <-
     reoptRecoveryLoop symAddrMap reoptOpts unnamedFunPrefix sysp debugTypeMap discState