@@ -19,12 +19,13 @@ import Data.ByteString.Internal
19import Data.IntMap (IntMap)
19import Data.IntMap (IntMap)
20import qualified Data.IntMap as IntMap
20import qualified Data.IntMap as IntMap
22−import Data.Int
22+import Data.Int (Int32)
23import Data.Serialize
23import Data.Serialize
24−import Data.Word
24+import Data.Word (Word8)
26import Foreign.ForeignPtr (withForeignPtr)
26import Foreign.ForeignPtr (withForeignPtr)
27import Foreign.Ptr (Ptr, plusPtr)
27import Foreign.Ptr (Ptr, plusPtr)
28+import Foreign.Storable (sizeOf)
29import Network.Socket
30import Network.Socket
30 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
31 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
@@ -112,10 +113,10 @@ mkSourceEnd host service chanId = withSocketsDo $ do
112 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
113 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
113 setSocketOption sock ReuseAddr 1
114 setSocketOption sock ReuseAddr 1
114 N.connect sock (addrAddress serverAddr)
115 N.connect sock (addrAddress serverAddr)
115− NBS.sendAll sock $ encode (fromIntegral chanId :: Int64)
116+ NBS.sendAll sock $ encode (fromIntegral chanId :: Int32)
116 return SourceEnd
117 return SourceEnd
117 { send = {-# SCC "send" #-} \bss -> do
118 { 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
119 sizeStr | size < 255 = encode (fromIntegral size :: Word8)
120 sizeStr | size < 255 = encode (fromIntegral size :: Word8)
120 | otherwise = BS.cons (toEnum 255) (encode size)
121 | otherwise = BS.cons (toEnum 255) (encode size)
121 NBS.sendMany sock (sizeStr:bss)
122 NBS.sendMany sock (sizeStr:bss)
@@ -145,11 +146,11 @@ procConnections chans sock = forever $ do
145 let err m = error $ printf "procConnections: %s" m
146 let err m = error $ printf "procConnections: %s" m
146 (clientSock, _clientAddr) <- accept sock
147 (clientSock, _clientAddr) <- accept sock
147 -- decode the first message to find the correct chanId
148 -- decode the first message to find the correct chanId
148− bss <- recvExact clientSock 8
149+ bss <- recvExact clientSock (fromIntegral $ (sizeOf (0 :: Int32)))
149 case bss of
150 case bss of
150 Left _ -> err "inbound chanId aborted"
151 Left _ -> err "inbound chanId aborted"
151 Right bs -> do
152 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)
153 (chanId', chanMap) <- takeMVar chans
154 (chanId', chanMap) <- takeMVar chans
154 case IntMap.lookup chanId chanMap of
155 case IntMap.lookup chanId chanMap of
155 Nothing -> do
156 Nothing -> do
@@ -175,7 +176,7 @@ procMessages chans chanId chan sock = do
175 Right sizeBS -> do
176 Right sizeBS -> do
176 case decode sizeBS :: Either String Word8 of
177 case decode sizeBS :: Either String Word8 of
177 Right 255 -> do
178 Right 255 -> do
178− esizeBS' <- recvExact sock 8
179+ esizeBS' <- recvExact sock (fromIntegral $ sizeOf (0 :: Int32))
179 case esizeBS' of
180 case esizeBS' of
180 Left _ -> closeSocket
181 Left _ -> closeSocket
181 Right sizeBS' -> procMessage $ decode sizeBS'
182 Right sizeBS' -> procMessage $ decode sizeBS'
@@ -194,7 +195,7 @@ procMessages chans chanId chan sock = do
194 let chanMap' = IntMap.insert chanId (chan, socks') chanMap
195 let chanMap' = IntMap.insert chanId (chan, socks') chanMap
195 putMVar chans (chanId', chanMap')
196 putMVar chans (chanId', chanMap')
196 sClose sock
197 sClose sock
197− procMessage :: Either String Int64 -> IO ()
198+ procMessage :: Either String Int32 -> IO ()
198 procMessage (Left msg) = err msg
199 procMessage (Left msg) = err msg
199 procMessage (Right size) = do
200 procMessage (Right size) = do
200 ebs <- recvExact sock size
201 ebs <- recvExact sock size
@@ -216,19 +217,20 @@ procMessages chans chanId chan sock = do
216-- a partial retrieval since the socket was closed, and `return Right bs`
217-- a partial retrieval since the socket was closed, and `return Right bs`
217-- indicates success. This hasn't been implemented, since a Transport `receive`
218-- indicates success. This hasn't been implemented, since a Transport `receive`
218-- represents an atomic receipt of a message.
219-- represents an atomic receipt of a message.
219−recvExact :: Socket -> Int64 -> IO (Either ByteString ByteString)
220+recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
220recvExact sock l = do
221recvExact sock l = do
221 res <- createAndTrim (fromIntegral l) (go 0)
222 res <- createAndTrim (fromIntegral l) (go 0)
222 case BS.length res of
223 case BS.length res of
223 n | n == (fromIntegral l) -> return $ Right res
224 n | n == (fromIntegral l) -> return $ Right res
225− printf "recvExact: expected %dB, got %dB\n" l n
226+-- printf "recvExact: expected %dB, got %dB\n" l n
226 return $ Left res
227 return $ Left res
228 go :: Int -> Ptr Word8 -> IO Int
229 go :: Int -> Ptr Word8 -> IO Int
229 go n ptr | n == (fromIntegral l) = return n
230 go n ptr | n == (fromIntegral l) = return n
230 go n ptr = do
231 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
232 if len == 0
234 if len == 0
233 then return n
235 then return n
234 else withForeignPtr p $ \p -> do
236 else withForeignPtr p $ \p -> do