echidna/lib/Echidna/Solidity.hs

{-# LANGUAGE FlexibleContexts, OverloadedStrings, LambdaCase, TupleSections #-}


module Echidna.Solidity where

import Control.Lens               ((^.), (%=), _1, assign, use, view)
import Control.Exception          (Exception)
import Control.Monad              (liftM2)
import Control.Monad.Catch        (MonadThrow(..))
import Control.Monad.IO.Class     (MonadIO(..))
import Control.Monad.State.Strict (MonadState, execState, modify, runState)
import Data.Foldable              (toList)
import Data.List                  (find, partition)
import Data.Map                   (insert)
import Data.Maybe                 (isNothing, fromMaybe)
import Data.Monoid                ((<>))
import Data.Text                  (Text, isPrefixOf, unpack)
import System.Process             (readProcess)
import System.IO.Temp             (writeSystemTempFile)

import qualified Data.Map as Map (lookup)

import Echidna.ABI (SolSignature)

import EVM
  (Contract, VM, VMResult(..), contract, contracts, env, gas, loadContract, replaceCodeOfSelf, resetState, state)
import EVM.Concrete (Blob(..))
import EVM.Exec     (exec, vmForEthrunCreation)
import EVM.Keccak   (newContractAddress)
import EVM.Solidity (abiMap, contractName, creationCode, methodInputs, methodName, readSolc, SolcContract)
import EVM.Types    (Addr)

data EchidnaException = BadAddr Addr
                      | CompileFailure
                      | NoContracts
                      | TestArgsFound Text
                      | ContractNotFound Text

instance Show EchidnaException where
  show = \case
    BadAddr a            -> "No contract at " ++ show a ++ " exists"
    CompileFailure       -> "Couldn't compile given file"
    NoContracts          -> "No contracts found in given file"
    (ContractNotFound c) -> "Given contract " ++ show c ++ " not found in given file"
    (TestArgsFound t)    -> "Test " ++ show t ++ " has arguments, aborting"

instance Exception EchidnaException

-- | parses additional solc arguments
solcArguments :: FilePath -> Maybe Text -> [String]
solcArguments filePath argStr = args <> fromMaybe [] additional
  where args = ["--combined-json=bin-runtime,bin,srcmap,srcmap-runtime,abi,ast", filePath]
        additional = words . unpack <$> argStr

-- | reads all contracts within the solidity file at `filepath` and passes optional solc params to compiler
readContracts :: (MonadIO m, MonadThrow m) => FilePath -> Maybe Text -> m [SolcContract]
readContracts filePath solcArgs = liftIO solc >>= \case
  Nothing -> throwM CompileFailure
  Just m  -> return $ toList $ fst m
  where solc = readSolc =<< writeSystemTempFile "" =<< readProcess
          "solc" (solcArguments filePath solcArgs) ""

-- | reads either the first contract found or the contract named `selectedContractName` within the solidity file at `filepath`
readContract :: (MonadIO m, MonadThrow m) => FilePath -> Maybe Text -> Maybe Text -> m SolcContract
readContract filePath selectedContractName solcArgs = do
    cs <- readContracts filePath solcArgs
    c <- chooseContract cs selectedContractName
    warn (isNothing selectedContractName && 1 < length cs) $
      "Multiple contracts found in file, only analyzing the first"
    liftIO $ print $ "Analyzing contract: " <> c ^. contractName
    return c
  where chooseContract :: (MonadThrow m) => [SolcContract] -> Maybe Text -> m SolcContract
        chooseContract [] _ = throwM NoContracts
        chooseContract (c:_) Nothing = return c
        chooseContract cs (Just name) = case find (\x -> name == x ^. contractName) cs of
          Nothing -> throwM $ ContractNotFound name
          Just c  -> return c
        warn :: (MonadIO m) => Bool -> Text -> m ()
        warn p s = if p then liftIO $ print s else pure ()

-- | loads the solidity file at `filePath` and selects either the default or specified contract to analyze
loadSolidity :: (MonadIO m, MonadThrow m) => FilePath -> Maybe Text -> Maybe Text -> m (VM, [SolSignature], [Text])
loadSolidity filePath selectedContract solcArgs = do
    c <- readContract filePath selectedContract solcArgs
    let (VMSuccess (B bc), vm) = runState exec . vmForEthrunCreation $ c ^. creationCode
        load = do resetState
                  assign (state . gas) 0xffffffffffffffff
                  loadContract (vm ^. state . contract)
        loaded = execState load $ execState (replaceCodeOfSelf bc) vm
        abi = map (liftM2 (,) (view methodName) (map snd . view methodInputs)) . toList $ c ^. abiMap
        (tests, funs) = partition (isPrefixOf "echidna_" . fst) abi
    case find (not . null . snd) tests of
      Nothing      -> return (loaded, funs, fst <$> tests)
      (Just (t,_)) -> throwM $ TestArgsFound t


insertContract :: MonadState VM m => Contract -> m ()
insertContract c = do a <- (`newContractAddress` 1) <$> use (state . contract)
                      env . contracts %= insert a c
                      modify . execState $ loadContract a

currentContract :: MonadThrow m => VM -> m Contract
currentContract v = let a = v ^. state . contract in
  maybe (throwM $ BadAddr a) pure . Map.lookup a $ v ^. env . contracts

addSolidity :: (MonadIO m, MonadThrow m, MonadState VM m) => FilePath -> Maybe Text -> Maybe Text -> m ()
addSolidity f mc ma = insertContract =<< currentContract =<< view _1 <$> loadSolidity f mc ma