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:
parent
3dce6a080e
commit
15dc63d47f
1 changed files with 33 additions and 7 deletions
40
P2P/IO.hs
40
P2P/IO.hs
|
@ -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
|
||||
|
|
Loading…
Reference in a new issue