@@ -14,6 +14,7 @@ import Control.Monad (forever, forM_)
1414
1515import Data.ByteString (ByteString)
1616import qualified Data.ByteString as BS
17+import Data.ByteString.Internal
1718
1819import Data.Conduit
1920import qualified Data.Conduit.Binary as CB
@@ -26,6 +27,9 @@ import Data.Int
2627import Data.Serialize
2728import Data.Word
2829
30+import Foreign.ForeignPtr (withForeignPtr)
31+import Foreign.Ptr (Ptr, plusPtr)
32+
2933import Network.Socket
3034 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
3135 , SocketType (Stream), SocketOption (ReuseAddr)
@@ -151,9 +155,8 @@ procConnections chans sock = forever $ do
151155 -- decode the first message to find the correct chanId
152156 bss <- recvExact clientSock 8
153157 case bss of
154− [] -> err "inbound chanId aborted"
155− bss -> do
156− let bs = BS.concat bss
158+ Left _ -> err "inbound chanId aborted"
159+ Right bs -> do
157160 let chanId = either err fromIntegral (decode bs :: Either String Int64)
158161 (chanId', chanMap) <- takeMVar chans
159162 case IntMap.lookup chanId chanMap of
@@ -174,16 +177,16 @@ procConnections chans sock = forever $ do
174177-- socket has closed.
175178procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
176179procMessages chans chanId chan sock = do
177− sizeBSs <- recvExact sock 1
178− case sizeBSs of
179− [] -> closeSocket
180− _ -> do
181− case decode . BS.concat $ sizeBSs :: Either String Word8 of
180+ esizeBS <- recvExact sock 1
181+ case esizeBS of
182+ Left _ -> closeSocket
183+ Right sizeBS -> do
184+ case decode sizeBS :: Either String Word8 of
182185 Right 255 -> do
183− sizeBSs' <- recvExact sock 8
184− case sizeBSs' of
185− [] -> closeSocket
186− _ -> procMessage (decode . BS.concat $ sizeBSs')
186+ esizeBS' <- recvExact sock 8
187+ case esizeBS' of
188+ Left _ -> closeSocket
189+ Right sizeBS' -> procMessage $ decode sizeBS'
187190 esize -> procMessage (fromIntegral <$> esize)
188191 where
189192 err = error $ printf "procMessages: %s"
@@ -202,11 +205,11 @@ procMessages chans chanId chan sock = do
202205 procMessage :: Either String Int64 -> IO ()
203206 procMessage (Left msg) = err msg
204207 procMessage (Right size) = do
205− bss <- recvExact sock size
206− case bss of
207− [] -> closeSocket
208− _ -> do
209− writeChan chan bss
208+ ebs <- recvExact sock size
209+ case ebs of
210+ Left _ -> closeSocket
211+ Right bs -> do
212+ writeChan chan [bs]
210213 procMessages chans chanId chan sock
211214
212215-- | The result of `recvExact sock n` is a `[ByteString]` whose concatenation
@@ -221,14 +224,20 @@ procMessages chans chanId chan sock = do
221224-- a partial retrieval since the socket was closed, and `return Right bs`
222225-- indicates success. This hasn't been implemented, since a Transport `receive`
223226-- represents an atomic receipt of a message.
224−recvExact :: Socket -> Int64 -> IO [ByteString]
225−recvExact sock n = go [] sock (fromIntegral n)
227+recvExact :: Socket -> Int64 -> IO (Either ByteString ByteString)
228+recvExact sock l = do
229+ res <- createAndTrim (fromIntegral l) (go 0)
230+ case BS.length res of
231+ n | n == (fromIntegral l) -> return $ Right res
232+ _ -> return $ Left res
226233 where
227− go :: [ByteString] -> Socket -> Int -> IO [ByteString]
228− go bss _ 0 = return (reverse bss)
229− go bss sock n = do
230− bs <- NBS.recv sock n
231− if BS.null bs
232− then return []
233− else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
234+ go :: Int -> Ptr Word8 -> IO Int
235+ go n ptr | n == (fromIntegral l) = return n
236+ go n ptr = do
237+ (p, off, len) <- toForeignPtr <$> NBS.recv sock n
238+ if len == 0
239+ then return n
240+ else withForeignPtr p $ \p -> do
241+ memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)
242+ go (n+len) ptr
234243