Merge pull request #52 from avieth/avieth/limit_length

Optional limits on address and data length

Author: facundominguez

src/Network/Transport/TCP.hs

@@ -56,6 +56,7 @@ import Network.Transport.TCP.Internal
   , decodeConnectionRequestResponse
   , forkServer
   , recvWithLength
+  , recvWithLengthFold
   , recvWord32
   , encodeWord32
   , tryCloseSocket
@@ -482,6 +483,18 @@ data TCPParameters = TCPParameters {
   , transportConnectTimeout :: Maybe Int
     -- | Create a QDisc for an EndPoint.
   , tcpNewQDisc :: forall t . IO (QDisc t)
+    -- | 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 :: 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 :: Word32
+    -- | Maximum length (in bytes) of a 'Received' event payload.
+  , tcpMaxChunkSize :: Word32
   }
 
 -- | Internal functionality we expose for unit testing
@@ -588,6 +601,9 @@ defaultTCPParameters = TCPParameters {
   , tcpUserTimeout     = Nothing
   , tcpNewQDisc        = simpleUnboundedQDisc
   , transportConnectTimeout = Nothing
+  , tcpMaxAddressLength = maxBound
+  , tcpMaxReceiveLength = maxBound
+  , tcpMaxChunkSize     = maxBound
   }
 
 --------------------------------------------------------------------------------
@@ -864,7 +880,9 @@ handleConnectionRequest transport sock = handle handleException $ do
     forM_ (tcpUserTimeout $ transportParams transport) $
       N.setSocketOption sock N.UserTimeout
     ourEndPointId <- recvWord32 sock
-    theirAddress  <- EndPointAddress . BS.concat <$> recvWithLength sock
+    let maxAddressLength = tcpMaxAddressLength $ transportParams transport
+    theirAddress  <- EndPointAddress . BS.concat <$>
+      recvWithLength maxAddressLength sock
     let ourAddress = encodeEndPointAddress (transportHost transport)
                                            (transportPort transport)
                                            ourEndPointId
@@ -914,7 +932,7 @@ handleConnectionRequest transport sock = handle handleException $ do
       -- been recorded as part of the remote endpoint. Either way, we no longer
       -- have to worry about closing the socket on receiving an asynchronous
       -- exception from this point forward.
-      forM_ mEndPoint $ handleIncomingMessages . (,) ourEndPoint
+      forM_ mEndPoint $ handleIncomingMessages (transportParams transport) . (,) ourEndPoint
 
     handleException :: SomeException -> IO ()
     handleException ex = do
@@ -957,8 +975,8 @@ handleConnectionRequest transport sock = handle handleException $ do
 --
 -- Returns only if the remote party closes the socket or if an error occurs.
 -- This runs in a thread that will never be killed.
-handleIncomingMessages :: EndPointPair -> IO ()
-handleIncomingMessages (ourEndPoint, theirEndPoint) = do
+handleIncomingMessages :: TCPParameters -> EndPointPair -> IO ()
+handleIncomingMessages params (ourEndPoint, theirEndPoint) = do
     mSock <- withMVar theirState $ \st ->
       case st of
         RemoteEndPointInvalid _ ->
@@ -1175,7 +1193,8 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
     -- overhead
     readMessage :: N.Socket -> LightweightConnectionId -> IO ()
     readMessage sock lcid =
-      recvWithLength 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 ()
@@ -1190,6 +1209,8 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
     ourQueue    = localQueue ourEndPoint
     theirState  = remoteState theirEndPoint
     theirAddr   = remoteAddress theirEndPoint
+    recvLimit   = tcpMaxReceiveLength params
+    chunkLimit  = tcpMaxChunkSize params
 
     -- Deal with a premature exit
     prematureExit :: N.Socket -> IOException -> IO ()
@@ -1365,7 +1386,7 @@ setupRemoteEndPoint params (ourEndPoint, theirEndPoint) connTimeout = do
         return False
 
     when didAccept $ void $ forkIO $
-      handleIncomingMessages (ourEndPoint, theirEndPoint)
+      handleIncomingMessages params (ourEndPoint, theirEndPoint)
     return $ either (const Nothing) (Just . snd) result
   where
     ourAddress      = localAddress ourEndPoint

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,9 +62,10 @@ 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)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as BS (length, concat, null)
 import Data.ByteString.Lazy.Internal (smallChunkSize)
@@ -179,13 +183,44 @@ forkServer host port backlog reuseAddr terminationHandler requestHandler = do
                                         (tryCloseSocket . fst)
                                         (requestHandler . fst)
 
+-- | 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 :: N.Socket -> IO [ByteString]
-recvWithLength sock = recvWord32 sock >>= recvExact sock
+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 ()
@@ -196,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

@@ -17,6 +17,7 @@ import Network.Transport.TCP ( createTransport
                              , createTransportExposeInternals
                              , TransportInternals(..)
                              , encodeEndPointAddress
+                             , TCPParameters(..)
                              , defaultTCPParameters
                              , LightweightConnectionId
                              )
@@ -164,7 +165,7 @@ testEarlyDisconnect = do
       (clientPort, _) <- forkServer "127.0.0.1" "0" 5 True throwIO $ \sock -> do
         -- Initial setup
         0 <- recvWord32 sock
-        _ <- recvWithLength sock
+        _ <- recvWithLength maxBound sock
         sendMany sock [encodeWord32 (encodeConnectionRequestResponse ConnectionRequestAccepted)]
 
         -- Server opens  a logical connection
@@ -173,7 +174,7 @@ testEarlyDisconnect = do
 
         -- Server sends a message
         1024 <- recvWord32 sock
-        ["ping"] <- recvWithLength sock
+        ["ping"] <- recvWithLength maxBound sock
 
         -- Reply
         sendMany sock [
@@ -276,7 +277,7 @@ testEarlyCloseSocket = do
       (clientPort, _) <- forkServer "127.0.0.1" "0" 5 True throwIO $ \sock -> do
         -- Initial setup
         0 <- recvWord32 sock
-        _ <- recvWithLength sock
+        _ <- recvWithLength maxBound sock
         sendMany sock [encodeWord32 (encodeConnectionRequestResponse ConnectionRequestAccepted)]
 
         -- Server opens a logical connection
@@ -285,7 +286,7 @@ testEarlyCloseSocket = do
 
         -- Server sends a message
         1024 <- recvWord32 sock
-        ["ping"] <- recvWithLength sock
+        ["ping"] <- recvWithLength maxBound sock
 
         -- Reply
         sendMany sock [
@@ -619,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 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)
@@ -638,7 +639,7 @@ testReconnect = do
           -- Client sends a message
           Right connId' <- tryIO $ (recvWord32 sock :: IO LightweightConnectionId)
           True <- return $ connId == connId'
-          Right ["ping"] <- tryIO $ recvWithLength sock
+          Right ["ping"] <- tryIO $ recvWithLength maxBound sock
           putMVar serverDone ()
 
     Right () <- tryIO $ N.sClose sock
@@ -711,15 +712,15 @@ testUnidirectionalError = do
     -- would shutdown the socket in the other direction)
     void . (try :: IO () -> IO (Either SomeException ())) $ do
       0 <- recvWord32 sock
-      _ <- recvWithLength 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 sock
+      ["ping"] <- recvWithLength maxBound sock
       putMVar serverGotPing ()
 
   -- Client
@@ -831,10 +832,77 @@ testUseRandomPort = do
      putMVar testDone ()
    takeMVar testDone
 
+-- | Verify that if a peer sends an address or data which exceeds the maximum
+--   length, that peer's connection will be terminated, but other peers will
+--   not be affected.
+testMaxLength :: IO ()
+testMaxLength = do
+
+  Right serverTransport <- createTransport "127.0.0.1" "9998" ((,) "127.0.0.1") $ defaultTCPParameters {
+      -- 17 bytes should fit every valid address at 127.0.0.1.
+      -- 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 = 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
+
+  serverAddress <- newEmptyMVar
+  testDone <- newEmptyMVar
+  goodClientConnected <- newEmptyMVar
+  goodClientDone <- newEmptyMVar
+  badClientDone <- newEmptyMVar
+
+  forkTry $ do
+    Right serverEp <- newEndPoint serverTransport
+    putMVar serverAddress (address serverEp)
+    readMVar badClientDone
+    ConnectionOpened _ _ _ <- receive serverEp
+    Received _ _ <- receive serverEp
+    -- Will lose the connection when the good client sends 9 bytes.
+    ErrorEvent (TransportError (EventConnectionLost _) _) <- receive serverEp
+    readMVar goodClientDone
+    putMVar testDone ()
+
+  forkTry $ do
+    Right badClientEp <- newEndPoint badClientTransport
+    address <- readMVar serverAddress
+    -- Wait until the good client connects, then try to connect. It'll fail,
+    -- but the good client should still be OK.
+    readMVar goodClientConnected
+    Left (TransportError ConnectFailed _)
+      <- connect badClientEp address ReliableOrdered defaultConnectHints
+    closeEndPoint badClientEp
+    putMVar badClientDone ()
+
+  forkTry $ do
+    Right goodClientEp <- newEndPoint goodClientTransport
+    address <- readMVar serverAddress
+    Right conn <- connect goodClientEp address ReliableOrdered defaultConnectHints
+    putMVar goodClientConnected ()
+    -- Wait until the bad client has tried and failed to connect before
+    -- attempting a send, to ensure that its failure did not affect us.
+    readMVar badClientDone
+    Right () <- send conn ["00000000"]
+    -- The send which breaches the limit does not appear to fail, but the
+    -- (heavyweight) connection is now severed. We can reliably determine that
+    -- by receiving.
+    Right () <- send conn ["000000000"]
+    ErrorEvent (TransportError (EventConnectionLost _) _) <- receive goodClientEp
+    closeEndPoint goodClientEp
+    putMVar goodClientDone ()
+
+  readMVar testDone
+  closeTransport badClientTransport
+  closeTransport goodClientTransport
+  closeTransport serverTransport
+
 main :: IO ()
 main = do
   tcpResult <- tryIO $ runTests
-           [ ("Use random port"       , testUseRandomPort)
+           [ ("Use random port",        testUseRandomPort)
            , ("EarlyDisconnect",        testEarlyDisconnect)
            , ("EarlyCloseSocket",       testEarlyCloseSocket)
            , ("IgnoreCloseSocket",      testIgnoreCloseSocket)
@@ -847,6 +915,7 @@ main = do
            , ("Reconnect",              testReconnect)
            , ("UnidirectionalError",    testUnidirectionalError)
            , ("InvalidCloseConnection", testInvalidCloseConnection)
+           , ("MaxLength",              testMaxLength)
            ]
   -- Run the generic tests even if the TCP specific tests failed..
   testTransport (either (Left . show) (Right) <$>