make sure that the specified number of bytes of DATA are always sent

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.
This commit is contained in:
Joey Hess 2016-12-02 13:45:45 -04:00
parent 3dce6a080e
commit 15dc63d47f
No known key found for this signature in database
GPG key ID: C910D9222512E3C7

View file

@ -5,10 +5,11 @@
- Licensed under the GNU GPL version 3 or higher.
-}
{-# LANGUAGE RankNTypes, ScopedTypeVariables, FlexibleContexts, CPP #-}
{-# LANGUAGE RankNTypes, FlexibleContexts, BangPatterns, CPP #-}
module P2P.IO
( RunEnv(..)
( RunProto
, RunEnv(..)
, runNetProto
, runNet
) where
@ -55,7 +56,7 @@ runNetProto runenv = go
go (Free (Net n)) = runNet runenv go n
go (Free (Local _)) = return Nothing
-- Interprater of the Net part of Proto.
-- Interpreter of the Net part of Proto.
--
-- An interpreter of Proto has to be provided, to handle the rest of Proto
-- actions.
@ -78,13 +79,14 @@ runNet runenv runner f = case f of
let e = ERROR $ "protocol parse error: " ++ show l
net $ sendMessage e
next e
SendBytes _len b next -> do
SendBytes len b next -> do
v <- liftIO $ tryIO $ do
L.hPut (runOhdl runenv) b
ok <- sendExactly len b (runOhdl runenv)
hFlush (runOhdl runenv)
return ok
case v of
Left _e -> return Nothing
Right () -> runner next
Right True -> runner next
_ -> return Nothing
ReceiveBytes (Len n) next -> do
v <- liftIO $ tryIO $ L.hGet (runIhdl runenv) (fromIntegral n)
case v of
@ -109,6 +111,30 @@ runNet runenv runner f = case f of
-- 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