Configurable limit on bytes in Received event

Since 'QDisc's have been added, it's now possible to control the rate of
incoming events. However, it's not so easy to control the rate of
incoming data, because one 'Received' event could include as many as
2^32 bytes. This patch gives a limit on the number of bytes in any
single 'Received' event. The 'recvExact' function remains, using
'maxBound :: Word32' as the limit, which was the implicit limit anyway.

Also made the maximum data payload and address lengths non-optional; the
default is 'maxBound :: Word32'.

Author: avieth

src/Network/Transport/TCP.hs

@@ -56,6 +56,7 @@ import Network.Transport.TCP.Internal
   , decodeConnectionRequestResponse
   , forkServer
   , recvWithLength
+  , recvWithLengthFold
   , recvWord32
   , encodeWord32
   , tryCloseSocket
@@ -482,16 +483,18 @@ data TCPParameters = TCPParameters {
   , transportConnectTimeout :: Maybe Int
     -- | Create a QDisc for an EndPoint.
   , tcpNewQDisc :: forall t . IO (QDisc t)
-    -- | Optional maximum length (in bytes) for a peer's address.
+    -- | Maximum length (in bytes) for a peer's address.
     -- If a peer attempts to send an address of length exceeding the limit,
     -- the connection will be refused (socket will close).
-  , tcpMaxAddressLength :: Maybe Word32
-    -- | Optional maximum length (in bytes) to receive from a peer.
+  , tcpMaxAddressLength :: Word32
+    -- | Maximum length (in bytes) to receive from a peer.
     -- If a peer attempts to send data on a lightweight connection exceeding
     -- the limit, the heavyweight connection which carries that lightweight
     -- connection will go down. The peer and the local node will get an
     -- EventConnectionLost.
-  , tcpMaxReceiveLength :: Maybe Word32
+  , tcpMaxReceiveLength :: Word32
+    -- | Maximum length (in bytes) of a 'Received' event payload.
+  , tcpMaxChunkSize :: Word32
   }
 
 -- | Internal functionality we expose for unit testing
@@ -598,8 +601,9 @@ defaultTCPParameters = TCPParameters {
   , tcpUserTimeout     = Nothing
   , tcpNewQDisc        = simpleUnboundedQDisc
   , transportConnectTimeout = Nothing
-  , tcpMaxAddressLength = Nothing
-  , tcpMaxReceiveLength = Nothing
+  , tcpMaxAddressLength = maxBound
+  , tcpMaxReceiveLength = maxBound
+  , tcpMaxChunkSize     = maxBound
   }
 
 --------------------------------------------------------------------------------
@@ -1189,8 +1193,8 @@ handleIncomingMessages params (ourEndPoint, theirEndPoint) = do
     -- overhead
     readMessage :: N.Socket -> LightweightConnectionId -> IO ()
     readMessage sock lcid =
-      recvWithLength recvLimit sock >>=
-        qdiscEnqueue' ourQueue theirAddr . Received (connId lcid)
+      recvWithLengthFold recvLimit chunkLimit sock () $ \bs _ ->
+        qdiscEnqueue' ourQueue theirAddr (Received (connId lcid) bs)
 
     -- Stop probing a connection as a result of receiving a probe ack.
     stopProbing :: IO ()
@@ -1206,6 +1210,7 @@ handleIncomingMessages params (ourEndPoint, theirEndPoint) = do
     theirState  = remoteState theirEndPoint
     theirAddr   = remoteAddress theirEndPoint
     recvLimit   = tcpMaxReceiveLength params
+    chunkLimit  = tcpMaxChunkSize params
 
     -- Deal with a premature exit
     prematureExit :: N.Socket -> IOException -> IO ()

src/Network/Transport/TCP/Internal.hs

@@ -1,3 +1,5 @@
+{-# LANGUAGE BangPatterns #-}
+
 -- | Utility functions for TCP sockets
 module Network.Transport.TCP.Internal
   ( ControlHeader(..)
@@ -8,6 +10,7 @@ module Network.Transport.TCP.Internal
   , decodeConnectionRequestResponse
   , forkServer
   , recvWithLength
+  , recvWithLengthFold
   , recvExact
   , recvWord32
   , encodeWord32
@@ -59,7 +62,7 @@ import qualified Network.Socket.ByteString as NBS (recv)
 import Control.Concurrent (ThreadId)
 import Data.Word (Word32)
 
-import Control.Monad (forever, when)
+import Control.Monad (forever, when, unless)
 import Control.Exception (SomeException, catch, bracketOnError, throwIO, mask_)
 import Control.Applicative ((<$>), (<*>))
 import Data.Word (Word32)
@@ -180,20 +183,44 @@ forkServer host port backlog reuseAddr terminationHandler requestHandler = do
                                         (tryCloseSocket . fst)
                                         (requestHandler . fst)
 
--- | Read a length and then a payload of that length, subject to an optional
---   limit on the length.
-recvWithLength :: Maybe Word32 -> N.Socket -> IO [ByteString]
-recvWithLength mlimit sock = case mlimit of
-  Nothing -> recvWord32 sock >>= recvExact sock
-  Just limit -> do
-    length <- recvWord32 sock
-    when (length > limit) $
-      throwIO (userError "recvWithLength: limit exceeded")
-    recvExact sock length
+-- | Read a length, then 1 or more payloads each less than some maximum
+-- length in bytes, such that the sum of their lengths is the length that was
+-- read.
+recvWithLengthFold
+  :: Word32                      -- ^ Maximum total size.
+  -> Word32                      -- ^ Maximum chunk size.
+  -> N.Socket
+  -> t                           -- ^ Start element for the fold.
+  -> ([ByteString] -> t -> IO t) -- ^ Run this every time we get data of at
+                                 -- most the maximum size.
+  -> IO t
+recvWithLengthFold maxSize maxChunk sock base folder = do
+  len <- recvWord32 sock
+  when (len > maxSize) $
+    throwIO (userError "recvWithLengthFold: limit exceeded")
+  loop base len
+  where
+  loop !base !total = do
+    (bs, received) <- recvExact sock (min maxChunk total)
+    base' <- folder bs base
+    let remaining = total - received
+    when (received > total) $ throwIO (userError "recvWithLengthFold: got more bytes than requested")
+    if remaining == 0
+    then return base'
+    else loop base' remaining
+
+-- | Read a length and then a payload of that length
+recvWithLength
+  :: Word32          -- ^ Maximum total size.
+  -> N.Socket
+  -> IO [ByteString]
+recvWithLength maxSize sock = fmap (concat . reverse) $
+  recvWithLengthFold maxSize maxBound sock [] $
+    \bs lst -> return (bs : lst)
 
 -- | Receive a 32-bit unsigned integer
 recvWord32 :: N.Socket -> IO Word32
-recvWord32 = fmap (decodeWord32 . BS.concat) . flip recvExact 4
+recvWord32 = fmap (decodeWord32 . BS.concat . fst) . flip recvExact 4
 
 -- | Close a socket, ignoring I/O exceptions.
 tryCloseSocket :: N.Socket -> IO ()
@@ -204,16 +231,22 @@ tryCloseSocket sock = void . tryIO $
 --
 -- Throws an I/O exception if the socket closes before the specified
 -- number of bytes could be read
-recvExact :: N.Socket                -- ^ Socket to read from
-          -> Word32                  -- ^ Number of bytes to read
-          -> IO [ByteString]
+recvExact :: N.Socket                  -- ^ Socket to read from
+          -> Word32                    -- ^ Number of bytes to read
+          -> IO ([ByteString], Word32) -- ^ Data and number of bytes read
 recvExact _ len | len < 0 = throwIO (userError "recvExact: Negative length")
-recvExact sock len = go [] len
+recvExact sock len = go [] 0 len
   where
-    go :: [ByteString] -> Word32 -> IO [ByteString]
-    go acc 0 = return (reverse acc)
-    go acc l = do
+    go :: [ByteString] -> Word32 -> Word32 -> IO ([ByteString], Word32)
+    go acc !n 0 = return (reverse acc, n)
+    go acc !n l = do
       bs <- NBS.recv sock (fromIntegral l `min` smallChunkSize)
       if BS.null bs
         then throwIO (userError "recvExact: Socket closed")
-        else go (bs : acc) (l - fromIntegral (BS.length bs))
+        else do
+          let received  = fromIntegral (BS.length bs)
+              remaining = l - received
+              total     = n + received
+          -- Check for underflow. Shouldn't be possible but let's make sure.
+          when (received > l) $ throwIO (userError "recvExact: got more bytes than requested")
+          go (bs : acc) total remaining

tests/TestTCP.hs

@@ -165,7 +165,7 @@ testEarlyDisconnect = do
       (clientPort, _) <- forkServer "127.0.0.1" "0" 5 True throwIO $ \sock -> do
         -- Initial setup
         0 <- recvWord32 sock
-        _ <- recvWithLength Nothing sock
+        _ <- recvWithLength maxBound sock
         sendMany sock [encodeWord32 (encodeConnectionRequestResponse ConnectionRequestAccepted)]
 
         -- Server opens  a logical connection
@@ -174,7 +174,7 @@ testEarlyDisconnect = do
 
         -- Server sends a message
         1024 <- recvWord32 sock
-        ["ping"] <- recvWithLength Nothing sock
+        ["ping"] <- recvWithLength maxBound sock
 
         -- Reply
         sendMany sock [
@@ -277,7 +277,7 @@ testEarlyCloseSocket = do
       (clientPort, _) <- forkServer "127.0.0.1" "0" 5 True throwIO $ \sock -> do
         -- Initial setup
         0 <- recvWord32 sock
-        _ <- recvWithLength Nothing sock
+        _ <- recvWithLength maxBound sock
         sendMany sock [encodeWord32 (encodeConnectionRequestResponse ConnectionRequestAccepted)]
 
         -- Server opens a logical connection
@@ -286,7 +286,7 @@ testEarlyCloseSocket = do
 
         -- Server sends a message
         1024 <- recvWord32 sock
-        ["ping"] <- recvWithLength Nothing sock
+        ["ping"] <- recvWithLength maxBound sock
 
         -- Reply
         sendMany sock [
@@ -620,7 +620,7 @@ testReconnect = do
   (serverPort, _) <- forkServer "127.0.0.1" "0" 5 True throwIO $ \sock -> do
     -- Accept the connection
     Right 0  <- tryIO $ recvWord32 sock
-    Right _  <- tryIO $ recvWithLength Nothing sock
+    Right _  <- tryIO $ recvWithLength maxBound sock
 
     -- The first time we close the socket before accepting the logical connection
     count <- modifyMVar counter $ \i -> return (i + 1, i)
@@ -639,7 +639,7 @@ testReconnect = do
           -- Client sends a message
           Right connId' <- tryIO $ (recvWord32 sock :: IO LightweightConnectionId)
           True <- return $ connId == connId'
-          Right ["ping"] <- tryIO $ recvWithLength Nothing sock
+          Right ["ping"] <- tryIO $ recvWithLength maxBound sock
           putMVar serverDone ()
 
     Right () <- tryIO $ N.sClose sock
@@ -712,15 +712,15 @@ testUnidirectionalError = do
     -- would shutdown the socket in the other direction)
     void . (try :: IO () -> IO (Either SomeException ())) $ do
       0 <- recvWord32 sock
-      _ <- recvWithLength Nothing sock
+      _ <- recvWithLength maxBound sock
       () <- sendMany sock [encodeWord32 (encodeConnectionRequestResponse ConnectionRequestAccepted)]
 
       Just CreatedNewConnection <- decodeControlHeader <$> recvWord32 sock
       connId <- recvWord32 sock :: IO LightweightConnectionId
 
       connId' <- recvWord32 sock :: IO LightweightConnectionId
       True <- return $ connId == connId'
-      ["ping"] <- recvWithLength Nothing sock
+      ["ping"] <- recvWithLength maxBound sock
       putMVar serverGotPing ()
 
   -- Client
@@ -843,8 +843,8 @@ testMaxLength = do
       -- Port is at most 5 bytes (65536) and id is a base-10 Word32 so
       -- at most 10 bytes. We'll have one client with a 5-byte port to push it
       -- over the chosen limit of 16
-      tcpMaxAddressLength = Just 16
-    , tcpMaxReceiveLength = Just 8
+      tcpMaxAddressLength = 16
+    , tcpMaxReceiveLength = 8
     }
   Right goodClientTransport <- createTransport "127.0.0.1" "9999" ((,) "127.0.0.1") defaultTCPParameters
   Right badClientTransport <- createTransport "127.0.0.1" "10000" ((,) "127.0.0.1") defaultTCPParameters