diff --git a/lib/Echidna/Campaign.hs b/lib/Echidna/Campaign.hs index 029ab8815..2845586bc 100644 --- a/lib/Echidna/Campaign.hs +++ b/lib/Echidna/Campaign.hs @@ -29,6 +29,8 @@ import Data.Time (LocalTime) import Data.Vector qualified as V import System.Random (mkStdGen) +import Echidna.MCP (runMCPServer) + import EVM (cheatCode) import EVM.ABI (getAbi, AbiType(AbiAddressType, AbiTupleType), AbiValue(AbiAddress, AbiTuple), abiValueType) import EVM.Dapp (DappInfo(..)) @@ -57,7 +59,7 @@ import Echidna.Types.Test import Echidna.Types.Test qualified as Test import Echidna.Types.Tx (TxCall(..), Tx(..)) import Echidna.Types.Worker -import Echidna.Worker +import Echidna.Worker instance MonadThrow m => MonadThrow (RandT g m) where throwM = lift . throwM @@ -156,7 +158,7 @@ runSymWorker callback vm dict workerId _ name = do shrinkAndRandomlyExplore transactions (10 :: Int) listenerFunc _ = pure () - shrinkAndRandomlyExplore _ 0 = do + shrinkAndRandomlyExplore _ 0 = do testRefs <- asks (.testRefs) tests <- liftIO $ traverse readIORef testRefs CampaignConf{shrinkLimit} <- asks (.cfg.campaignConf) @@ -166,7 +168,7 @@ runSymWorker callback vm dict workerId _ name = do testRefs <- asks (.testRefs) tests <- liftIO $ traverse readIORef testRefs CampaignConf{stopOnFail, shrinkLimit} <- asks (.cfg.campaignConf) - if stopOnFail && any final tests then + if stopOnFail && any final tests then lift callback -- >> pure FastFailed else if any shrinkable tests then do shrinkLoop shrinkLimit @@ -191,7 +193,7 @@ runSymWorker callback vm dict workerId _ name = do shrinkLoop 0 = return () - shrinkLoop n = do + shrinkLoop n = do lift callback updateTests $ \test -> do if test.workerId == Just workerId then @@ -235,7 +237,7 @@ runSymWorker callback vm dict workerId _ name = do contract <- chooseContract cs name failedTests <- findFailedTests let failedTestSignatures = map getAssertionSignature failedTests - case tx of + case tx of Nothing -> getRandomTargetMethod contract conf.campaignConf.symExecTargets failedTestSignatures >>= \case Nothing -> do return () @@ -245,7 +247,7 @@ runSymWorker callback vm dict workerId _ name = do return () Just method -> do exploreAndVerify contract method vm' txsBase - + exploreAndVerify contract method vm' txsBase = do (threadId, symTxsChan) <- exploreContract contract method vm' modify' (\ws -> ws { runningThreads = [threadId] }) @@ -686,6 +688,12 @@ spawnListener handler = do liftIO $ void $ forkFinally (listenerLoop handler chan nworkers) (const $ putMVar stopVar ()) pure stopVar +spawnMCPServer :: (MonadReader Env m, MonadIO m) => m () +spawnMCPServer = do + env <- ask + forM_ env.cfg.campaignConf.mcpPort $ \port -> + liftIO $ void $ forkIO (runMCPServer env (fromIntegral port)) + -- | Repeatedly run 'handler' on events from 'chan'. -- Stops once 'workersAlive' workers stop. listenerLoop diff --git a/lib/Echidna/Config.hs b/lib/Echidna/Config.hs index f03b6e531..a26012cc4 100644 --- a/lib/Echidna/Config.hs +++ b/lib/Echidna/Config.hs @@ -99,6 +99,7 @@ instance FromJSON EConfigWithUsage where <*> v ..:? "coverageFormats" ..!= [Txt,Html,Lcov] <*> v ..:? "workers" <*> v ..:? "server" + <*> v ..:? "mcp" <*> v ..:? "symExec" ..!= False <*> smtSolver <*> v ..:? "symExecTargets" ..!= Nothing diff --git a/lib/Echidna/MCP.hs b/lib/Echidna/MCP.hs new file mode 100644 index 000000000..e0ba95fe8 --- /dev/null +++ b/lib/Echidna/MCP.hs @@ -0,0 +1,78 @@ +module Echidna.MCP (runMCPServer) where + +import Control.Monad.IO.Class (liftIO) +import Data.Aeson (encode) +import Data.IORef (readIORef) +import Data.List (find, isInfixOf) +import Data.Maybe (fromMaybe) +import Data.Set (toList) +import Data.Text (pack, unpack) +import Data.Text.Lazy (toStrict) +import Data.Text.Lazy.Encoding (decodeUtf8) +import MCP.Server +import Network.Wai.Handler.Warp (Port) + +import Echidna.Types.Config (Env(..)) +import Echidna.Types.Corpus (corpusSize) +import Echidna.Types.Tx (Tx) +import Echidna.Pretty (ppTx) + +getCorpusSize :: Env -> ToolCallHandler IO +getCorpusSize env _ _ = do + corpus <- liftIO $ readIORef env.corpusRef + pure $ Right $ ContentText $ pack $ show $ corpusSize corpus + +inspectCorpusTransactions :: Env -> ToolCallHandler IO +inspectCorpusTransactions env _ args = do + corpus <- liftIO $ readIORef env.corpusRef + let + sequence' = read $ unpack $ fromMaybe "0" $ lookup "sequence" args + page = read $ unpack $ fromMaybe "0" $ lookup "page" args + txs = fromMaybe [] $ snd <$> find (\(i, _) -> i == sequence') (toList corpus) + paginatedTxs = take 10 $ drop (page * 10) txs + pure $ Right $ ContentText $ pack $ unlines $ map (ppTx mempty) paginatedTxs + +findTransactionInCorpus :: Env -> ToolCallHandler IO +findTransactionInCorpus env _ args = do + corpus <- liftIO $ readIORef env.corpusRef + let + query = unpack $ fromMaybe "" $ lookup "query" args + results = + [ (seq', i `div` 10) + | (seq', txs) <- toList corpus + , (tx, i) <- zip txs [0..] + , query `isInfixOf` ppTx mempty tx + ] + pure $ Right $ ContentText $ pack $ show results + +runMCPServer :: Env -> Port -> IO () +runMCPServer env port = do + let + info = McpServerInfo "Echidna" "2.2.7" "Echidna MCP server" + tools' = + [ ToolDefinition "getCorpusSize" "Get the current corpus size" (InputSchemaDefinitionObject [] []) Nothing + , ToolDefinition "inspectCorpusTransactions" "Inspect corpus transactions" + (InputSchemaDefinitionObject + [ ("sequence", InputSchemaDefinitionProperty "integer" "Sequence number") + , ("page", InputSchemaDefinitionProperty "integer" "Page number (10 txs per page)") + ] + ["sequence", "page"]) + Nothing + , ToolDefinition "findTransactionInCorpus" "Find transaction in corpus" + (InputSchemaDefinitionObject + [ ("query", InputSchemaDefinitionProperty "string" "String to search") + ] + ["query"]) + Nothing + ] + toolHandler :: ToolCallHandler IO + toolHandler "getCorpusSize" args = getCorpusSize env "getCorpusSize" args + toolHandler "inspectCorpusTransactions" args = inspectCorpusTransactions env "inspectCorpusTransactions" args + toolHandler "findTransactionInCorpus" args = findTransactionInCorpus env "findTransactionInCorpus" args + toolHandler _ _ = pure $ Left $ InternalError "Unknown tool" + handlers = McpServerHandlers + { prompts = Nothing + , resources = Nothing + , tools = Just (pure tools', toolHandler) + } + runMcpServerHttpWithConfig (HttpConfig {httpPort = port, httpHost = "127.0.0.1", httpEndpoint = "/mcp", httpVerbose = False}) info handlers \ No newline at end of file diff --git a/lib/Echidna/Pretty.hs b/lib/Echidna/Pretty.hs index c9415e5a0..c1620f7af 100644 --- a/lib/Echidna/Pretty.hs +++ b/lib/Echidna/Pretty.hs @@ -10,7 +10,7 @@ import EVM.Types (Addr) import Echidna.ABI (ppAbiValue) import Echidna.Types.Signature (SolCall) -import Echidna.Types.Tx (TxCall(..)) +import Echidna.Types.Tx (Tx(..), TxCall(..)) -- | Pretty-print some 'AbiCall'. ppSolCall :: Map Addr Text -> SolCall -> String @@ -24,3 +24,7 @@ ppTxCall _ (SolCreate _) = "" ppTxCall labels (SolCall x) = ppSolCall labels x ppTxCall _ NoCall = "*wait*" ppTxCall _ (SolCalldata x) = BSC8.unpack $ "0x" <> BS16.encode x + +-- | Pretty-print some 'Tx' +ppTx :: Map Addr Text -> Tx -> String +ppTx labels tx = ppTxCall labels (tx.call) diff --git a/lib/Echidna/Types/Campaign.hs b/lib/Echidna/Types/Campaign.hs index fc4ae0314..33573cfed 100644 --- a/lib/Echidna/Types/Campaign.hs +++ b/lib/Echidna/Types/Campaign.hs @@ -39,6 +39,8 @@ data CampaignConf = CampaignConf -- ^ Number of fuzzing workers , serverPort :: Maybe Word16 -- ^ Server-Sent Events HTTP port number, if missing server is not ran + , mcpPort :: Maybe Word16 + -- ^ MCP server port number, if missing server is not ran , symExec :: Bool -- ^ Whether to add an additional symbolic execution worker , symExecSMTSolver :: Solver diff --git a/lib/Echidna/UI.hs b/lib/Echidna/UI.hs index 326e7d78e..b57bd3340 100644 --- a/lib/Echidna/UI.hs +++ b/lib/Echidna/UI.hs @@ -32,7 +32,7 @@ import UnliftIO.Concurrent hiding (killThread, threadDelay) import EVM.Types (Addr, Contract, VM, VMType(Concrete), W256) import Echidna.ABI -import Echidna.Campaign (runWorker, spawnListener) +import Echidna.Campaign (runWorker, spawnListener, spawnMCPServer) import Echidna.Output.Corpus (saveCorpusEvent) import Echidna.Output.JSON qualified import Echidna.Server (runSSEServer) @@ -43,7 +43,7 @@ import Echidna.Types.Corpus qualified as Corpus import Echidna.Types.Coverage (coverageStats) import Echidna.Types.Test (EchidnaTest(..), didFail, isOptimizationTest) import Echidna.Types.Tx (Tx) -import Echidna.Types.Worker +import Echidna.Types.Worker import Echidna.UI.Report import Echidna.UI.Widgets import Echidna.Utility (timePrefix, getTimestamp) @@ -94,6 +94,8 @@ ui vm dict initialCorpus cliSelectedContract = do corpusSaverStopVar <- spawnListener (saveCorpusEvent env) + spawnMCPServer + workers <- forM (zip corpusChunks [0..(nworkers-1)]) $ uncurry (spawnWorker env perWorkerTestLimit) diff --git a/package.yaml b/package.yaml index 99bf5e138..534b96035 100644 --- a/package.yaml +++ b/package.yaml @@ -79,6 +79,7 @@ library: - word-wrap - xml-conduit - yaml + - mcp-server executables: echidna: diff --git a/src/Main.hs b/src/Main.hs index 5839e8db7..562fe4f6a 100644 --- a/src/Main.hs +++ b/src/Main.hs @@ -122,6 +122,7 @@ data Options = Options { cliFilePath :: NE.NonEmpty FilePath , cliWorkers :: Maybe Word8 , cliServerPort :: Maybe Word16 + , cliMcpPort :: Maybe Word16 , cliSelectedContract :: Maybe Text , cliConfigFilepath :: Maybe FilePath , cliOutputFormat :: Maybe OutputFormat @@ -168,6 +169,9 @@ options = Options . NE.fromList <*> optional (option auto $ long "server" <> metavar "PORT" <> help "Run events server on the given port") + <*> optional (option auto $ long "mcp" + <> metavar "PORT" + <> help "Run MCP server on the given port") <*> optional (option str $ long "contract" <> metavar "CONTRACT" <> help "Contract to analyze") @@ -275,6 +279,7 @@ overrideConfig config Options{..} = do , seed = cliSeed <|> campaignConf.seed , workers = cliWorkers <|> campaignConf.workers , serverPort = cliServerPort <|> campaignConf.serverPort + , mcpPort = cliMcpPort <|> campaignConf.mcpPort , symExec = fromMaybe campaignConf.symExec cliSymExec , symExecTargets = (\ t -> Just [t]) =<< cliSymExecTargets , symExecTimeout = fromMaybe campaignConf.symExecTimeout cliSymExecTimeout diff --git a/stack.yaml b/stack.yaml index f6b4936d7..4de3cdd4c 100644 --- a/stack.yaml +++ b/stack.yaml @@ -3,7 +3,7 @@ resolver: lts-23.21 packages: - '.' - +- './haskell-mcp-server' extra-deps: - git: https://github.com/ethereum/hevm.git commit: 9982c580ed19b88ebab9744d29d940fd2f0bd8c6