Decompile: basic storage rewrites supported

src/Decompile.hs

@@ -19,20 +19,19 @@ This module decompiles EVM bytecode into an Act spec. It operates as follows
 -}
 module Decompile where
 
-import Debug.Trace
-
 import Prelude hiding (LT, GT)
 
 import Control.Monad.Except
 import Control.Monad.Extra
-import Data.Bifunctor
+import Data.List
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Maybe (fromJust)
 import Data.Text (Text)
 import Data.Text qualified as T
+import Data.Tuple.Extra
 import Data.Typeable
 import EVM.ABI (AbiKind(..), abiKind)
 import EVM.Fetch qualified as Fetch
@@ -69,9 +68,10 @@ summarize solvers contract = do
   runExceptT $ do
     ctor <- ExceptT creation
     behvs <- ExceptT runtime
+    -- TODO: raise error if we have packed slots
     layout <- ExceptT . pure . toErr "missing storage layout in solc output" $ contract.storageLayout
     pure $ EVMContract
-      { name = contract.contractName
+      { name = snd $ T.breakOnEnd ":" contract.contractName
       , storageLayout = layout
       , runtime = behvs
       , creation = ctor
@@ -120,16 +120,13 @@ mkStore :: EVMContract -> Store
 mkStore c = Map.singleton (T.unpack c.name) (Map.mapKeys T.unpack $ fmap fromitem c.storageLayout)
   where
     fromitem item = (convslot item.slotType, toInteger item.slot)
-    convslot (EVM.StorageMapping a b) = StorageMapping (fmap PrimitiveType a) (PrimitiveType b)
-    convslot (EVM.StorageValue a) = StorageValue (PrimitiveType a)
 
 -- | Build an act constructor spec from the summarized bytecode
 mkConstructor :: EVMContract -> Either Text Constructor
 mkConstructor cs
   | Set.size (snd cs.creation) == 1 =
       case head (Set.elems (snd cs.creation)) of
         EVM.Success props _ _ _ -> do
-          traceShowM props
           ps <- mapM fromProp props
           pure $ Constructor
             { _cname = T.unpack cs.name
@@ -148,7 +145,7 @@ mkBehvs :: EVMContract -> Either Text [Behaviour]
 mkBehvs c = concatMapM (\(i, bs) -> mapM (mkbehv i) (Set.toList bs)) (Map.toList c.runtime)
   where
     mkbehv :: Method -> EVM.Expr EVM.End -> Either Text Behaviour
-    mkbehv method (EVM.Success props _ retBuf _) = do
+    mkbehv method (EVM.Success props _ retBuf state) = do
       pres <- mapM fromProp props
       ret <- case method.output of
         [] -> Right Nothing
@@ -161,18 +158,72 @@ mkBehvs c = concatMapM (\(i, bs) -> mapM (mkbehv i) (Set.toList bs)) (Map.toList
               pure . Just . TExp SInteger $ v
           Dynamic -> Left "cannot decompile methods that return dynamically sized types"
         _ -> Left "cannot decompile methods with multiple return types"
+      rewrites <- case Map.toList state of
+        [(EVM.SymAddr "entrypoint", contract)] -> do
+          partitioned <- partitionStorage contract.storage
+          mkRewrites c.name (invertLayout c.storageLayout) partitioned
+        [(_, _)] -> error $ "Internal Error: state contains a single entry for an unexpected contract:\n" <> show state
+        [] -> error "Internal Error: unexpected empty state"
+        _ -> Left "cannot decompile methods that update storage on other contracts"
       pure $ Behaviour
         { _contract = T.unpack c.name
         , _interface = behvIface method
         , _name = T.unpack method.name
         , _preconditions = pres
         , _caseconditions = mempty -- TODO: what to do here?
         , _postconditions = mempty
-        , _stateUpdates = mempty -- TODO
+        , _stateUpdates = rewrites
         , _returns = ret
         }
     mkbehv _ _ = Left "unexpected unsucessful branch"
 
+-- TODO: we probably need to diff against the prestore or smth to be sound here
+mkRewrites :: Text -> Map (Integer, Integer) (Text, SlotType) -> DistinctStore -> Either Text [Rewrite]
+mkRewrites cname layout (DistinctStore writes) = forM (Map.toList writes) $ \(slot,item) ->
+    case Map.lookup (toInteger slot, 0) layout of
+      Just (name, typ) -> case typ of
+        StorageValue v -> case v of
+          PrimitiveType t -> case abiKind t of
+            Static -> case t of
+              AbiTupleType _ -> Left "cannot decompile methods that write to tuple in storage"
+              AbiFunctionType -> Left "cannot decompile methods that store function pointers"
+              _ -> do
+                  val <- fromWord item
+                  pure (Rewrite (Update SInteger (Item SInteger v (SVar nowhere (T.unpack cname) (T.unpack name))) val))
+            Dynamic -> Left "cannot decompile methods that store dynamically sized types"
+          ContractType {} -> Left "cannot decompile contracts that have contract types in storage"
+        StorageMapping {} -> Left "cannot decompile contracts that write to mappings"
+      Nothing -> Left $ "write to a storage location that is not mentioned in the solc layout: " <> (T.pack $ show slot)
+
+newtype DistinctStore = DistinctStore (Map (EVM.W256) (EVM.Expr EVM.EWord))
+
+-- | Attempts to decompose an input storage expression into a map from provably distinct keys to abstract words
+-- currently only supports stores where all writes are to concrete locations
+-- supporting writes to symbolic locations would requires calling an smt solver
+partitionStorage :: EVM.Expr EVM.Storage -> Either Text DistinctStore
+partitionStorage = go mempty
+  where
+    go :: Map EVM.W256 (EVM.Expr EVM.EWord) -> EVM.Expr EVM.Storage -> Either Text DistinctStore
+    go curr = \case
+      EVM.AbstractStore _ -> pure $ DistinctStore curr
+      EVM.ConcreteStore s -> do
+        let s' = Map.toList (fmap EVM.Lit s)
+            new = foldl' checkedInsert curr s'
+        pure $ DistinctStore new
+      EVM.SStore (EVM.Lit k) v base -> go (checkedInsert curr (k,v)) base
+      EVM.SStore {} -> Left "cannot decompile contracts with writes to symbolic storage slots"
+      EVM.GVar _ -> error "Internal Error: unexpected GVar"
+
+    -- this is safe because:
+    --   1. we traverse top down
+    --   2. we only consider lit keys in go
+    checkedInsert curr (key, val) = case Map.lookup key curr of
+      -- if a key was already written to higher in the write chain, ignore this write
+      Just _ -> curr
+      -- if this is the first time we have seen a key then insert it
+      Nothing -> Map.insert key val curr
+
+
 -- | Convert an HEVM Prop into a boolean Exp
 fromProp :: EVM.Prop -> Either Text (Exp ABoolean)
 fromProp p = case p of
@@ -204,7 +255,9 @@ fromWord w = case w of
 
 -- Verification ------------------------------------------------------------------------------------
 
+
 -- | Verify that the decompiled spec is equivalent to the input bytecodes
+-- This compiles the generated act spec back down to an Expr and then checks that the two are equivalent
 verifyDecompilation :: ByteString -> ByteString -> Act -> IO ()
 verifyDecompilation creation runtime spec =
   withSolvers CVC5 4 Nothing $ \solvers -> do
@@ -230,6 +283,16 @@ ctorIface args = Interface "constructor" (fmap (\(n, t) -> Decl t (T.unpack n))
 behvIface :: Method -> Interface
 behvIface method = Interface (T.unpack method.name) (fmap (\(n, t) -> Decl t (T.unpack n)) method.inputs)
 
+convslot :: EVM.SlotType -> SlotType
+convslot (EVM.StorageMapping a b) = StorageMapping (fmap PrimitiveType a) (PrimitiveType b)
+convslot (EVM.StorageValue a) = StorageValue (PrimitiveType a)
+
+invertLayout :: Map Text StorageItem -> Map (Integer, Integer) (Text, SlotType)
+invertLayout = Map.fromList . fmap go . Map.toList
+  where
+    go :: (Text, StorageItem) -> ((Integer, Integer), (Text, SlotType))
+    go (n, i) = ((toInteger i.slot, toInteger i.offset), (n, convslot i.slotType))
+
 
 -- Repl Stuff --------------------------------------------------------------------------------------