15dc63d47f
It's possible, in direct or thin mode, that an object file gets truncated or appended to as it's being sent. This would break the protocol badly, so make sure never to send too many bytes, and to close the protocol connection if too few bytes are available.
258 lines
7 KiB
Haskell
258 lines
7 KiB
Haskell
{- P2P protocol, IO implementation
|
|
-
|
|
- Copyright 2016 Joey Hess <id@joeyh.name>
|
|
-
|
|
- Licensed under the GNU GPL version 3 or higher.
|
|
-}
|
|
|
|
{-# LANGUAGE RankNTypes, FlexibleContexts, BangPatterns, CPP #-}
|
|
|
|
module P2P.IO
|
|
( RunProto
|
|
, RunEnv(..)
|
|
, runNetProto
|
|
, runNet
|
|
) where
|
|
|
|
import P2P.Protocol
|
|
import Utility.Process
|
|
import Git
|
|
import Git.Command
|
|
import Utility.AuthToken
|
|
import Utility.SafeCommand
|
|
import Utility.SimpleProtocol
|
|
import Utility.Exception
|
|
|
|
import Control.Monad
|
|
import Control.Monad.Free
|
|
import Control.Monad.IO.Class
|
|
import System.Exit (ExitCode(..))
|
|
import System.IO
|
|
import Control.Concurrent
|
|
import Control.Concurrent.Async
|
|
import qualified Data.ByteString as B
|
|
import qualified Data.ByteString.Lazy as L
|
|
|
|
-- Type of interpreters of the Proto free monad.
|
|
type RunProto m = forall a. (MonadIO m, MonadMask m) => Proto a -> m (Maybe a)
|
|
|
|
data RunEnv = RunEnv
|
|
{ runRepo :: Repo
|
|
, runCheckAuth :: (AuthToken -> Bool)
|
|
, runIhdl :: Handle
|
|
, runOhdl :: Handle
|
|
}
|
|
|
|
-- Purposefully incomplete interpreter of Proto.
|
|
--
|
|
-- This only runs Net actions. No Local actions will be run
|
|
-- (those need the Annex monad) -- if the interpreter reaches any,
|
|
-- it returns Nothing.
|
|
runNetProto :: RunEnv -> Proto a -> IO (Maybe a)
|
|
runNetProto runenv = go
|
|
where
|
|
go :: RunProto IO
|
|
go (Pure v) = pure (Just v)
|
|
go (Free (Net n)) = runNet runenv go n
|
|
go (Free (Local _)) = return Nothing
|
|
|
|
-- Interpreter of the Net part of Proto.
|
|
--
|
|
-- An interpreter of Proto has to be provided, to handle the rest of Proto
|
|
-- actions.
|
|
runNet :: (MonadIO m, MonadMask m) => RunEnv -> RunProto m -> NetF (Proto a) -> m (Maybe a)
|
|
runNet runenv runner f = case f of
|
|
SendMessage m next -> do
|
|
v <- liftIO $ tryIO $ do
|
|
hPutStrLn (runOhdl runenv) (unwords (formatMessage m))
|
|
hFlush (runOhdl runenv)
|
|
case v of
|
|
Left _e -> return Nothing
|
|
Right () -> runner next
|
|
ReceiveMessage next -> do
|
|
v <- liftIO $ tryIO $ hGetLine (runIhdl runenv)
|
|
case v of
|
|
Left _e -> return Nothing
|
|
Right l -> case parseMessage l of
|
|
Just m -> runner (next m)
|
|
Nothing -> runner $ do
|
|
let e = ERROR $ "protocol parse error: " ++ show l
|
|
net $ sendMessage e
|
|
next e
|
|
SendBytes len b next -> do
|
|
v <- liftIO $ tryIO $ do
|
|
ok <- sendExactly len b (runOhdl runenv)
|
|
hFlush (runOhdl runenv)
|
|
return ok
|
|
case v of
|
|
Right True -> runner next
|
|
_ -> return Nothing
|
|
ReceiveBytes (Len n) next -> do
|
|
v <- liftIO $ tryIO $ L.hGet (runIhdl runenv) (fromIntegral n)
|
|
case v of
|
|
Left _e -> return Nothing
|
|
Right b -> runner (next b)
|
|
CheckAuthToken _u t next -> do
|
|
let authed = runCheckAuth runenv t
|
|
runner (next authed)
|
|
Relay hin hout next -> do
|
|
v <- liftIO $ runRelay runnerio hin hout
|
|
case v of
|
|
Nothing -> return Nothing
|
|
Just exitcode -> runner (next exitcode)
|
|
RelayService service next -> do
|
|
v <- liftIO $ runRelayService runenv runnerio service
|
|
case v of
|
|
Nothing -> return Nothing
|
|
Just () -> runner next
|
|
where
|
|
-- This is only used for running Net actions when relaying,
|
|
-- so it's ok to use runNetProto, despite it not supporting
|
|
-- all Proto actions.
|
|
runnerio = runNetProto runenv
|
|
|
|
-- Send exactly the specified number of bytes or returns False.
|
|
--
|
|
-- The ByteString can be larger or smaller than the specified length.
|
|
-- For example, it can be lazily streaming from a file that gets
|
|
-- appended to, or truncated.
|
|
--
|
|
-- Must avoid sending too many bytes as it would confuse the other end.
|
|
-- This is easily dealt with by truncating it.
|
|
--
|
|
-- If too few bytes are sent, the only option is to give up on this
|
|
-- connection. False is returned to indicate this problem.
|
|
--
|
|
-- We can't check the length of the whole lazy bytestring without buffering
|
|
-- it in memory. Instead, process it one chunk at a time, and sum the length
|
|
-- of the chunks.
|
|
sendExactly :: Len -> L.ByteString -> Handle -> IO Bool
|
|
sendExactly (Len l) lb h = go 0 $ L.toChunks $ L.take (fromIntegral l) lb
|
|
where
|
|
go n [] = return (toInteger n == l)
|
|
go n (b:bs) = do
|
|
B.hPut h b
|
|
let !n' = n + B.length b
|
|
go n' bs
|
|
|
|
runRelay :: RunProto IO -> RelayHandle -> RelayHandle -> IO (Maybe ExitCode)
|
|
runRelay runner (RelayHandle hout) (RelayHandle hin) = bracket setup cleanup go
|
|
where
|
|
setup = do
|
|
v <- newEmptyMVar
|
|
void $ async $ relayFeeder runner v
|
|
void $ async $ relayReader v hout
|
|
return v
|
|
|
|
cleanup _ = do
|
|
hClose hin
|
|
hClose hout
|
|
|
|
go v = relayHelper runner v hin
|
|
|
|
runRelayService :: RunEnv -> RunProto IO -> Service -> IO (Maybe ())
|
|
runRelayService runenv runner service = bracket setup cleanup go
|
|
where
|
|
cmd = case service of
|
|
UploadPack -> "upload-pack"
|
|
ReceivePack -> "receive-pack"
|
|
|
|
serviceproc = gitCreateProcess
|
|
[ Param cmd
|
|
, File (repoPath (runRepo runenv))
|
|
] (runRepo runenv)
|
|
|
|
setup = do
|
|
(Just hin, Just hout, _, pid) <- createProcess serviceproc
|
|
{ std_out = CreatePipe
|
|
, std_in = CreatePipe
|
|
}
|
|
v <- newEmptyMVar
|
|
void $ async $ relayFeeder runner v
|
|
void $ async $ relayReader v hout
|
|
waiter <- async $ waitexit v pid
|
|
return (v, waiter, hin, hout, pid)
|
|
|
|
cleanup (_, waiter, hin, hout, pid) = do
|
|
hClose hin
|
|
hClose hout
|
|
cancel waiter
|
|
void $ waitForProcess pid
|
|
|
|
go (v, _, hin, _, _) = do
|
|
r <- relayHelper runner v hin
|
|
case r of
|
|
Nothing -> return Nothing
|
|
Just exitcode -> runner $ net $ relayToPeer (RelayDone exitcode)
|
|
|
|
waitexit v pid = putMVar v . RelayDone =<< waitForProcess pid
|
|
|
|
-- Processes RelayData as it is put into the MVar.
|
|
relayHelper :: RunProto IO -> MVar RelayData -> Handle -> IO (Maybe ExitCode)
|
|
relayHelper runner v hin = loop
|
|
where
|
|
loop = do
|
|
d <- takeMVar v
|
|
case d of
|
|
RelayFromPeer b -> do
|
|
L.hPut hin b
|
|
hFlush hin
|
|
loop
|
|
RelayToPeer b -> do
|
|
r <- runner $ net $ relayToPeer (RelayToPeer b)
|
|
case r of
|
|
Nothing -> return Nothing
|
|
Just () -> loop
|
|
RelayDone exitcode -> do
|
|
_ <- runner $ net $ relayToPeer (RelayDone exitcode)
|
|
return (Just exitcode)
|
|
|
|
-- Takes input from the peer, and puts it into the MVar for processing.
|
|
-- Repeats until the peer tells it it's done or hangs up.
|
|
relayFeeder :: RunProto IO -> MVar RelayData -> IO ()
|
|
relayFeeder runner v = loop
|
|
where
|
|
loop = do
|
|
mrd <- runner $ net relayFromPeer
|
|
case mrd of
|
|
Nothing -> putMVar v (RelayDone (ExitFailure 1))
|
|
Just rd -> do
|
|
putMVar v rd
|
|
case rd of
|
|
RelayDone _ -> return ()
|
|
_ -> loop
|
|
|
|
-- Reads input from the Handle and puts it into the MVar for relaying to
|
|
-- the peer. Continues until EOF on the Handle.
|
|
relayReader :: MVar RelayData -> Handle -> IO ()
|
|
relayReader v hout = loop
|
|
where
|
|
loop = do
|
|
bs <- getsome []
|
|
case bs of
|
|
[] -> return ()
|
|
_ -> do
|
|
putMVar v $ RelayToPeer (L.fromChunks bs)
|
|
loop
|
|
|
|
-- Waiit for the first available chunk. Then, without blocking,
|
|
-- try to get more chunks, in case a stream of chunks is being
|
|
-- written in close succession.
|
|
--
|
|
-- On Windows, hGetNonBlocking is broken, so avoid using it there.
|
|
getsome [] = do
|
|
b <- B.hGetSome hout chunk
|
|
if B.null b
|
|
then return []
|
|
#ifndef mingw32_HOST_OS
|
|
else getsome [b]
|
|
#else
|
|
else return [b]
|
|
#endif
|
|
getsome bs = do
|
|
b <- B.hGetNonBlocking hout chunk
|
|
if B.null b
|
|
then return (reverse bs)
|
|
else getsome (b:bs)
|
|
|
|
chunk = 65536
|