summaryrefslogtreecommitdiff
path: root/src/Script/Shell.hs
blob: 1c052b08fed86c130c6cf6a1fc9254fc37e70584 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
module Script.Shell (
    ShellScript(..),
    ShellStatement(..),
    ShellPipeline(..),
    ShellCommand(..),
    withShellProcess,
) where

import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import Control.Monad.Except
import Control.Monad.IO.Class
import Control.Monad.Reader

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
import Network.Ip
import Output
import Process
import Run.Monad
import Script.Var


newtype ShellScript = ShellScript [ ShellStatement ]

data ShellStatement = ShellStatement
    { shellPipeline :: ShellPipeline
    , shellSourceLine :: SourceLine
    }

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
    }


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 (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 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 seiProcName) $ "failed at: " <> textSourceLine cmdSourceLine
                    liftIO $ putMVar seiStatusVar status
                    throwError Failed

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 []
    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
            hClose pstderr

    let procKillWith = Nothing
    let process = Process {..}

    void $ forkTest $ lineReadingLoop process hout $ \line -> do
        outProc OutputChildStdout process line
        liftIO $ atomically $ modifyTVar procOutput (++ [ line ])
    void $ forkTest $ lineReadingLoop process herr $ \line -> do
        outProc OutputChildStderr process line

    return process

withShellProcess :: Node -> ProcName -> ShellScript -> (Process -> TestRun a) -> TestRun a
withShellProcess node pname script inner = do
    procVar <- asks $ teProcesses . fst

    process <- spawnShell node pname script
    liftIO $ modifyMVar_ procVar $ return . (process:)

    inner process `finally` 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"