From c7475433b5842fb3fff8a1478be0f80de3b5f358 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Wed, 10 Jul 2024 16:07:54 -0400
Subject: [PATCH] catch mvar errors in UI.hs

---
 lib/Echidna/CatchMVar.hs | 14 +++++++++++++-
 lib/Echidna/UI.hs        | 19 +++++++++++--------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/lib/Echidna/CatchMVar.hs b/lib/Echidna/CatchMVar.hs
index b842ed63..793da578 100644
--- a/lib/Echidna/CatchMVar.hs
+++ b/lib/Echidna/CatchMVar.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE TemplateHaskell #-}
 
-module Echidna.CatchMVar (catchMVar, putMVar_, takeMVar_, writeChan_, readChan_) where
+module Echidna.CatchMVar where
 
 import Control.Exception
 import Language.Haskell.TH
@@ -28,11 +28,23 @@ catchMVarTempl = [| catchMVar $printLocation |]
 putMVar_ :: Q Exp
 putMVar_ = [| ($catchMVarTempl .) . putMVar |]
 
+tryPutMVar_ :: Q Exp
+tryPutMVar_ = [| ($catchMVarTempl .) . tryPutMVar |]
+
 takeMVar_ :: Q Exp
 takeMVar_ = [| $catchMVarTempl . takeMVar |]
 
+readMVar_ :: Q Exp
+readMVar_ = [| $catchMVarTempl . readMVar |]
+
 writeChan_ :: Q Exp
 writeChan_ = [| ($catchMVarTempl .) . writeChan |]
 
 readChan_ :: Q Exp
 readChan_ = [| $catchMVarTempl . readChan |]
+
+writeBChan_ :: Q Exp
+writeBChan_ = [| ($catchMVarTempl .) . writeBChan |]
+
+writeBChanNonBlocking_ :: Q Exp
+writeBChanNonBlocking_ = [| ($catchMVarTempl .) . writeBChanNonBlocking |]
diff --git a/lib/Echidna/UI.hs b/lib/Echidna/UI.hs
index 5387aed5..ff0b77e9 100644
--- a/lib/Echidna/UI.hs
+++ b/lib/Echidna/UI.hs
@@ -1,8 +1,11 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE TemplateHaskell #-}
 {-# OPTIONS_GHC -Wno-unused-imports #-}
 
 module Echidna.UI where
 
+import Echidna.CatchMVar
+
 #ifdef INTERACTIVE_UI
 import Brick
 import Brick.BChan
@@ -97,7 +100,7 @@ ui vm dict initialCorpus cliSelectedContract = do
     Interactive -> do
       -- Channel to push events to update UI
       uiChannel <- liftIO $ newBChan 1000
-      let forwardEvent = void . writeBChanNonBlocking uiChannel . EventReceived
+      let forwardEvent = void . $writeBChanNonBlocking_ uiChannel . EventReceived
       uiEventsForwarderStopVar <- spawnListener forwardEvent
 
       ticker <- liftIO . forkIO . forever $ do
@@ -106,12 +109,12 @@ ui vm dict initialCorpus cliSelectedContract = do
         now <- getTimestamp
         tests <- traverse readIORef env.testRefs
         states <- workerStates workers
-        writeBChan uiChannel (CampaignUpdated now tests states)
+        $writeBChan_ uiChannel (CampaignUpdated now tests states)
 
         -- TODO: remove and use events for this
         c <- readIORef env.fetchContractCache
         s <- readIORef env.fetchSlotCache
-        writeBChan uiChannel (FetchCacheUpdated c s)
+        $writeBChan_ uiChannel (FetchCacheUpdated c s)
 
       -- UI initialization
       let buildVty = do
@@ -149,7 +152,7 @@ ui vm dict initialCorpus cliSelectedContract = do
       stopWorkers workers
 
       -- wait for all events to be processed
-      forM_ [uiEventsForwarderStopVar, corpusSaverStopVar] takeMVar
+      liftIO $ forM_ [uiEventsForwarderStopVar, corpusSaverStopVar] $takeMVar_
 
       liftIO $ killThread ticker
 
@@ -168,7 +171,7 @@ ui vm dict initialCorpus cliSelectedContract = do
       liftIO $ forM_ [sigINT, sigTERM] $ \sig ->
         let handler = Catch $ do
               stopWorkers workers
-              void $ tryPutMVar serverStopVar ()
+              void $ $tryPutMVar_ serverStopVar ()
         in installHandler sig handler Nothing
 #endif
       let forwardEvent ev = putStrLn =<< runReaderT (ppLogLine vm ev) env
@@ -190,7 +193,7 @@ ui vm dict initialCorpus cliSelectedContract = do
         printStatus
 
       -- wait for all events to be processed
-      forM_ [uiEventsForwarderStopVar, corpusSaverStopVar] takeMVar
+      liftIO $ forM_ [uiEventsForwarderStopVar, corpusSaverStopVar] $takeMVar_
 
       liftIO $ killThread ticker
 
@@ -200,7 +203,7 @@ ui vm dict initialCorpus cliSelectedContract = do
       when (isJust conf.campaignConf.serverPort) $ do
         -- wait until we send all SSE events
         liftIO $ putStrLn "Waiting until all SSE are received..."
-        readMVar serverStopVar
+        liftIO $ $readMVar_ serverStopVar
 
       states <- liftIO $ workerStates workers
 
@@ -237,7 +240,7 @@ ui vm dict initialCorpus cliSelectedContract = do
         ]
 
       time <- liftIO getTimestamp
-      writeChan env.eventQueue (time, WorkerEvent workerId workerType (WorkerStopped stopReason))
+      liftIO $ $writeChan_ env.eventQueue (time, WorkerEvent workerId workerType (WorkerStopped stopReason))
 
     pure (threadId, stateRef)