@@ -14,6 +14,7 @@ import Control.Monad (forever, forM_)
15import Data.ByteString (ByteString)
15import Data.ByteString (ByteString)
16import qualified Data.ByteString as BS
16import qualified Data.ByteString as BS
17+import Data.ByteString.Internal
18import Data.Conduit
19import Data.Conduit
19import qualified Data.Conduit.Binary as CB
20import qualified Data.Conduit.Binary as CB
@@ -26,6 +27,9 @@ import Data.Int
26import Data.Serialize
27import Data.Serialize
27import Data.Word
28import Data.Word
30+import Foreign.ForeignPtr (withForeignPtr)
31+import Foreign.Ptr (Ptr, plusPtr)
29import Network.Socket
33import Network.Socket
30 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
34 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
31 , SocketType (Stream), SocketOption (ReuseAddr)
35 , SocketType (Stream), SocketOption (ReuseAddr)
@@ -151,9 +155,8 @@ procConnections chans sock = forever $ do
151 -- decode the first message to find the correct chanId
155 -- decode the first message to find the correct chanId
152 bss <- recvExact clientSock 8
156 bss <- recvExact clientSock 8
153 case bss of
157 case bss of
154− [] -> err "inbound chanId aborted"
158+ Left _ -> err "inbound chanId aborted"
155− bss -> do
159+ Right bs -> do
156− let bs = BS.concat bss
157 let chanId = either err fromIntegral (decode bs :: Either String Int64)
160 let chanId = either err fromIntegral (decode bs :: Either String Int64)
158 (chanId', chanMap) <- takeMVar chans
161 (chanId', chanMap) <- takeMVar chans
159 case IntMap.lookup chanId chanMap of
162 case IntMap.lookup chanId chanMap of
@@ -174,16 +177,16 @@ procConnections chans sock = forever $ do
174-- socket has closed.
177-- socket has closed.
175procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
178procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
176procMessages chans chanId chan sock = do
179procMessages chans chanId chan sock = do
177− sizeBSs <- recvExact sock 1
180+ esizeBS <- recvExact sock 1
178− case sizeBSs of
181+ case esizeBS of
179− [] -> closeSocket
182+ Left _ -> closeSocket
180− _ -> do
183+ Right sizeBS -> do
181− case decode . BS.concat $ sizeBSs :: Either String Word8 of
184+ case decode sizeBS :: Either String Word8 of
182 Right 255 -> do
185 Right 255 -> do
183− sizeBSs' <- recvExact sock 8
186+ esizeBS' <- recvExact sock 8
184− case sizeBSs' of
187+ case esizeBS' of
185− [] -> closeSocket
188+ Left _ -> closeSocket
186− _ -> procMessage (decode . BS.concat $ sizeBSs')
189+ Right sizeBS' -> procMessage $ decode sizeBS'
187 esize -> procMessage (fromIntegral <$> esize)
190 esize -> procMessage (fromIntegral <$> esize)
189 err = error $ printf "procMessages: %s"
192 err = error $ printf "procMessages: %s"
@@ -202,11 +205,11 @@ procMessages chans chanId chan sock = do
202 procMessage :: Either String Int64 -> IO ()
205 procMessage :: Either String Int64 -> IO ()
203 procMessage (Left msg) = err msg
206 procMessage (Left msg) = err msg
204 procMessage (Right size) = do
207 procMessage (Right size) = do
205− bss <- recvExact sock size
208+ ebs <- recvExact sock size
206− case bss of
209+ case ebs of
207− [] -> closeSocket
210+ Left _ -> closeSocket
208− _ -> do
211+ Right bs -> do
209− writeChan chan bss
212+ writeChan chan [bs]
210 procMessages chans chanId chan sock
213 procMessages chans chanId chan sock
212-- | The result of `recvExact sock n` is a `[ByteString]` whose concatenation
215-- | The result of `recvExact sock n` is a `[ByteString]` whose concatenation
@@ -221,14 +224,20 @@ procMessages chans chanId chan sock = do
221-- a partial retrieval since the socket was closed, and `return Right bs`
224-- a partial retrieval since the socket was closed, and `return Right bs`
222-- indicates success. This hasn't been implemented, since a Transport `receive`
225-- indicates success. This hasn't been implemented, since a Transport `receive`
223-- represents an atomic receipt of a message.
226-- represents an atomic receipt of a message.
224−recvExact :: Socket -> Int64 -> IO [ByteString]
227+recvExact :: Socket -> Int64 -> IO (Either ByteString ByteString)
225−recvExact sock n = go [] sock (fromIntegral n)
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
227− go :: [ByteString] -> Socket -> Int -> IO [ByteString]
234+ go :: Int -> Ptr Word8 -> IO Int
228− go bss _ 0 = return (reverse bss)
235+ go n ptr | n == (fromIntegral l) = return n
229− go bss sock n = do
236+ go n ptr = do
230− bs <- NBS.recv sock n
237+ (p, off, len) <- toForeignPtr <$> NBS.recv sock n
231− if BS.null bs
238+ if len == 0
232− then return []
239+ then return n
233− else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
240+ else withForeignPtr p $ \p -> do
241+ memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)