Foundry.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE TemplateHaskell   #-}

module Echidna.Output.Foundry (foundryTest) where

import Data.Aeson (Value(..), object, (.=))
import Data.List (elemIndex, nub)
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Text (Text, unpack)
import Data.Text.Lazy (fromStrict)
import qualified Data.ByteString.Base16 as B16
import qualified Data.ByteString.Char8 as C8
import Data.Vector as V hiding ((++), map, zipWith, elemIndex, mapMaybe, concatMap, length)
import Numeric (showHex)
import qualified Data.Text.Lazy as TL
import Text.Mustache (Template, substituteValue, toMustache)
import Text.Mustache.Compile (embedTemplate)

import EVM.ABI (AbiValue(..))
import EVM.Types (W256, Addr)

import Echidna.Types.Test (EchidnaTest(..), TestType(..))
import Echidna.Types.Tx (Tx(..), TxCall(..))

template :: Template
template = $(embedTemplate ["lib/Echidna/Output/assets"] "foundry.mustache")

-- | Generate a Foundry test from an EchidnaTest result.
-- For property tests, psender is the address used to call the property function.
foundryTest :: Maybe Text -> Maybe Addr -> EchidnaTest -> TL.Text
foundryTest mContractName mPsender test =
  case test.testType of
    AssertionTest{} ->
      let testData = createTestData mContractName Nothing test
      in fromStrict $ substituteValue template (toMustache testData)
    PropertyTest name _ ->
      let testData = createTestData mContractName (Just (name, mPsender)) test
      in fromStrict $ substituteValue template (toMustache testData)
    _ -> ""

-- | Create an Aeson Value from test data for the Mustache template.
-- When a property name and psender are provided, a final assertion is added
-- to call the property from psender and check it returns false.
createTestData :: Maybe Text -> Maybe (Text, Maybe Addr) -> EchidnaTest -> Value
createTestData mContractName mProperty test =
  let
    senders = nub $ map (.src) test.reproducer
    actors = zipWith actorObject senders [1..]
    repro = mapMaybe (foundryTx senders) test.reproducer
    cName = fromMaybe "YourContract" mContractName
    propAssertion = mProperty >>= \(name, mAddr) ->
      let prank = case mAddr of
            Just addr -> "        vm.stopPrank();\n        vm.prank(" ++ formatAddr addr ++ ");\n"
            Nothing   -> "        vm.stopPrank();\n"
      in Just $ prank ++ "        assertFalse(Target." ++ unpack name ++ "());"
  in
  object
    [ "testName"     .= ("FoundryTest" :: Text)
    , "contractName" .= cName
    , "actors"       .= actors
    , "reproducer"   .= repro
    , "propertyAssertion" .= propAssertion
    ]

-- | Create a JSON object for an actor.
actorObject :: Addr -> Int -> Value
actorObject sender i = object
  [ "name"    .= ("USER" ++ show i :: String)
  , "address" .= formatAddr sender
  ]

-- | Format an address for Solidity.
formatAddr :: Addr -> String
formatAddr addr = "address(0x" ++ showHex (fromIntegral addr :: W256) "" ++ ")"

-- | Generate a single transaction line for the reproducer.
foundryTx :: [Addr] -> Tx -> Maybe Value
foundryTx senders tx =
  case tx.call of
    SolCall (name, args) ->
      let
        (time, blocks) = tx.delay
        senderName =
          case elemIndex tx.src senders of
            Just i -> "USER" ++ show (i + 1)
            Nothing -> formatAddr tx.src
        prelude =
          (if time > 0 || blocks > 0 then "    _delay(" ++ show time ++ ", " ++ show blocks ++ ");\n" else "") ++
          "    _setUpActor(" ++ senderName ++ ");"
        -- Handle fallback function (empty name).
        call = if unpack name == ""
          then "    address(Target).call(\"\");"
          else "    Target." ++ unpack name ++ "(" ++ foundryArgs (map abiValueToString args) ++ ");"
      in Just $ object ["prelude" .= prelude, "call" .= call]
    _ -> Nothing

-- | Format arguments for a Solidity call.
foundryArgs :: [String] -> String
foundryArgs [] = ""
foundryArgs [x] = x
foundryArgs (x:xs) = x ++ ", " ++ foundryArgs xs

-- | Convert an AbiValue to its string representation for Solidity.
abiValueToString :: AbiValue -> String
abiValueToString (AbiUInt _ w) = show w
abiValueToString (AbiInt _ w) = show w
abiValueToString (AbiAddress a) = "address(0x" ++ showHex (fromIntegral a :: W256) "" ++ ")"
abiValueToString (AbiBool b) = if b then "true" else "false"
abiValueToString (AbiBytes _ bs) = "hex\"" ++ C8.unpack (B16.encode bs) ++ "\""
abiValueToString (AbiString s) = show s
abiValueToString (AbiTuple vs) = "(" ++ foundryArgs (map abiValueToString (V.toList vs)) ++ ")"
abiValueToString _ = ""