Shell pipesHEADmaster

Changelog: Implemented pipes in shell scripts

3 files changed, 130 insertions, 28 deletions

diff --git a/src/Parser/Shell.hs b/src/Parser/Shell.hs
index 7720ee9..22d47ed 100644
--- a/src/Parser/Shell.hs
+++ b/src/Parser/Shell.hs
@@ -74,17 +74,35 @@ parseArgument = lexeme $ fmap (App AnnNone (Pure T.concat) <$> foldr (liftA2 (:)
 parseArguments :: TestParser (Expr [ Text ])
 parseArguments = foldr (liftA2 (:)) (Pure []) <$> many parseArgument
 
-shellStatement :: TestParser (Expr [ ShellStatement ])
-shellStatement = label "shell statement" $ do
+parseCommand :: TestParser (Expr ShellCommand)
+parseCommand = label "shell statement" $ do
     line <- getSourceLine
     command <- parseArgument
     args <- parseArguments
-    return $ fmap (: []) $ ShellStatement
+    return $ ShellCommand
         <$> command
         <*> args
         <*> pure line
 
+parsePipeline :: Expr (Maybe ShellPipeline) -> TestParser (Expr ShellPipeline)
+parsePipeline upper = do
+    cmd <- parseCommand
+    let pipeline = ShellPipeline <$> cmd <*> upper
+    choice
+        [ do
+            osymbol "|"
+            parsePipeline (Just <$> pipeline)
+
+        , do
+            return pipeline
+        ]
+
+parseStatement :: TestParser (Expr [ ShellStatement ])
+parseStatement = do
+    line <- getSourceLine
+    fmap ((: []) . flip ShellStatement line) <$> parsePipeline (pure Nothing)
+
 shellScript :: TestParser (Expr ShellScript)
 shellScript = do
     indent <- L.indentLevel
-    fmap ShellScript <$> blockOf indent shellStatement
+    fmap ShellScript <$> blockOf indent parseStatement
diff --git a/src/Script/Shell.hs b/src/Script/Shell.hs
index 95d4fe4..1c052b0 100644
--- a/src/Script/Shell.hs
+++ b/src/Script/Shell.hs
@@ -1,6 +1,8 @@
 module Script.Shell (
-    ShellStatement(..),
     ShellScript(..),
+    ShellStatement(..),
+    ShellPipeline(..),
+    ShellCommand(..),
     withShellProcess,
 ) where
 
@@ -15,8 +17,15 @@ import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Text.IO qualified as T
 
+import Foreign.C.Types
+import Foreign.Ptr
+import Foreign.Marshal.Array
+import Foreign.Storable
+
 import System.Exit
 import System.IO
+import System.Posix.IO qualified as P
+import System.Posix.Types
 import System.Process hiding (ShellCommand)
 
 import Network
@@ -27,48 +36,101 @@ import Run.Monad
 import Script.Var
 
 
+newtype ShellScript = ShellScript [ ShellStatement ]
+
 data ShellStatement = ShellStatement
-    { shellCommand :: Text
-    , shellArguments :: [ Text ]
+    { shellPipeline :: ShellPipeline
     , shellSourceLine :: SourceLine
     }
 
-newtype ShellScript = ShellScript [ ShellStatement ]
+data ShellPipeline = ShellPipeline
+    { pipeCommand :: ShellCommand
+    , pipeUpstream :: Maybe ShellPipeline
+    }
+
+data ShellCommand = ShellCommand
+    { cmdCommand :: Text
+    , cmdArguments :: [ Text ]
+    , cmdSourceLine :: SourceLine
+    }
+
+
+data ShellExecInfo = ShellExecInfo
+    { seiNode :: Node
+    , seiProcName :: ProcName
+    , seiStatusVar :: MVar ExitCode
+    }
 
 
-executeScript :: Node -> ProcName -> MVar ExitCode -> Handle -> Handle -> Handle -> ShellScript -> TestRun ()
-executeScript node pname statusVar pstdin pstdout pstderr (ShellScript statements) = do
-    setNetworkNamespace $ getNetns node
-    forM_ statements $ \ShellStatement {..} -> case shellCommand of
+data HandleHandling
+    = CloseHandle Handle
+    | KeepHandle Handle
+
+closeIfRequested :: MonadIO m => HandleHandling -> m ()
+closeIfRequested (CloseHandle h) = liftIO $ hClose h
+closeIfRequested (KeepHandle _) = return ()
+
+handledHandle :: HandleHandling -> Handle
+handledHandle (CloseHandle h) = h
+handledHandle (KeepHandle h) = h
+
+
+executeCommand :: ShellExecInfo -> HandleHandling -> HandleHandling -> HandleHandling -> ShellCommand -> TestRun ()
+executeCommand ShellExecInfo {..} pstdin pstdout pstderr ShellCommand {..} = do
+    case cmdCommand of
         "echo" -> liftIO $ do
-            T.hPutStrLn pstdout $ T.intercalate " " shellArguments
-            hFlush pstdout
+            T.hPutStrLn (handledHandle pstdout) $ T.intercalate " " cmdArguments
+            hFlush (handledHandle pstdout)
+            mapM_ closeIfRequested [ pstdin, pstdout, pstderr ]
         cmd -> do
             (_, _, _, phandle) <- liftIO $ createProcess_ "shell"
-                (proc (T.unpack cmd) (map T.unpack shellArguments))
-                    { std_in = UseHandle pstdin
-                    , std_out = UseHandle pstdout
-                    , std_err = UseHandle pstderr
-                    , cwd = Just (nodeDir node)
+                (proc (T.unpack cmd) (map T.unpack cmdArguments))
+                    { std_in = UseHandle $ handledHandle pstdin
+                    , std_out = UseHandle $ handledHandle pstdout
+                    , std_err = UseHandle $ handledHandle pstderr
+                    , cwd = Just (nodeDir seiNode)
                     , env = Just []
                     }
+            mapM_ closeIfRequested [ pstdin, pstdout, pstderr ]
             liftIO (waitForProcess phandle) >>= \case
                 ExitSuccess -> return ()
                 status -> do
-                    outLine OutputChildFail (Just $ textProcName pname) $ "failed at: " <> textSourceLine shellSourceLine
-                    liftIO $ putMVar statusVar status
+                    outLine OutputChildFail (Just $ textProcName seiProcName) $ "failed at: " <> textSourceLine cmdSourceLine
+                    liftIO $ putMVar seiStatusVar status
                     throwError Failed
-    liftIO $ putMVar statusVar ExitSuccess
+
+executePipeline :: ShellExecInfo -> HandleHandling -> HandleHandling -> HandleHandling -> ShellPipeline -> TestRun ()
+executePipeline sei pstdin pstdout pstderr ShellPipeline {..} = do
+    case pipeUpstream of
+        Nothing -> do
+            executeCommand sei pstdin pstdout pstderr pipeCommand
+
+        Just upstream -> do
+            ( pipeRead, pipeWrite ) <- createPipeCloexec
+            void $ forkTestUsing forkOS $ do
+                executePipeline sei pstdin (CloseHandle pipeWrite) (KeepHandle $ handledHandle pstderr) upstream
+
+            executeCommand sei (CloseHandle pipeRead) pstdout (KeepHandle $ handledHandle pstderr) pipeCommand
+            closeIfRequested pstderr
+
+executeScript :: ShellExecInfo -> Handle -> Handle -> Handle -> ShellScript -> TestRun ()
+executeScript sei@ShellExecInfo {..} pstdin pstdout pstderr (ShellScript statements) = do
+    setNetworkNamespace $ getNetns seiNode
+    forM_ statements $ \ShellStatement {..} -> do
+        executePipeline sei (KeepHandle pstdin) (KeepHandle pstdout) (KeepHandle pstderr) shellPipeline
+    liftIO $ putMVar seiStatusVar ExitSuccess
 
 spawnShell :: Node -> ProcName -> ShellScript -> TestRun Process
 spawnShell procNode procName script = do
     procOutput <- liftIO $ newTVarIO []
-    statusVar <- liftIO $ newEmptyMVar
-    ( pstdin, procStdin ) <- liftIO $ createPipe
-    ( hout, pstdout ) <- liftIO $ createPipe
-    ( herr, pstderr ) <- liftIO $ createPipe
-    procHandle <- fmap (Right . (, statusVar)) $ forkTestUsing forkOS $ do
-        executeScript procNode procName statusVar pstdin pstdout pstderr script
+    seiStatusVar <- liftIO $ newEmptyMVar
+    ( pstdin, procStdin ) <- createPipeCloexec
+    ( hout, pstdout ) <- createPipeCloexec
+    ( herr, pstderr ) <- createPipeCloexec
+    procHandle <- fmap (Right . (, seiStatusVar)) $ forkTestUsing forkOS $ do
+        let seiNode = procNode
+            seiProcName = procName
+        executeScript ShellExecInfo {..} pstdin pstdout pstderr script
         liftIO $ do
             hClose pstdin
             hClose pstdout
@@ -96,3 +158,17 @@ withShellProcess node pname script inner = do
         ps <- liftIO $ takeMVar procVar
         closeTestProcess process `finally` do
             liftIO $ putMVar procVar $ filter (/=process) ps
+
+
+foreign import ccall "shell_pipe_cloexec" c_pipe_cloexec :: Ptr Fd -> IO CInt
+
+createPipeCloexec :: (MonadIO m, MonadFail m) => m ( Handle, Handle )
+createPipeCloexec = liftIO $ do
+    allocaArray 2 $ \ptr -> do
+        c_pipe_cloexec ptr >>= \case
+            0 -> do
+                rh <- P.fdToHandle =<< peekElemOff ptr 0
+                wh <- P.fdToHandle =<< peekElemOff ptr 1
+                return ( rh, wh )
+            _ -> do
+                fail $ "failed to create pipe"
diff --git a/src/shell.c b/src/shell.c
new file mode 100644
index 0000000..d832078
--- /dev/null
+++ b/src/shell.c
@@ -0,0 +1,8 @@
+#define _GNU_SOURCE
+#include <fcntl.h>
+#include <unistd.h>
+
+int shell_pipe_cloexec( int pipefd[ 2 ] )
+{
+	return pipe2( pipefd, O_CLOEXEC );
+}