Skip to content

Commit 00b96ee

Browse files
gustavo-griecogustavo-grieco
committed
new command
1 parent b256cc7 commit 00b96ee

File tree

6 files changed

+92
-43
lines changed

6 files changed

+92
-43
lines changed

‎lib/Echidna/Agent/Fuzzer.hs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ import Control.Monad.Trans (lift)
1616
import Control.Monad.IO.Class (MonadIO)
1717
import System.Random (mkStdGen)
1818
import Data.IORef (IORef, writeIORef, readIORef, atomicModifyIORef')
19+
import Data.Text (Text)
20+
import EVM.ABI (AbiValue)
1921
import Data.Map (Map)
2022
import qualified Data.Map as Map
2123
import System.Directory (getCurrentDirectory)
@@ -197,9 +199,12 @@ fuzzerLoop callback vm testLimit bus = do
197199
putStrLn $ "Fuzzer " ++ show workerId ++ ": dumped LCOV coverage."
198200
pure ()
199201
Just (WrappedMessage _ (ToFuzzer tid (PrioritizeFunction funcName))) -> do
202+
-- Deprecated
203+
pure ()
204+
Just (WrappedMessage _ (ToFuzzer tid (FuzzTransaction funcName args))) -> do
200205
workerId <- gets (.workerId)
201206
when (tid == workerId) $ do
202-
modify' $ \s -> s { prioritizedFunctions = funcName : s.prioritizedFunctions }
207+
modify' $ \s -> s { prioritizedFunctions = (funcName, args) : s.prioritizedFunctions }
203208
pure ()
204209
Just (WrappedMessage _ (ToFuzzer tid ClearPrioritization)) -> do
205210
workerId <- gets (.workerId)

‎lib/Echidna/MCP.hs‎

Lines changed: 46 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -57,11 +57,11 @@ inspectCorpusTransactionsTool args env _ _ = do
5757
let corpusList = Set.toList c
5858
startIndex = (page - 1) * pageSize
5959
pageItems = take pageSize $ drop startIndex corpusList
60-
61-
ppSequence (i, txs) =
60+
61+
ppSequence (i, txs) =
6262
printf "Sequence (value: %d):\n%s" i (unlines $ map (ppTx Map.empty) txs)
6363

64-
return $ if null pageItems
64+
return $ if null pageItems
6565
then "No more transactions found."
6666
else intercalate "\n" (map ppSequence pageItems)
6767
where
@@ -81,12 +81,28 @@ splitOn c s = case break (== c) s of
8181
(_:r) -> splitOn c r
8282

8383
parseArg :: String -> Maybe AbiValue
84-
parseArg s =
84+
parseArg s =
8585
let s' = trim s
8686
in if "0x" `isPrefixOf` s'
8787
then AbiAddress . fromIntegral <$> (readMaybe s' :: Maybe Integer)
8888
else AbiUInt 256 . fromIntegral <$> (readMaybe s' :: Maybe Integer)
8989

90+
parseFuzzArg :: String -> Maybe (Maybe AbiValue)
91+
parseFuzzArg s =
92+
let s' = trim s
93+
in if s' == "?"
94+
then Just Nothing
95+
else Just <$> parseArg s'
96+
97+
parseFuzzCall :: String -> Maybe (Text, [Maybe AbiValue])
98+
parseFuzzCall s = do
99+
let (fname, rest) = break (== '(') s
100+
if null rest then Nothing else do
101+
let argsS = take (length rest - 2) (drop 1 rest) -- remove parens
102+
let argParts = if all isSpace argsS then [] else splitOn ',' argsS
103+
args <- mapM parseFuzzArg argParts
104+
return (pack fname, args)
105+
90106
parseCall :: String -> Maybe (String, [AbiValue])
91107
parseCall s = do
92108
let (fname, rest) = break (== '(') s
@@ -126,10 +142,10 @@ injectTransactionTool args env bus _ = do
126142
Just p -> Data.Maybe.fromMaybe 0 (readMaybe (unpack p))
127143
Nothing -> 0
128144
txStr = maybe "" unpack (lookup "transaction" args)
129-
145+
130146
c <- readIORef env.corpusRef
131147
let corpusList = Set.toList c
132-
148+
133149
if idx < 0 || idx >= length corpusList
134150
then return "Error: Invalid sequence index."
135151
else do
@@ -139,16 +155,16 @@ injectTransactionTool args env bus _ = do
139155
else do
140156
let contextTx = case originalSeq of
141157
[] -> Nothing
142-
(x:xs) -> Just (if pos > 0 && pos <= length (x:xs)
143-
then (x:xs) !! (pos - 1)
158+
(x:xs) -> Just (if pos > 0 && pos <= length (x:xs)
159+
then (x:xs) !! (pos - 1)
144160
else x)
145161
case parseTx contextTx txStr of
146162
Nothing -> return "Error: Failed to parse transaction string."
147163
Just newTx -> do
148164
let newSeq = take pos originalSeq ++ [newTx]
149165
replyVar <- newEmptyTMVarIO
150166
atomically $ writeTChan bus (WrappedMessage AIId (ToFuzzer 0 (ExecuteSequence newSeq (Just replyVar))))
151-
167+
152168
-- Wait for reply
153169
found <- atomically $ takeTMVar replyVar
154170
if found
@@ -163,13 +179,16 @@ dumpLcovTool _ env _ _ = do
163179
filename <- saveLcovHook env dir env.sourceCache contracts
164180
return $ "Dumped LCOV coverage to " ++ filename
165181

166-
-- | Implementation of prioritize_function tool
167-
prioritizeFunctionTool :: ToolExecution
168-
prioritizeFunctionTool args env bus _ = do
169-
let msg = Data.Maybe.fromMaybe "" (lookup "function" args)
170-
let nWorkers = getNFuzzWorkers env.cfg.campaignConf
171-
mapM_ (\i -> atomically $ writeTChan bus (WrappedMessage AIId (ToFuzzer i (PrioritizeFunction (unpack msg))))) [0 .. nWorkers - 1]
172-
return $ printf "Requested prioritization of function '%s' on %d fuzzers" (unpack msg) nWorkers
182+
-- | Implementation of fuzz_transaction tool
183+
fuzzTransactionTool :: ToolExecution
184+
fuzzTransactionTool args env bus _ = do
185+
let txStr = Data.Maybe.fromMaybe "" (lookup "transaction" args)
186+
case parseFuzzCall (unpack txStr) of
187+
Nothing -> return "Error: Failed to parse transaction string."
188+
Just (fname, fuzzArgs) -> do
189+
let nWorkers = getNFuzzWorkers env.cfg.campaignConf
190+
mapM_ (\i -> atomically $ writeTChan bus (WrappedMessage AIId (ToFuzzer i (FuzzTransaction fname fuzzArgs)))) [0 .. nWorkers - 1]
191+
return $ printf "Requested fuzzing of transaction '%s' on %d fuzzers" (unpack txStr) nWorkers
173192

174193
-- | Implementation of clear_priorities tool
175194
clearPrioritiesTool :: ToolExecution
@@ -182,7 +201,11 @@ clearPrioritiesTool _ env bus _ = do
182201
readLogsTool :: ToolExecution
183202
readLogsTool _ _ _ logsRef = do
184203
logs <- readIORef logsRef
185-
return $ unpack $ T.unlines $ reverse logs
204+
-- Get last 100 logs
205+
-- logs is [Newest, ..., Oldest]
206+
-- We want to take the 100 newest, and show them in chronological order
207+
let logsToShow = reverse $ take 100 logs
208+
return $ unpack $ T.unlines $ logsToShow
186209

187210
-- | Implementation of show_coverage tool
188211
showCoverageTool :: ToolExecution
@@ -235,7 +258,7 @@ availableTools =
235258
, Tool "inspect_corpus_transactions" "Browse the corpus transactions" inspectCorpusTransactionsTool
236259
, Tool "inject_transaction" "Inject a transaction into a sequence and execute it" injectTransactionTool
237260
, Tool "dump_lcov" "Dump coverage in LCOV format" dumpLcovTool
238-
, Tool "prioritize_function" "Prioritize a function for fuzzing" prioritizeFunctionTool
261+
, Tool "fuzz_transaction" "Fuzz a transaction with optional concrete arguments" fuzzTransactionTool
239262
, Tool "clear_priorities" "Clear the function prioritization list" clearPrioritiesTool
240263
, Tool "read_logs" "Read the last 100 log messages" readLogsTool
241264
, Tool "show_coverage" "Show coverage report for a particular contract" showCoverageTool
@@ -254,7 +277,7 @@ runMCPServer env port logsRef = do
254277
let serverInfo = McpServerInfo
255278
{ serverName = "Echidna MCP Server"
256279
, serverVersion = "1.0.0"
257-
, serverInstructions = "Echidna Agent Interface. Available tools: read_corpus, inspect_corpus_transactions, dump_lcov, prioritize_function, clear_priorities, read_logs, show_coverage"
280+
, serverInstructions = "Echidna Agent Interface. Available tools: read_corpus, inspect_corpus_transactions, dump_lcov, fuzz_transaction, clear_priorities, read_logs, show_coverage"
258281
}
259282

260283
let mkToolDefinition :: Tool -> ToolDefinition
@@ -267,7 +290,7 @@ runMCPServer env port logsRef = do
267290
, required = ["page"]
268291
}
269292
"inject_transaction" -> InputSchemaDefinitionObject
270-
{ properties =
293+
{ properties =
271294
[ ("sequence_index", InputSchemaDefinitionProperty "string" "The index of the sequence in the corpus")
272295
, ("position", InputSchemaDefinitionProperty "string" "The position to insert the transaction at")
273296
, ("transaction", InputSchemaDefinitionProperty "string" "The transaction string (e.g. 'func(arg1, arg2)')")
@@ -278,9 +301,9 @@ runMCPServer env port logsRef = do
278301
{ properties = []
279302
, required = []
280303
}
281-
"prioritize_function" -> InputSchemaDefinitionObject
282-
{ properties = [("function", InputSchemaDefinitionProperty "string" "The name of the function to prioritize")]
283-
, required = ["function"]
304+
"fuzz_transaction" -> InputSchemaDefinitionObject
305+
{ properties = [("transaction", InputSchemaDefinitionProperty "string" "The transaction string (e.g. 'func(arg1, ?, arg3)')")]
306+
, required = ["transaction"]
284307
}
285308
"clear_priorities" -> InputSchemaDefinitionObject
286309
{ properties = []

‎lib/Echidna/Transaction.hs‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
module Echidna.Transaction where
55

6-
import Control.Monad (join, when)
6+
import Control.Monad (join, when, zipWithM)
77
import Control.Monad.IO.Class (MonadIO, liftIO)
88
import Control.Monad.Random.Strict (MonadRandom, getRandom, getRandomR, uniform)
99
import Control.Monad.Reader (MonadReader, ask)
@@ -14,7 +14,6 @@ import Data.Map (Map, toList)
1414
import Data.Maybe (catMaybes)
1515
import Data.Set (Set)
1616
import Data.Set qualified as Set
17-
import qualified Data.Text as T
1817
import Data.Vector qualified as V
1918
import Optics.Core
2019
import Optics.State.Operators
@@ -68,28 +67,46 @@ genTx world deployedContracts = do
6867
let txConf = env.cfg.txConf
6968
genDict <- gets (.genDict)
7069
prioritized <- gets (.prioritizedFunctions)
71-
let prioritizedTxt = map T.pack prioritized
7270
sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap
7371
sender <- rElem' world.senders
7472
contractAList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts)
7573
let allContracts = catMaybes contractAList
76-
(dstAddr, dstAbis) <- if null prioritizedTxt
77-
then rElem' $ Set.fromList allContracts
74+
75+
(dstAddr, solCall) <- if null prioritized
76+
then do
77+
(addr, sigs) <- rElem' $ Set.fromList allContracts
78+
call <- genInteractionsM genDict sigs
79+
pure (addr, call)
7880
else do
79-
let isPrioritized n = any (`T.isInfixOf` n) prioritizedTxt
80-
let prioritizedContracts = filter (\(_, sigs) -> any (\(n,_) -> isPrioritized n) sigs) allContracts
8181
usePrioritized <- (<= (0.9 :: Double)) <$> getRandom
82-
if usePrioritized && not (null prioritizedContracts)
82+
if usePrioritized
8383
then do
84-
(addr, sigs) <- rElem' $ Set.fromList prioritizedContracts
85-
-- Filter sigs to only prioritized ones
86-
let pSigs = NE.filter (\(n, _) -> isPrioritized n) sigs
87-
case NE.nonEmpty pSigs of
88-
Just pSigsNE -> pure (addr, pSigsNE)
89-
Nothing -> pure (addr, sigs) -- Should not happen
90-
else rElem' $ Set.fromList allContracts
84+
(pName, pArgs) <- rElem (NE.fromList prioritized)
85+
-- Find contracts containing this function with matching arity
86+
let isMatch (_, sigs) = any (\(n, ts) -> n == pName && length ts == length pArgs) sigs
87+
matchingContracts = filter isMatch allContracts
88+
89+
if null matchingContracts
90+
then do
91+
(addr, sigs) <- rElem' $ Set.fromList allContracts
92+
call <- genInteractionsM genDict sigs
93+
pure (addr, call)
94+
else do
95+
(addr, sigs) <- rElem' $ Set.fromList matchingContracts
96+
-- Pick the matching signature
97+
let matchingSigs = NE.filter (\(n, ts) -> n == pName && length ts == length pArgs) sigs
98+
(name, types) <- rElem (NE.fromList matchingSigs)
99+
100+
-- Generate arguments
101+
let genArg (Just val) _ = pure val
102+
genArg Nothing t = genAbiValueM' genDict name 0 t
91103

92-
solCall <- genInteractionsM genDict dstAbis
104+
vals <- zipWithM genArg pArgs types
105+
pure (addr, (name, vals))
106+
else do
107+
(addr, sigs) <- rElem' $ Set.fromList allContracts
108+
call <- genInteractionsM genDict sigs
109+
pure (addr, call)
93110
value <- genValue txConf.maxValue genDict.dictValues world.payableSigs solCall
94111
ts <- (,) <$> genDelay txConf.maxTimeDelay genDict.dictValues
95112
<*> genDelay txConf.maxBlockDelay genDict.dictValues

‎lib/Echidna/Types/Campaign.hs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import Data.Word (Word8, Word16)
66
import GHC.Conc (numCapabilities)
77

88
import EVM.Solvers (Solver(..))
9+
import EVM.ABI (AbiValue)
910

1011
import Echidna.ABI (GenDict, emptyDict)
1112
import Echidna.Types
@@ -86,7 +87,7 @@ data WorkerState = WorkerState
8687
, runningThreads :: [ThreadId]
8788
-- ^ Extra threads currently being run,
8889
-- aside from the main worker thread
89-
, prioritizedFunctions :: ![String]
90+
, prioritizedFunctions :: ![(Text, [Maybe AbiValue])]
9091
-- ^ Functions to prioritize during fuzzing
9192
}
9293

‎lib/Echidna/Types/InterWorker.hs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ module Echidna.Types.InterWorker where
33
import Control.Concurrent.STM
44
import Data.Text (Text)
55

6+
import EVM.ABI (AbiValue)
67
import Echidna.Types.Tx (Tx)
78
import Echidna.Types.Test (EchidnaTest)
89

@@ -14,14 +15,16 @@ data AgentId = FuzzerId Int | SymbolicId | AIId
1415
data FuzzerCmd
1516
= DumpLcov
1617
| SolutionFound [Tx]
17-
| PrioritizeFunction String
18+
| PrioritizeFunction String -- Deprecated
19+
| FuzzTransaction Text [Maybe AbiValue]
1820
| ClearPrioritization
1921
| ExecuteSequence [Tx] (Maybe (TMVar Bool))
2022

2123
instance Show FuzzerCmd where
2224
show DumpLcov = "DumpLcov"
2325
show (SolutionFound txs) = "SolutionFound " ++ show txs
2426
show (PrioritizeFunction s) = "PrioritizeFunction " ++ show s
27+
show (FuzzTransaction s args) = "FuzzTransaction " ++ show s ++ " " ++ show args
2528
show ClearPrioritization = "ClearPrioritization"
2629
show (ExecuteSequence txs _) = "ExecuteSequence " ++ show txs
2730

‎lib/Echidna/UI.hs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ ui vm dict initialCorpus cliSelectedContract = do
185185

186186
let forwardEvent ev = do
187187
msg <- runReaderT (ppLogLine vm ev) env
188-
liftIO $ atomicModifyIORef' logBuffer (\logs -> (take 100 (pack msg : logs), ()))
188+
liftIO $ atomicModifyIORef' logBuffer (\logs -> (pack msg : logs, ()))
189189
putStrLn msg
190190
uiEventsForwarderStopVar <- spawnListener forwardEvent
191191

0 commit comments

Comments
 (0)