summaryrefslogtreecommitdiff
path: root/Propellor/SimpleSh.hs
blob: 0999be9a6a8124e204d729847db63d49cf92409a (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
-- | Simple server, using a named pipe. Client connects, sends a command,
-- and gets back all the output from the command, in a stream.
--
-- This is useful for eg, docker.

module Propellor.SimpleSh where

import Network.Socket
import Control.Concurrent.Chan
import Control.Concurrent.Async
import System.Process (std_in, std_out, std_err)
import System.Exit

import Propellor
import Utility.FileMode
import Utility.ThreadScheduler

data Cmd = Cmd String [String]
	deriving (Read, Show)

data Resp = StdoutLine String | StderrLine String | Done ExitCode
	deriving (Read, Show)

simpleSh :: FilePath -> IO ()
simpleSh namedpipe = do
	nukeFile namedpipe
	let dir = takeDirectory namedpipe
	createDirectoryIfMissing True dir
	modifyFileMode dir (removeModes otherGroupModes)
	s <- socket AF_UNIX Stream defaultProtocol
	bind s (SockAddrUnix namedpipe)
	listen s 2
	forever $ do
		(client, _addr) <- accept s
		h <- socketToHandle client ReadWriteMode
		hSetBuffering h LineBuffering
		maybe noop (run h) . readish =<< hGetLine h
  where
	run h (Cmd cmd params) = do
		let p = (proc cmd params)
                	{ std_in = Inherit
	                , std_out = CreatePipe
			, std_err = CreatePipe
			}
		(Nothing, Just outh, Just errh, pid) <- createProcess p
		chan <- newChan

		let runwriter = do
			v <- readChan chan
			hPutStrLn h (show v)
			case v of
				Done _ -> noop
				_ -> runwriter
		writer <- async runwriter

		let mkreader t from = maybe noop (const $ mkreader t from) 
			=<< catchMaybeIO (writeChan chan . t =<< hGetLine from)
		void $ concurrently
			(mkreader StdoutLine outh)
			(mkreader StderrLine errh)

		writeChan chan . Done =<< waitForProcess pid

		wait writer

		hClose outh
		hClose errh
		hClose h

simpleShClient :: FilePath -> String -> [String] -> ([Resp] -> IO a) -> IO a
simpleShClient namedpipe cmd params handler = do
	s <- socket AF_UNIX Stream defaultProtocol
	connect s (SockAddrUnix namedpipe)
	h <- socketToHandle s ReadWriteMode
	hSetBuffering h LineBuffering
	hPutStrLn h $ show $ Cmd cmd params
	resps <- catMaybes . map readish . lines <$> hGetContents h
	hClose h `after` handler resps

simpleShClientRetry :: Int -> FilePath -> String -> [String] -> ([Resp] -> IO a) -> IO a
simpleShClientRetry retries namedpipe cmd params handler = go retries
  where
	run = simpleShClient namedpipe cmd params handler
	go n
		| n < 1 = run
		| otherwise = do
			v <- tryIO run
			case v of
				Right r -> return r
				Left _ -> do
					threadDelaySeconds (Seconds 1)
					go (n - 1)

getStdout :: Resp -> Maybe String
getStdout (StdoutLine s) = Just s
getStdout _ = Nothing