@@ -12,6 +12,7 @@ import Control.Monad (forever, forM_)
12import Data.ByteString.Lazy.Char8 (ByteString)
12import Data.ByteString.Lazy.Char8 (ByteString)
13import Data.IntMap (IntMap)
13import Data.IntMap (IntMap)
14import Data.Int
14import Data.Int
15import Network.Socket
16import Network.Socket
16 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
17 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
17 , SocketType (Stream), SocketOption (ReuseAddr)
18 , SocketType (Stream), SocketOption (ReuseAddr)
@@ -99,7 +100,12 @@ mkSourceEnd host service chanId = withSocketsDo $ do
99 return SourceEnd
100 return SourceEnd
100 { send = \bss -> do
101 { send = \bss -> do
101 let size = fromIntegral (sum . map BS.length $ bss) :: Int64
102 let size = fromIntegral (sum . map BS.length $ bss) :: Int64
103+ if size <= fromIntegral (maxBound :: Word8)
105+ NBS.sendAll sock (B.encode (fromIntegral size :: Word8))
107+ NBS.sendAll sock (B.encode (0 :: Word8))
102− NBS.sendAll sock (B.encode size)
108+ NBS.sendAll sock (B.encode size)
103 mapM_ (NBS.sendAll sock) bss
109 mapM_ (NBS.sendAll sock) bss
104 , closeSourceEnd = sClose sock
110 , closeSourceEnd = sClose sock
@@ -122,60 +128,88 @@ mkTargetEnd chans chanId chan = TargetEnd
122 putMVar chans (chanId', chanMap')
128 putMVar chans (chanId', chanMap')
131+-- | This function waits for inbound connections. If a connection fails
132+-- for some reason, an error is raised.
125procConnections :: Chans -> Socket -> IO ()
133procConnections :: Chans -> Socket -> IO ()
126procConnections chans sock = forever $ do
134procConnections chans sock = forever $ do
127 (clientSock, _clientAddr) <- accept sock
135 (clientSock, _clientAddr) <- accept sock
128 -- decode the first message to find the correct chanId
136 -- decode the first message to find the correct chanId
129− bs <- recvExact clientSock 8
137+ mBs <- recvExact clientSock 8
130− let chanId = fromIntegral (B.decode bs :: Int64)
138+ case mBs of
131− (chanId', chanMap) <- takeMVar chans
132− case IntMap.lookup chanId chanMap of
134− putMVar chans (chanId', chanMap)
135− error "procConnections: cannot find chanId"
139+ Nothing -> error "procConnections: inbound chanId aborted"
136− Just (chan, socks) -> do
140+ Just bs -> do
137− threadId <- forkIO $ procMessages chans chanId chan clientSock
141+ let chanId = fromIntegral (B.decode bs :: Int64)
138− let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
139− putMVar chans (chanId', chanMap')
141−-- | This function first extracts a header of type Int64, which determines
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
145−procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
146−procMessages chans chanId chan sock = do
147− sizeBS <- recvExact sock 8
150 (chanId', chanMap) <- takeMVar chans
142 (chanId', chanMap) <- takeMVar chans
151 case IntMap.lookup chanId chanMap of
143 case IntMap.lookup chanId chanMap of
152− Nothing -> do
144+ Nothing -> do
153 putMVar chans (chanId', chanMap)
145 putMVar chans (chanId', chanMap)
154− error "procMessages: chanId not found."
146+ error "procConnections: cannot find chanId"
155 Just (chan, socks) -> do
147 Just (chan, socks) -> do
156− let socks' = filter ((/= sock) . snd) socks
148+ threadId <- forkIO $ procMessages chans chanId chan clientSock
157− let chanMap' = IntMap.insert chanId (chan, socks') chanMap
149+ let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
158 putMVar chans (chanId', chanMap')
150 putMVar chans (chanId', chanMap')
161− let size = B.decode sizeBS
162− bs <- recvExact sock size
164− procMessages chans chanId chan sock
152+-- | This function first extracts a header of type Word8, which determines
153+-- the size of the ByteString that follows. If this size is 0, this indicates
154+-- that the ByteString is large, so the next value is an Int64, which
155+-- determines the size of the ByteString that follows. The ByteString is then
156+-- extracted from the socket, and then written to the Chan only when
157+-- complete. If either of the first header size is null this indicates the
158+-- socket has closed.
159+procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
160+procMessages chans chanId chan sock = do
161+ mSizeBS <- recvExact sock 1
163+ Nothing -> closeSocket
165+ let size = fromIntegral (B.decode sizeBS :: Word8)
168+ mSizeBS' <- recvExact sock 8
170+ Nothing -> closeSocket
171+ Just sizeBS' -> procMessage (B.decode sizeBS' :: Int64)
172+ else procMessage size
174+ closeSocket :: IO ()
176+ (chanId', chanMap) <- takeMVar chans
177+ case IntMap.lookup chanId chanMap of
179+ putMVar chans (chanId', chanMap)
180+ error "procMessages: chanId not found."
181+ Just (chan, socks) -> do
182+ let socks' = filter ((/= sock) . snd) socks
183+ let chanMap' = IntMap.insert chanId (chan, socks') chanMap
184+ putMVar chans (chanId', chanMap')
186+ procMessage :: Int64 -> IO ()
187+ procMessage size = do
188+ mBs <- recvExact sock size
190+ Nothing -> closeSocket
193+ procMessages chans chanId chan sock
166−-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
195+-- | The result of `recvExact sock n` is a `Maybe ByteString` of length `n`,
167−-- from `sock`. No more bytes than necessary are read from the socket.
196+-- received from `sock`. No more bytes than necessary are read from the socket.
168-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
197-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
169-- superfluous input depending on the socket type. Also note that
198-- superfluous input depending on the socket type. Also note that
170-- if `recv` returns an empty `ByteString` then this means that the socket
199-- if `recv` returns an empty `ByteString` then this means that the socket
171-- was closed: in this case, we return the empty `ByteString`.
200-- was closed: in this case, we return the empty `ByteString`.
201+-- NB: It may be appropriate to change the return type to
202+-- IO (Either ByteString ByteString), where `return Left bs` indicates
203+-- a partial retrieval since the socket was closed, and `return Right bs`
204+-- indicates success. This hasn't been implemented, since a Transport `receive`
205+-- represents an atomic receipt of a message.
172−recvExact :: Socket -> Int64 -> IO ByteString
206+recvExact :: Socket -> Int64 -> IO (Maybe ByteString)
173recvExact sock n = go BS.empty sock n
207recvExact sock n = go BS.empty sock n
175− go bs sock 0 = return bs
209+ go bs sock 0 = return (Just bs)
176 go bs sock n = do
210 go bs sock n = do
177 bs' <- NBS.recv sock n
211 bs' <- NBS.recv sock n
178 if BS.null bs'
212 if BS.null bs'
179− then return bs'
213+ then return Nothing
180 else go (BS.append bs bs') sock (n - BS.length bs')
214 else go (BS.append bs bs') sock (n - BS.length bs')