1 files changed, 83 insertions, 30 deletions

diff --git a/src/Process.hs b/src/Process.hs
index 61a9fe8..1389987 100644
--- a/src/Process.hs
+++ b/src/Process.hs
@@ -5,9 +5,14 @@ module Process (
     send,
     outProc,
     lineReadingLoop,
+    startProcessIOLoops,
     spawnOn,
     closeProcess,
+    closeTestProcess,
     withProcess,
+
+    IgnoreProcessOutput(..),
+    flushProcessOutput,
 ) where
 
 import Control.Arrow
@@ -18,9 +23,11 @@ import Control.Monad.Except
 import Control.Monad.Reader
 
 import Data.Function
+import Data.Maybe
+import Data.Scientific
 import Data.Text (Text)
-import qualified Data.Text as T
-import qualified Data.Text.IO as T
+import Data.Text qualified as T
+import Data.Text.IO qualified as T
 
 import System.Directory
 import System.Environment
@@ -36,13 +43,16 @@ import Network
 import Network.Ip
 import Output
 import Run.Monad
+import Script.Expr
 import Script.Expr.Class
+import Script.Object
 
 data Process = Process
     { procName :: ProcName
     , procHandle :: Either ProcessHandle ( ThreadId, MVar ExitCode )
     , procStdin :: Handle
-    , procOutput :: TVar [Text]
+    , procOutput :: TVar [ Text ]
+    , procIgnore :: TVar ( Int, [ ( Int, Maybe Regex ) ] )
     , procKillWith :: Maybe Signal
     , procNode :: Node
     }
@@ -83,15 +93,38 @@ outProc otype p line = outLine otype (Just $ textProcName $ procName p) line
 lineReadingLoop :: MonadOutput m => Process -> Handle -> (Text -> m ()) -> m ()
 lineReadingLoop process h act =
     liftIO (tryIOError (T.hGetLine h)) >>= \case
-        Left err
-            | isEOFError err -> return ()
-            | otherwise -> outProc OutputChildFail process $ T.pack $ "IO error: " ++ show err
+        Left err -> do
+            when (not (isEOFError err)) $ do
+                outProc OutputChildFail process $ T.pack $ "IO error: " ++ show err
+            liftIO $ hClose h
         Right line -> do
             act line
             lineReadingLoop process h act
 
+startProcessIOLoops :: Process -> Handle -> Handle -> TestRun ()
+startProcessIOLoops process@Process {..} hout herr = do
+
+    void $ forkTest $ lineReadingLoop process hout $ \line -> do
+        outProc OutputChildStdout process line
+        liftIO $ atomically $ do
+            ignores <- map snd . snd <$> readTVar procIgnore
+            when (not $ any (matches line) ignores) $ do
+                modifyTVar procOutput (++ [ line ])
+
+    void $ forkTest $ lineReadingLoop process herr $ \line -> do
+        case procName of
+            ProcNameTcpdump -> return ()
+            _ -> outProc OutputChildStderr process line
+
+  where
+    matches _    Nothing
+        = True
+    matches line (Just re)
+        | Right (Just _) <- regexMatch re line = True
+        | otherwise = False
+
 spawnOn :: Either Network Node -> ProcName -> Maybe Signal -> String -> TestRun Process
-spawnOn target pname killWith cmd = do
+spawnOn target procName procKillWith cmd = do
     -- When executing command given with relative path, turn it to absolute one,
     -- because working directory will be changed for the shell wrapper.
     cmd' <- liftIO $ do
@@ -105,39 +138,28 @@ spawnOn target pname killWith cmd = do
 
     let netns = either getNetns getNetns target
     currentEnv <- liftIO $ getEnvironment
-    (Just hin, Just hout, Just herr, handle) <- liftIO $ do
+    (Just procStdin, Just hout, Just herr, handle) <- liftIO $ do
         runInNetworkNamespace netns $ createProcess (shell cmd')
             { std_in = CreatePipe, std_out = CreatePipe, std_err = CreatePipe
             , cwd = Just (either netDir nodeDir target)
             , env = Just $ ( "EREBOS_DIR", "." ) : currentEnv
             }
-    pout <- liftIO $ newTVarIO []
-
-    let process = Process
-            { procName = pname
-            , procHandle = Left handle
-            , procStdin = hin
-            , procOutput = pout
-            , procKillWith = killWith
-            , procNode = either (const undefined) id target
-            }
+    let procHandle = Left handle
+    procOutput <- liftIO $ newTVarIO []
+    procIgnore <- liftIO $ newTVarIO ( 0, [] )
+    let procNode = either (const undefined) id target
+    let process = Process {..}
 
-    void $ forkTest $ lineReadingLoop process hout $ \line -> do
-        outProc OutputChildStdout process line
-        liftIO $ atomically $ modifyTVar pout (++[line])
-    void $ forkTest $ lineReadingLoop process herr $ \line -> do
-        case pname of
-             ProcNameTcpdump -> return ()
-             _ -> outProc OutputChildStderr process line
+    startProcessIOLoops process hout herr
 
     asks (teGDB . fst) >>= maybe (return Nothing) (liftIO . tryReadMVar) >>= \case
-        Just gdb | ProcName _ <- pname -> addInferior gdb process
+        Just gdb | ProcName _ <- procName -> addInferior gdb process
         _ -> return ()
 
     return process
 
-closeProcess :: (MonadIO m, MonadOutput m, MonadError Failed m) => Process -> m ()
-closeProcess p = do
+closeProcess :: (MonadIO m, MonadOutput m, MonadError Failed m) => Scientific -> Process -> m ()
+closeProcess timeout p = do
     liftIO $ hClose $ procStdin p
     case procKillWith p of
         Nothing -> return ()
@@ -146,7 +168,7 @@ closeProcess p = do
             Just pid -> signalProcess sig pid
 
     liftIO $ void $ forkIO $ do
-        threadDelay 1000000
+        threadDelay $ floor $ 1000000 * timeout
         either terminateProcess (killThread . fst) $ procHandle p
     liftIO (either waitForProcess (takeMVar . snd) (procHandle p)) >>= \case
         ExitSuccess -> return ()
@@ -154,6 +176,11 @@ closeProcess p = do
             outProc OutputChildFail p $ T.pack $ "exit code: " ++ show code
             throwError Failed
 
+closeTestProcess :: Process -> TestRun ()
+closeTestProcess process = do
+    timeout <- getCurrentTimeout
+    closeProcess timeout process
+
 withProcess :: Either Network Node -> ProcName -> Maybe Signal -> String -> (Process -> TestRun a) -> TestRun a
 withProcess target pname killWith cmd inner = do
     procVar <- asks $ teProcesses . fst
@@ -163,5 +190,31 @@ withProcess target pname killWith cmd inner = do
 
     inner process `finally` do
         ps <- liftIO $ takeMVar procVar
-        closeProcess process `finally` do
+        closeTestProcess process `finally` do
             liftIO $ putMVar procVar $ filter (/=process) ps
+
+
+data IgnoreProcessOutput = IgnoreProcessOutput Process Int
+
+instance ObjectType TestRun IgnoreProcessOutput where
+    type ConstructorArgs IgnoreProcessOutput = ( Process, Maybe Regex )
+
+    textObjectType _ _ = "IgnoreProcessOutput"
+    textObjectValue _ (IgnoreProcessOutput _ _) = "<IgnoreProcessOutput>"
+
+    createObject oid ( process@Process {..}, regex ) = do
+        liftIO $ atomically $ do
+            flushProcessOutput process regex
+            ( iid, list ) <- readTVar procIgnore
+            writeTVar procIgnore ( iid + 1, ( iid, regex ) : list )
+            return $ Object oid $ IgnoreProcessOutput process iid
+
+    destroyObject Object { objImpl = IgnoreProcessOutput Process {..} iid } = do
+        liftIO $ atomically $ do
+            writeTVar procIgnore . fmap (filter ((iid /=) . fst)) =<< readTVar procIgnore
+
+flushProcessOutput :: Process -> Maybe Regex -> STM ()
+flushProcessOutput p mbre = do
+    writeTVar (procOutput p) =<< case mbre of
+        Nothing -> return []
+        Just re -> filter (either error isNothing . regexMatch re) <$> readTVar (procOutput p)