@@ -12,6 +12,7 @@ import Control.Monad (forever, forM_)
1212import Data.ByteString.Lazy.Char8 (ByteString)
1313import Data.IntMap (IntMap)
1414import Data.Int
15+import Data.Word
1516import Network.Socket
1617 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
1718 , SocketType (Stream), SocketOption (ReuseAddr)
@@ -99,7 +100,12 @@ mkSourceEnd host service chanId = withSocketsDo $ do
99100 return SourceEnd
100101 { send = \bss -> do
101102 let size = fromIntegral (sum . map BS.length $ bss) :: Int64
102− NBS.sendAll sock (B.encode size)
103+ if size <= fromIntegral (maxBound :: Word8)
104+ then
105+ NBS.sendAll sock (B.encode (fromIntegral size :: Word8))
106+ else do
107+ NBS.sendAll sock (B.encode (0 :: Word8))
108+ NBS.sendAll sock (B.encode size)
103109 mapM_ (NBS.sendAll sock) bss
104110 , closeSourceEnd = sClose sock
105111 }
@@ -122,60 +128,88 @@ mkTargetEnd chans chanId chan = TargetEnd
122128 putMVar chans (chanId', chanMap')
123129 }
124130
131+-- | This function waits for inbound connections. If a connection fails
132+-- for some reason, an error is raised.
125133procConnections :: Chans -> Socket -> IO ()
126134procConnections chans sock = forever $ do
127135 (clientSock, _clientAddr) <- accept sock
128136 -- decode the first message to find the correct chanId
129− bs <- recvExact clientSock 8
130− let chanId = fromIntegral (B.decode bs :: Int64)
131− (chanId', chanMap) <- takeMVar chans
132− case IntMap.lookup chanId chanMap of
133− Nothing -> do
134− putMVar chans (chanId', chanMap)
135− error "procConnections: cannot find chanId"
136− Just (chan, socks) -> do
137− threadId <- forkIO $ procMessages chans chanId chan clientSock
138− let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
139− putMVar chans (chanId', chanMap')
140−
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
144−-- complete.
145−procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
146−procMessages chans chanId chan sock = do
147− sizeBS <- recvExact sock 8
148− if BS.null sizeBS
149− then do
137+ mBs <- recvExact clientSock 8
138+ case mBs of
139+ Nothing -> error "procConnections: inbound chanId aborted"
140+ Just bs -> do
141+ let chanId = fromIntegral (B.decode bs :: Int64)
150142 (chanId', chanMap) <- takeMVar chans
151143 case IntMap.lookup chanId chanMap of
152− Nothing -> do
144+ Nothing -> do
153145 putMVar chans (chanId', chanMap)
154− error "procMessages: chanId not found."
146+ error "procConnections: cannot find chanId"
155147 Just (chan, socks) -> do
156− let socks' = filter ((/= sock) . snd) socks
157− let chanMap' = IntMap.insert chanId (chan, socks') chanMap
148+ threadId <- forkIO $ procMessages chans chanId chan clientSock
149+ let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
158150 putMVar chans (chanId', chanMap')
159− sClose sock
160− else do
161− let size = B.decode sizeBS
162− bs <- recvExact sock size
163− writeChan chan bs
164− procMessages chans chanId chan sock
165151
166−-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
167−-- from `sock`. No more bytes than necessary are read from the socket.
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
162+ case mSizeBS of
163+ Nothing -> closeSocket
164+ Just sizeBS -> do
165+ let size = fromIntegral (B.decode sizeBS :: Word8)
166+ if size == 0
167+ then do
168+ mSizeBS' <- recvExact sock 8
169+ case mSizeBS' of
170+ Nothing -> closeSocket
171+ Just sizeBS' -> procMessage (B.decode sizeBS' :: Int64)
172+ else procMessage size
173+ where
174+ closeSocket :: IO ()
175+ closeSocket = do
176+ (chanId', chanMap) <- takeMVar chans
177+ case IntMap.lookup chanId chanMap of
178+ Nothing -> do
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')
185+ sClose sock
186+ procMessage :: Int64 -> IO ()
187+ procMessage size = do
188+ mBs <- recvExact sock size
189+ case mBs of
190+ Nothing -> closeSocket
191+ Just bs -> do
192+ writeChan chan bs
193+ procMessages chans chanId chan sock
194+
195+-- | The result of `recvExact sock n` is a `Maybe ByteString` of length `n`,
196+-- received from `sock`. No more bytes than necessary are read from the socket.
168197-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
169198-- superfluous input depending on the socket type. Also note that
170199-- if `recv` returns an empty `ByteString` then this means that the socket
171200-- was closed: in this case, we return the empty `ByteString`.
172−recvExact :: Socket -> Int64 -> IO 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.
206+recvExact :: Socket -> Int64 -> IO (Maybe ByteString)
173207recvExact sock n = go BS.empty sock n
174208 where
175− go bs sock 0 = return bs
209+ go bs sock 0 = return (Just bs)
176210 go bs sock n = do
177211 bs' <- NBS.recv sock n
178212 if BS.null bs'
179− then return bs'
213+ then return Nothing
180214 else go (BS.append bs bs') sock (n - BS.length bs')
181215