git.iko.soy/distributed-process
on 17ad36a help sign in
distributed-process/distributed-process-azure/src/Control/Distributed/Process/Backend/Azure.hs
335 lines 11 kB history raw
1module Control.Distributed.Process.Backend.Azure
2 ( -- * Initialization
3 Backend(..)
4 , AzureParameters(..)
5 , defaultAzureParameters
6 , initializeBackend
7 -- * Re-exports from Azure Service Management
8 , CloudService(..)
9 , VirtualMachine(..)
10 , Azure.cloudServices
11 -- * Remote and local processes
12 , ProcessPair(..)
13 , RemoteProcess
14 , LocalProcess
15 , localExpect
16 , remoteSend
17 , remoteThrow
18 , remoteSend'
19 , localSend
20 ) where
22import System.Environment (getEnv)
23import System.FilePath ((</>), takeFileName)
24import System.Environment.Executable (getExecutablePath)
25import System.IO (stdout, hFlush)
26import Data.Binary (encode, decode)
27import Data.Digest.Pure.MD5 (md5, MD5Digest)
28import qualified Data.ByteString as BSS
29 ( ByteString
30 , length
31 , concat
32 , hPut
33 )
34import qualified Data.ByteString.Char8 as BSSC (pack)
35import qualified Data.ByteString.Lazy as BSL
36 ( ByteString
37 , readFile
38 , length
39 , fromChunks
40 , toChunks
41 , hPut
42 )
43import qualified Data.ByteString.Lazy.Char8 as BSLC (unpack)
44import Data.Typeable (Typeable)
45import Control.Applicative ((<$>))
46import Control.Monad (void, unless)
47import Control.Monad.Reader (MonadReader, ReaderT, runReaderT, ask)
48import Control.Exception (Exception, catches, Handler(Handler))
49import Control.Monad.IO.Class (MonadIO, liftIO)
51-- Azure
52import Network.Azure.ServiceManagement
53 ( CloudService(..)
54 , VirtualMachine(..)
55 , Endpoint(..)
56 , AzureSetup
57 )
58import qualified Network.Azure.ServiceManagement as Azure
59 ( cloudServices
60 , azureSetup
61 , vmSshEndpoint
62 )
64-- SSH
65import qualified Network.SSH.Client.LibSSH2 as SSH
66 ( withSSH2
67 , scpSendFile
68 , withChannelBy
69 , Session
70 , readAllChannel
71 , writeAllChannel
72 , Channel
73 )
74import qualified Network.SSH.Client.LibSSH2.Foreign as SSH
75 ( openChannelSession
76 , channelExecute
77 , writeChannel
78 , readChannel
79 , channelSendEOF
80 )
81import qualified Network.SSH.Client.LibSSH2.Errors as SSH
82 ( ErrorCode
83 , NULL_POINTER
84 , getLastError
85 )
87-- CH
88import Control.Distributed.Process (Process, Closure)
89import Control.Distributed.Process.Serializable (Serializable)
90import Control.Distributed.Process.Internal.Types
91 ( messageToPayload
92 , createMessage
93 )
94import Network.Transport.Internal (encodeInt32, decodeInt32, prependLength)
96-- | Azure backend
97data Backend = Backend {
98 -- | Find virtual machines
99 findVMs :: IO [VirtualMachine]
100 -- | Copy the executable to a virtual machine
101 , copyToVM :: VirtualMachine -> IO ()
102 -- | Check the MD5 hash of the remote executable
103 , checkMD5 :: VirtualMachine -> IO Bool
104 -- | @runOnVM dict vm port p@ starts a CH node on port 'port' and runs 'p'
105 , callOnVM :: forall a.
106 VirtualMachine
107 -> String
108 -> ProcessPair a
109 -> IO a
110 -- | Create a new CH node and run the specified process in the background.
111 -- The CH node will exit when the process exists.
112 , spawnOnVM :: VirtualMachine
113 -> String
114 -> Closure (Backend -> Process ())
115 -> IO ()
116 } deriving (Typeable)
118data AzureParameters = AzureParameters {
119 azureSetup :: AzureSetup
120 , azureSshUserName :: FilePath
121 , azureSshPublicKey :: FilePath
122 , azureSshPrivateKey :: FilePath
123 , azureSshPassphrase :: String
124 , azureSshKnownHosts :: FilePath
125 , azureSshRemotePath :: FilePath
126 , azureSshLocalPath :: FilePath
127 }
129-- | Create default azure parameters
130defaultAzureParameters :: String -- ^ Azure subscription ID
131 -> FilePath -- ^ Path to X509 certificate
132 -> FilePath -- ^ Path to private key
133 -> IO AzureParameters
134defaultAzureParameters sid x509 pkey = do
135 home <- getEnv "HOME"
136 user <- getEnv "USER"
137 self <- getExecutablePath
138 setup <- Azure.azureSetup sid x509 pkey
139 return AzureParameters
140 { azureSetup = setup
141 , azureSshUserName = user
142 , azureSshPublicKey = home </> ".ssh" </> "id_rsa.pub"
143 , azureSshPrivateKey = home </> ".ssh" </> "id_rsa"
144 , azureSshPassphrase = ""
145 , azureSshKnownHosts = home </> ".ssh" </> "known_hosts"
146 , azureSshRemotePath = takeFileName self
147 , azureSshLocalPath = self
148 }
150-- | Initialize the backend
151initializeBackend :: AzureParameters -- ^ Connection parameters
152 -> String -- ^ Cloud service name
153 -> IO Backend
154initializeBackend params cloudService =
155 return Backend {
156 findVMs = apiFindVMs params cloudService
157 , copyToVM = apiCopyToVM params
158 , checkMD5 = apiCheckMD5 params
159 , callOnVM = apiCallOnVM params
160 , spawnOnVM = apiSpawnOnVM params
161 }
163-- | Find virtual machines
164apiFindVMs :: AzureParameters -> String -> IO [VirtualMachine]
165apiFindVMs params cloudService = do
166 css <- Azure.cloudServices (azureSetup params)
167 case filter ((== cloudService) . cloudServiceName) css of
168 [cs] -> return $ cloudServiceVMs cs
169 _ -> return []
171-- | Start a CH node on the given virtual machine
172apiCopyToVM :: AzureParameters -> VirtualMachine -> IO ()
173apiCopyToVM params vm =
174 void . withSSH2 params vm $ \s -> catchSshError s $
175 SSH.scpSendFile s 0o700 (azureSshLocalPath params) (azureSshRemotePath params)
177-- | Call a process on a VM
178apiCallOnVM :: AzureParameters
179 -> VirtualMachine
180 -> String
181 -> ProcessPair a
182 -> IO a
183apiCallOnVM params vm port ppair =
184 withSSH2 params vm $ \s -> do
185 let exe = "PATH=. " ++ azureSshRemotePath params
186 ++ " onvm run "
187 ++ " --host " ++ vmIpAddress vm
188 ++ " --port " ++ port
189 ++ " 2>&1"
190 (status, r) <- SSH.withChannelBy (SSH.openChannelSession s) id $ \ch -> do
191 SSH.channelExecute ch exe
192 _ <- SSH.writeChannel ch (encodeInt32 (BSL.length rprocEnc))
193 _ <- SSH.writeAllChannel ch rprocEnc
194 runLocalProcess (ppairLocal ppair) ch
195 if status == 0
196 then return r
197 else error "callOnVM: Non-zero exit status"
198 where
199 rprocEnc :: BSL.ByteString
200 rprocEnc = encode (ppairRemote ppair)
202apiSpawnOnVM :: AzureParameters
203 -> VirtualMachine
204 -> String
205 -> Closure (Backend -> Process ())
206 -> IO ()
207apiSpawnOnVM params vm port proc =
208 withSSH2 params vm $ \s -> do
209 let exe = "PATH=. " ++ azureSshRemotePath params
210 ++ " onvm run "
211 ++ " --host " ++ vmIpAddress vm
212 ++ " --port " ++ port
213 ++ " --background "
214 ++ " 2>&1"
215 (status, r) <- SSH.withChannelBy (SSH.openChannelSession s) id $ \ch -> do
216 SSH.channelExecute ch exe
217 _ <- SSH.writeChannel ch (encodeInt32 (BSL.length procEnc))
218 _ <- SSH.writeAllChannel ch procEnc
219 SSH.channelSendEOF ch
220 SSH.readAllChannel ch
221 unless (status == 0) $ error (BSLC.unpack r)
222 where
223 procEnc :: BSL.ByteString
224 procEnc = encode proc
226-- | Check the MD5 hash of the executable on the remote machine
227apiCheckMD5 :: AzureParameters -> VirtualMachine -> IO Bool
228apiCheckMD5 params vm = do
229 hash <- localHash params
230 withSSH2 params vm $ \s -> do
231 (r, _) <- SSH.withChannelBy (SSH.openChannelSession s) id $ \ch -> do
232 SSH.channelExecute ch "md5sum -c --status"
233 SSH.writeChannel ch . BSSC.pack $ show hash ++ " " ++ azureSshRemotePath params
234 SSH.channelSendEOF ch
235 SSH.readAllChannel ch
236 return (r == 0)
238withSSH2 :: AzureParameters -> VirtualMachine -> (SSH.Session -> IO a) -> IO a
239withSSH2 params (Azure.vmSshEndpoint -> Just ep) =
240 SSH.withSSH2 (azureSshKnownHosts params)
241 (azureSshPublicKey params)
242 (azureSshPrivateKey params)
243 (azureSshPassphrase params)
244 (azureSshUserName params)
245 (endpointVip ep)
246 (read $ endpointPort ep)
247withSSH2 _ vm =
248 error $ "withSSH2: No SSH endpoint for virtual machine " ++ vmName vm
250catchSshError :: SSH.Session -> IO a -> IO a
251catchSshError s io =
252 catches io [ Handler handleErrorCode
253 , Handler handleNullPointer
254 ]
255 where
256 handleErrorCode :: SSH.ErrorCode -> IO a
257 handleErrorCode _ = do
258 (_, str) <- SSH.getLastError s
259 error str
261 handleNullPointer :: SSH.NULL_POINTER -> IO a
262 handleNullPointer _ = do
263 (_, str) <- SSH.getLastError s
264 error str
266localHash :: AzureParameters -> IO MD5Digest
267localHash params = md5 <$> BSL.readFile (azureSshLocalPath params)
269--------------------------------------------------------------------------------
270-- Local and remote processes --
271--------------------------------------------------------------------------------
273data ProcessPair a = ProcessPair {
274 ppairRemote :: RemoteProcess ()
275 , ppairLocal :: LocalProcess a
276 }
278type RemoteProcess a = Closure (Backend -> Process a)
280newtype LocalProcess a = LocalProcess { unLocalProcess :: ReaderT SSH.Channel IO a }
281 deriving (Functor, Monad, MonadIO, MonadReader SSH.Channel)
283runLocalProcess :: LocalProcess a -> SSH.Channel -> IO a
284runLocalProcess = runReaderT . unLocalProcess
286localExpect :: Serializable a => LocalProcess a
287localExpect = LocalProcess $ do
288 ch <- ask
289 liftIO $ do
290 isE <- readIntChannel ch
291 len <- readIntChannel ch
292 msg <- readSizeChannel ch len
293 if isE /= 0
294 then error (decode msg)
295 else return (decode msg)
297localSend :: Serializable a => a -> LocalProcess ()
298localSend x = LocalProcess $ do
299 ch <- ask
300 liftIO $ mapM_ (SSH.writeChannel ch)
301 . prependLength
302 . messageToPayload
303 . createMessage
304 $ x
306remoteSend :: Serializable a => a -> Process ()
307remoteSend = liftIO . remoteSend' 0
309remoteThrow :: Exception e => e -> Process ()
310remoteThrow = liftIO . remoteSend' 1 . show
312remoteSend' :: Serializable a => Int -> a -> IO ()
313remoteSend' flags x = do
314 let enc = encode x
315 BSS.hPut stdout (encodeInt32 flags)
316 BSS.hPut stdout (encodeInt32 (BSL.length enc))
317 BSL.hPut stdout enc
318 hFlush stdout
320--------------------------------------------------------------------------------
321-- SSH utilities --
322--------------------------------------------------------------------------------
324readSizeChannel :: SSH.Channel -> Int -> IO BSL.ByteString
325readSizeChannel ch = go []
326 where
327 go :: [BSS.ByteString] -> Int -> IO BSL.ByteString
328 go acc 0 = return (BSL.fromChunks $ reverse acc)
329 go acc size = do
330 bs <- SSH.readChannel ch (fromIntegral (0x400 `min` size))
331 go (bs : acc) (size - BSS.length bs)
333readIntChannel :: SSH.Channel -> IO Int
334readIntChannel ch =
335 decodeInt32 . BSS.concat . BSL.toChunks <$> readSizeChannel ch 4