recvWithLength takes optional limit

Remove recvWithLengthLimited

Author: avieth

src/Network/Transport/TCP.hs

@@ -55,7 +55,7 @@ import Network.Transport.TCP.Internal
   , encodeConnectionRequestResponse
   , decodeConnectionRequestResponse
   , forkServer
-  , recvWithLengthLimited
+  , recvWithLength
   , recvWord32
   , encodeWord32
   , tryCloseSocket
@@ -878,7 +878,7 @@ handleConnectionRequest transport sock = handle handleException $ do
     ourEndPointId <- recvWord32 sock
     let maxAddressLength = tcpMaxAddressLength $ transportParams transport
     theirAddress  <- EndPointAddress . BS.concat <$>
-      recvWithLengthLimited maxAddressLength sock
+      recvWithLength maxAddressLength sock
     let ourAddress = encodeEndPointAddress (transportHost transport)
                                            (transportPort transport)
                                            ourEndPointId
@@ -1189,7 +1189,7 @@ handleIncomingMessages params (ourEndPoint, theirEndPoint) = do
     -- overhead
     readMessage :: N.Socket -> LightweightConnectionId -> IO ()
     readMessage sock lcid =
-      recvWithLengthLimited recvLimit sock >>=
+      recvWithLength recvLimit sock >>=
         qdiscEnqueue' ourQueue theirAddr . Received (connId lcid)
 
     -- Stop probing a connection as a result of receiving a probe ack.

src/Network/Transport/TCP/Internal.hs

@@ -8,7 +8,6 @@ module Network.Transport.TCP.Internal
   , decodeConnectionRequestResponse
   , forkServer
   , recvWithLength
-  , recvWithLengthLimited
   , recvExact
   , recvWord32
   , encodeWord32
@@ -181,19 +180,15 @@ forkServer host port backlog reuseAddr terminationHandler requestHandler = do
                                         (tryCloseSocket . fst)
                                         (requestHandler . fst)
 
--- | Read a length and then a payload of that length
-recvWithLength :: N.Socket -> IO [ByteString]
-recvWithLength sock = recvWord32 sock >>= recvExact sock
-
--- | Read a length and then a payload of that length, subject to a limit on
---   the length.
-recvWithLengthLimited :: Maybe Word32 -> N.Socket -> IO [ByteString]
-recvWithLengthLimited mlimit sock = case mlimit of
-  Nothing -> recvWithLength sock
+-- | 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 "recvWithLengthLimit: limit exceeded")
+      throwIO (userError "recvWithLength: limit exceeded")
     recvExact sock length
 
 -- | Receive a 32-bit unsigned integer

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 sock
+        _ <- recvWithLength Nothing 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 sock
+        ["ping"] <- recvWithLength Nothing 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 sock
+        _ <- recvWithLength Nothing 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 sock
+        ["ping"] <- recvWithLength Nothing 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 sock
+    Right _  <- tryIO $ recvWithLength Nothing 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 sock
+          Right ["ping"] <- tryIO $ recvWithLength Nothing 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 sock
+      _ <- recvWithLength Nothing 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 Nothing sock
       putMVar serverGotPing ()
 
   -- Client