@@ -19,12 +19,13 @@ import Data.ByteString.Internal
1919import Data.IntMap (IntMap)
2020import qualified Data.IntMap as IntMap
2121
22−import Data.Int
22+import Data.Int (Int32)
2323import Data.Serialize
24−import Data.Word
24+import Data.Word (Word8)
2525
2626import Foreign.ForeignPtr (withForeignPtr)
2727import Foreign.Ptr (Ptr, plusPtr)
28+import Foreign.Storable (sizeOf)
2829
2930import Network.Socket
3031 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
@@ -112,10 +113,10 @@ mkSourceEnd host service chanId = withSocketsDo $ do
112113 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
113114 setSocketOption sock ReuseAddr 1
114115 N.connect sock (addrAddress serverAddr)
115− NBS.sendAll sock $ encode (fromIntegral chanId :: Int64)
116+ NBS.sendAll sock $ encode (fromIntegral chanId :: Int32)
116117 return SourceEnd
117118 { send = {-# SCC "send" #-} \bss -> do
118− let size = fromIntegral (sum . map BS.length $ bss) :: Int64
119+ let size = fromIntegral (sum . map BS.length $ bss) :: Int32
119120 sizeStr | size < 255 = encode (fromIntegral size :: Word8)
120121 | otherwise = BS.cons (toEnum 255) (encode size)
121122 NBS.sendMany sock (sizeStr:bss)
@@ -145,11 +146,11 @@ procConnections chans sock = forever $ do
145146 let err m = error $ printf "procConnections: %s" m
146147 (clientSock, _clientAddr) <- accept sock
147148 -- decode the first message to find the correct chanId
148− bss <- recvExact clientSock 8
149+ bss <- recvExact clientSock (fromIntegral $ (sizeOf (0 :: Int32)))
149150 case bss of
150151 Left _ -> err "inbound chanId aborted"
151152 Right bs -> do
152− let chanId = either err fromIntegral (decode bs :: Either String Int64)
153+ let chanId = either err fromIntegral (decode bs :: Either String Int32)
153154 (chanId', chanMap) <- takeMVar chans
154155 case IntMap.lookup chanId chanMap of
155156 Nothing -> do
@@ -175,7 +176,7 @@ procMessages chans chanId chan sock = do
175176 Right sizeBS -> do
176177 case decode sizeBS :: Either String Word8 of
177178 Right 255 -> do
178− esizeBS' <- recvExact sock 8
179+ esizeBS' <- recvExact sock (fromIntegral $ sizeOf (0 :: Int32))
179180 case esizeBS' of
180181 Left _ -> closeSocket
181182 Right sizeBS' -> procMessage $ decode sizeBS'
@@ -194,7 +195,7 @@ procMessages chans chanId chan sock = do
194195 let chanMap' = IntMap.insert chanId (chan, socks') chanMap
195196 putMVar chans (chanId', chanMap')
196197 sClose sock
197− procMessage :: Either String Int64 -> IO ()
198+ procMessage :: Either String Int32 -> IO ()
198199 procMessage (Left msg) = err msg
199200 procMessage (Right size) = do
200201 ebs <- recvExact sock size
@@ -216,19 +217,20 @@ procMessages chans chanId chan sock = do
216217-- a partial retrieval since the socket was closed, and `return Right bs`
217218-- indicates success. This hasn't been implemented, since a Transport `receive`
218219-- represents an atomic receipt of a message.
219−recvExact :: Socket -> Int64 -> IO (Either ByteString ByteString)
220+recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
220221recvExact sock l = do
221222 res <- createAndTrim (fromIntegral l) (go 0)
222223 case BS.length res of
223224 n | n == (fromIntegral l) -> return $ Right res
224225 n -> do
225− printf "recvExact: expected %dB, got %dB\n" l n
226+-- printf "recvExact: expected %dB, got %dB\n" l n
226227 return $ Left res
227228 where
228229 go :: Int -> Ptr Word8 -> IO Int
229230 go n ptr | n == (fromIntegral l) = return n
230231 go n ptr = do
231− (p, off, len) <- toForeignPtr <$> NBS.recv sock (fromIntegral l)
232+ (p, off, len) <- toForeignPtr <$> NBS.recv sock (min (fromIntegral l-n) 4096)
233+-- printf "recvExact: received %d/%d\n" (n+len) l
232234 if len == 0
233235 then return n
234236 else withForeignPtr p $ \p -> do