Follow upstream changes -- Bytestring updates
[git-darcs-import.git] / src / Ssh.hs
blob0198ff7ca7bd6d445b6ecabefebcafd215ed7d21
1 {-# OPTIONS_GHC -cpp -fffi #-}
2 {-# LANGUAGE CPP, ForeignFunctionInterface #-}
4 module Ssh ( grabSSH, runSSH, getSSH, copySSH, copySSHs, SSHCmd(..) ) where
6 import Prelude hiding ( lookup, catch )
8 import System.Exit ( ExitCode(..) )
9 import System.Environment ( getEnv )
10 #ifndef WIN32
11 import System.Posix.Process ( getProcessID )
12 #else
13 import Darcs.Utils ( showHexLen )
14 import Data.Bits ( (.&.) )
15 import System.Random ( randomIO )
16 #endif
17 import System.IO ( Handle, hPutStr, hPutStrLn, hGetLine, hGetContents, hClose, hFlush )
18 import System.IO.Unsafe ( unsafePerformIO )
19 import System.Directory ( doesFileExist, createDirectoryIfMissing )
20 import Control.Monad ( when )
21 import System.Process ( runInteractiveProcess )
23 import Data.Map ( Map, empty, insert, lookup )
24 import Data.IORef ( IORef, newIORef, readIORef, modifyIORef )
26 import Darcs.SignalHandler ( catchNonSignal )
27 import Darcs.Utils ( withCurrentDirectory, breakCommand, prettyException, catchall )
28 import Darcs.Global ( atexit, sshControlMasterDisabled, darcsdir, withDebugMode )
29 import Darcs.Lock ( withTemp, withOpenTemp, tempdir_loc, removeFileMayNotExist )
30 import Exec ( exec, Redirects, Redirect(..), )
31 import Darcs.Progress ( withoutProgress, debugMessage, debugFail, progressList )
33 import qualified Data.ByteString as B (ByteString, hGet, writeFile, readFile)
34 import qualified Data.ByteString.Char8 as BC (unpack)
36 #include "impossible.h"
38 {-# NOINLINE sshConnections #-}
39 sshConnections :: IORef (Map String (Maybe Connection))
40 sshConnections = unsafePerformIO $ newIORef empty
42 data Connection = C { inp :: !Handle, out :: !Handle, err :: !Handle, deb :: String -> IO () }
44 withSSHConnection :: String -> (Connection -> IO a) -> IO a -> IO a
45 withSSHConnection x withconnection withoutconnection =
46 withoutProgress $
47 do cs <- readIORef sshConnections
48 let uhost = takeWhile (/= ':') x
49 url = cleanrepourl x
50 case lookup url (cs :: Map String (Maybe Connection)) of
51 Just Nothing -> withoutconnection
52 Just (Just c) -> withconnection c
53 Nothing ->
54 do mc <- do (ssh,sshargs_) <- getSSHOnly SSH
55 let sshargs = sshargs_ ++ [uhost,"darcs","transfer-mode","--repodir",cleanrepodir x]
56 debugMessage $ "ssh "++unwords sshargs
57 (i,o,e,_) <- runInteractiveProcess ssh sshargs Nothing Nothing
58 l <- hGetLine o
59 if l == "Hello user, I am darcs transfer mode"
60 then return ()
61 else debugFail "Couldn't start darcs transfer-mode on server"
62 let c = C { inp = i, out = o, err = e,
63 deb = \s -> debugMessage ("with ssh (transfer-mode) "++uhost++": "++s) }
64 modifyIORef sshConnections (insert url (Just c))
65 return $ Just c
66 `catchNonSignal`
67 \e -> do debugMessage $ "Failed to start ssh connection:\n "++
68 prettyException e
69 severSSHConnection x
70 debugMessage $ unlines $
71 [ "NOTE: the server may be running a version of darcs prior to 2.0.0."
72 , ""
73 , "Installing darcs 2 on the server will speed up ssh-based commands."
75 return Nothing
76 maybe withoutconnection withconnection mc
78 severSSHConnection :: String -> IO ()
79 severSSHConnection x = do debugMessage $ "Severing ssh failed connection to "++x
80 modifyIORef sshConnections (insert (cleanrepourl x) Nothing)
82 cleanrepourl :: String -> String
83 cleanrepourl zzz | take (length dd) zzz == dd = ""
84 where dd = darcsdir++"/"
85 cleanrepourl (z:zs) = z : cleanrepourl zs
86 cleanrepourl "" = ""
88 cleanrepodir :: String -> String
89 cleanrepodir = cleanrepourl . drop 1 . dropWhile (/= ':')
91 grabSSH :: String -> Connection -> IO B.ByteString
92 grabSSH x c = do
93 let dir = drop 1 $ dropWhile (/= ':') x
94 dd = darcsdir++"/"
95 clean zzz | take (length dd) zzz == dd = drop (length dd) zzz
96 clean (_:zs) = clean zs
97 clean "" = bug $ "Buggy path in grabSSH: "++x
98 file = clean dir
99 failwith e = do severSSHConnection x
100 eee <- hGetContents (err c) -- ratify hGetContents: it's okay
101 -- here because we're only grabbing
102 -- stderr, and we're also about to
103 -- throw the contents.
104 debugFail $ e ++ " grabbing ssh file "++x++"\n"++eee
105 deb c $ "get "++file
106 hPutStrLn (inp c) $ "get " ++ file
107 hFlush (inp c)
108 l2 <- hGetLine (out c)
109 if l2 == "got "++file
110 then do showlen <- hGetLine (out c)
111 case reads showlen of
112 [(len,"")] -> B.hGet (out c) len
113 _ -> failwith "Couldn't get length"
114 else if l2 == "error "++file
115 then do e <- hGetLine (out c)
116 case reads e of
117 (msg,_):_ -> debugFail $ "Error reading file remotely:\n"++msg
118 [] -> failwith "An error occurred"
119 else failwith "Error"
121 sshStdErrMode :: IO Redirect
122 sshStdErrMode = withDebugMode $ \amdebugging ->
123 return $ if amdebugging then AsIs else Null
125 copySSH :: String -> FilePath -> IO ()
126 copySSH uRaw f = withSSHConnection uRaw (\c -> grabSSH uRaw c >>= B.writeFile f) $
127 do let u = escape_dollar uRaw
128 stderr_behavior <- sshStdErrMode
129 r <- runSSH SCP u [] [u,f] (AsIs,AsIs,stderr_behavior)
130 when (r /= ExitSuccess) $
131 debugFail $ "(scp) failed to fetch: " ++ u
132 where {- '$' in filenames is troublesome for scp, for some reason.. -}
133 escape_dollar :: String -> String
134 escape_dollar = concatMap tr
135 where tr '$' = "\\$"
136 tr c = [c]
138 copySSHs :: String -> [String] -> FilePath -> IO ()
139 copySSHs u ns d =
140 withSSHConnection u (\c -> withCurrentDirectory d $
141 mapM_ (\n -> grabSSH (u++"/"++n) c >>= B.writeFile n) $
142 progressList "Copying via ssh" ns) $
143 do let path = drop 1 $ dropWhile (/= ':') u
144 host = takeWhile (/= ':') u
145 cd = "cd "++path++"\n"
146 input = cd++(unlines $ map ("get "++) ns)
147 withCurrentDirectory d $ withOpenTemp $ \(th,tn) ->
148 withTemp $ \sftpoutput ->
149 do hPutStr th input
150 hClose th
151 stderr_behavior <- sshStdErrMode
152 r <- runSSH SFTP u [] [host] (File tn, File sftpoutput, stderr_behavior)
153 let files = if length ns > 5
154 then (take 5 ns) ++ ["and "
155 ++ (show (length ns - 5)) ++ " more"]
156 else ns
157 hint = if take 1 path == "~"
158 then ["sftp doesn't expand ~, use path/ instead of ~/path/"]
159 else []
160 when (r /= ExitSuccess) $ do
161 outputPS <- B.readFile sftpoutput
162 debugFail $ unlines $
163 ["(sftp) failed to fetch files.",
164 "source directory: " ++ path,
165 "source files:"] ++ files ++
166 ["sftp output:",BC.unpack outputPS] ++
167 hint
169 -- ---------------------------------------------------------------------
170 -- older ssh helper functions
171 -- ---------------------------------------------------------------------
173 data SSHCmd = SSH | SCP | SFTP
175 instance Show SSHCmd where
176 show SSH = "ssh"
177 show SCP = "scp"
178 show SFTP = "sftp"
180 runSSH :: SSHCmd -> String -> [String] -> [String] -> Redirects -> IO ExitCode
181 runSSH cmd remoteAddr preArgs postArgs redirs =
182 do (ssh, args) <- getSSH cmd remoteAddr
183 exec ssh (preArgs ++ args ++ postArgs) redirs
185 -- | Return the command and arguments needed to run an ssh command
186 -- along with any extra features like use of the control master.
187 -- See 'getSSHOnly'
188 getSSH :: SSHCmd -> String -- ^ remote path
189 -> IO (String, [String])
190 getSSH cmd remoteAddr =
191 do (ssh, ssh_args) <- getSSHOnly cmd
192 cm_args <- if sshControlMasterDisabled
193 then return []
194 else do -- control master
195 cmPath <- controlMasterPath remoteAddr
196 hasLaunchedCm <- doesFileExist cmPath
197 when (not hasLaunchedCm) $ launchSSHControlMaster remoteAddr
198 hasCmFeature <- doesFileExist cmPath
199 return $ if hasCmFeature then [ "-o ControlPath=" ++ cmPath ] else []
200 let verbosity = case cmd of
201 SCP -> ["-q"] -- (p)scp is the only one that recognises -q
202 -- sftp and (p)sftp do not, and plink neither
203 _ -> []
205 return (ssh, verbosity ++ ssh_args ++ cm_args)
207 -- | Return the command and arguments needed to run an ssh command.
208 -- First try the appropriate darcs environment variable and SSH_PORT
209 -- defaulting to "ssh" and no specified port.
210 getSSHOnly :: SSHCmd -> IO (String, [String])
211 getSSHOnly cmd =
212 do ssh_command <- getEnv (evar cmd) `catchall` return (show cmd)
213 -- port
214 port <- (portFlag cmd `fmap` getEnv "SSH_PORT") `catchall` return []
215 let (ssh, ssh_args) = breakCommand ssh_command
217 return (ssh, ssh_args ++ port)
218 where
219 evar SSH = "DARCS_SSH"
220 evar SCP = "DARCS_SCP"
221 evar SFTP = "DARCS_SFTP"
222 portFlag SSH x = ["-p", x]
223 portFlag SCP x = ["-P", x]
224 portFlag SFTP x = ["-oPort="++x]
226 -- | Return True if this version of ssh has a ControlMaster feature
227 -- The ControlMaster functionality allows for ssh multiplexing
228 hasSSHControlMaster :: Bool
229 hasSSHControlMaster = unsafePerformIO hasSSHControlMasterIO
231 -- Because of the unsafePerformIO above, this can be called at any
232 -- point. It cannot rely on any state, not even the current directory.
233 hasSSHControlMasterIO :: IO Bool
234 hasSSHControlMasterIO = do
235 (ssh, _) <- getSSHOnly SSH
236 -- If ssh has the ControlMaster feature, it will recognise the
237 -- the -O flag, but exit with status 255 because of the nonsense
238 -- command. If it does not have the feature, it will simply dump
239 -- a help message on the screen and exit with 1.
240 sx <- exec ssh ["-O", "an_invalid_command"] (Null,Null,Null)
241 case sx of
242 ExitFailure 255 -> return True
243 _ -> return False
245 -- | Launch an SSH control master in the background, if available.
246 -- We don't have to wait for it or anything.
247 -- Note also that this will cleanup after itself when darcs exits
248 launchSSHControlMaster :: String -> IO ()
249 launchSSHControlMaster rawAddr =
250 when hasSSHControlMaster $ do
251 let addr = takeWhile (/= ':') rawAddr
252 (ssh, ssh_args) <- getSSHOnly SSH
253 cmPath <- controlMasterPath addr
254 removeFileMayNotExist cmPath
255 -- -f : put ssh in the background once it succeeds in logging you in
256 -- -M : launch as the control master for addr
257 -- -N : don't run any commands
258 -- -S : use cmPath as the ControlPath. Equivalent to -oControlPath=
259 exec ssh (ssh_args ++ [addr, "-S", cmPath, "-N", "-f", "-M"]) (Null,Null,AsIs)
260 atexit $ exitSSHControlMaster addr
261 return ()
263 -- | Tell the SSH control master for a given path to exit.
264 exitSSHControlMaster :: String -> IO ()
265 exitSSHControlMaster addr = do
266 (ssh, ssh_args) <- getSSHOnly SSH
267 cmPath <- controlMasterPath addr
268 exec ssh (ssh_args ++ [addr, "-S", cmPath, "-O", "exit"]) (Null,Null,Null)
269 return ()
271 -- | Create the directory ssh control master path for a given address
272 controlMasterPath :: String -- ^ remote path (foo\@bar.com:file is ok; the file part with be stripped)
273 -> IO FilePath
274 controlMasterPath rawAddr = do
275 let addr = takeWhile (/= ':') rawAddr
276 tmp <- (fmap (/// ".darcs") $ getEnv "HOME") `catchall` tempdir_loc
277 #ifdef WIN32
278 r <- randomIO
279 let suffix = (showHexLen 6 (r .&. 0xFFFFFF :: Int))
280 #else
281 suffix <- show `fmap` getProcessID
282 #endif
283 let tmpDarcsSsh = tmp /// "darcs-ssh"
284 createDirectoryIfMissing True tmpDarcsSsh
285 return $ tmpDarcsSsh /// addr ++ suffix
287 (///) :: FilePath -> FilePath -> FilePath
288 d /// f = d ++ "/" ++ f