@@ -138,16 +138,14 @@ procConnections chans sock = forever $ do
138 let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
138 let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
139 putMVar chans (chanId', chanMap')
139 putMVar chans (chanId', chanMap')
141−-- This function first extracts a header of type Int64, which determines
141+-- | This function first extracts a header of type Int64, which determines
142-- the size of the ByteString that follows. The ByteString is then
142-- the size of the ByteString that follows. The ByteString is then
143-- extracted from the socket, and then written to the Chan only when
143-- extracted from the socket, and then written to the Chan only when
144−-- complete. A length of -1 indicates that the socket has finished
144+-- complete.
146procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
145procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
147procMessages chans chanId chan sock = do
146procMessages chans chanId chan sock = do
148 sizeBS <- recvExact sock 8
147 sizeBS <- recvExact sock 8
149− let size = B.decode sizeBS
148+ if BS.null sizeBS
152 (chanId', chanMap) <- takeMVar chans
150 (chanId', chanMap) <- takeMVar chans
153 case IntMap.lookup chanId chanMap of
151 case IntMap.lookup chanId chanMap of
@@ -160,6 +158,7 @@ procMessages chans chanId chan sock = do
160 putMVar chans (chanId', chanMap')
158 putMVar chans (chanId', chanMap')
161 sClose sock
159 sClose sock
161+ let size = B.decode sizeBS
163 bs <- recvExact sock size
162 bs <- recvExact sock size
164 writeChan chan bs
163 writeChan chan bs
165 procMessages chans chanId chan sock
164 procMessages chans chanId chan sock
@@ -167,11 +166,16 @@ procMessages chans chanId chan sock = do
167-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
166-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
168-- from `sock`. No more bytes than necessary are read from the socket.
167-- from `sock`. No more bytes than necessary are read from the socket.
169-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
168-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
170−-- superfluous input depending on the socket type.
169+-- superfluous input depending on the socket type. Also note that
170+-- if `recv` returns an empty `ByteString` then this means that the socket
171+-- was closed: in this case, we return the empty `ByteString`.
171recvExact :: Socket -> Int64 -> IO ByteString
172recvExact :: Socket -> Int64 -> IO ByteString
172recvExact sock n = go BS.empty sock n
173recvExact sock n = go BS.empty sock n
174 go bs sock 0 = return bs
175 go bs sock 0 = return bs
175 go bs sock n = do
176 go bs sock n = do
176 bs' <- NBS.recv sock n
177 bs' <- NBS.recv sock n
177− go (BS.append bs bs') sock (n - BS.length bs')
180+ else go (BS.append bs bs') sock (n - BS.length bs')