diff options
author | Joey Hess <joeyh@joeyh.name> | 2016-12-02 13:45:45 -0400 |
---|---|---|
committer | Joey Hess <joeyh@joeyh.name> | 2016-12-02 13:45:45 -0400 |
commit | 81454244b07e00c0b250353bba1648f2d67715d9 (patch) | |
tree | 116fbdbef0a921a4ca1ec9f548d257f31b52aa59 | |
parent | 753793371aa6595e41b095952f5d67563e103290 (diff) |
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.
-rw-r--r-- | P2P/IO.hs | 40 |
1 files changed, 33 insertions, 7 deletions
@@ -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 |