@@ -5,23 +5,29 @@ module Network.Transport.TCP
55
66import Network.Transport
77
8−import Control.Concurrent (forkIO)
8+import Control.Concurrent (forkIO, ThreadId, killThread)
99import Control.Concurrent.Chan
1010import Control.Concurrent.MVar
11−import Control.Monad (forever, unless)
12−import Data.ByteString.Char8 (ByteString)
11+import Control.Monad (forever, forM_)
12+import Data.ByteString.Lazy.Char8 (ByteString)
1313import Data.IntMap (IntMap)
14+import Data.Int
1415import Data.Word
1516import Network.Socket
17+ ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
18+ , SocketType (Stream), SocketOption (ReuseAddr)
19+ , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
20+ , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
1621import Safe
1722
18−import qualified Data.ByteString.Char8 as BS
23+import qualified Data.Binary as B
24+import qualified Data.ByteString.Lazy.Char8 as BS
1925import qualified Data.IntMap as IntMap
2026import qualified Network.Socket as N
21−import qualified Network.Socket.ByteString as NBS
27+import qualified Network.Socket.ByteString.Lazy as NBS
2228
23−type Chans = MVar (Int, IntMap (Chan ByteString))
2429type ChanId = Int
30+type Chans = MVar (ChanId, IntMap (Chan [ByteString], [(ThreadId, Socket)]))
2531
2632-- | This deals with several different configuration properties:
2733-- * Buffer size, specified in Hints
@@ -37,7 +43,7 @@ data TCPConfig = TCPConfig Hints HostName ServiceName
3743-- Messages are all queued using an unbounded Chan.
3844mkTransport :: TCPConfig -> IO Transport
3945mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
40− channels <- newMVar (0, IntMap.empty)
46+ chans <- newMVar (0, IntMap.empty)
4147 serverAddrs <- getAddrInfo
4248 (Just (N.defaultHints { addrFlags = [AI_PASSIVE] }))
4349 Nothing
@@ -46,69 +52,163 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
4652 [] -> error "mkTransport: getAddrInfo returned []"
4753 as -> head as
4854 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
55+ setSocketOption sock ReuseAddr 1
4956 bindSocket sock (addrAddress serverAddr)
5057 listen sock 5
51− forkIO $ procConnections channels sock
58+ threadId <- forkIO $ procConnections chans sock
5259
5360 return Transport
54− { newConnectionWith = \_ -> do
55− (chanId, chanMap) <- takeMVar channels
61+ { newConnectionWith = {-# SCC "newConnectionWith" #-}\_ -> do
62+ (chanId, chanMap) <- takeMVar chans
5663 chan <- newChan
57− putMVar channels (chanId + 1, IntMap.insert chanId chan chanMap)
58− return (mkSourceAddr host service chanId, mkTargetEnd chan)
59− , newMulticastWith = undefined
60− , deserialize = \bs ->
61− case readMay . BS.unpack $ bs of
62− Nothing -> error "deserialize: cannot parse"
63− Just (host, service, chanId) -> Just $ mkSourceAddr host service chanId
64+ putMVar chans (chanId + 1, IntMap.insert chanId (chan, []) chanMap)
65+ return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
66+ , newMulticastWith = error "newMulticastWith: not defined"
67+ , deserialize = {-# SCC "deserialize" #-} \bs ->
68+ let (host, service, chanId) = B.decode bs in
69+ Just $ mkSourceAddr host service chanId
70+ , closeTransport = {-# SCC "closeTransport" #-} do
71+ -- Kill the transport channel process
72+ killThread threadId
73+ sClose sock
74+ -- Kill all target end processes
75+ (chanId, chanMap) <- takeMVar chans
76+ forM_ (IntMap.elems chanMap) (\(chan, socks) ->
77+ forM_ socks (\(threadId', sock') -> do
78+ killThread threadId'
79+ sClose sock'))
6480 }
6581
66− where
67− mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
68− mkSourceAddr host service chanId = SourceAddr
69− { connectWith = \_ -> mkSourceEnd host service chanId
70− , serialize = BS.pack . show $ (host, service, chanId)
71− }
82+mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
83+mkSourceAddr host service chanId = SourceAddr
84+ { connectWith = {-# SCC "connectWith" #-} \_ -> mkSourceEnd host service chanId
85+ , serialize = {-# SCC "serialize" #-} B.encode (host, service, chanId)
86+ }
7287
73− mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
74− mkSourceEnd host service chanId = withSocketsDo $ do
75− serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
76− let serverAddr = case serverAddrs of
77− [] -> error "mkSourceEnd: getAddrInfo returned []"
78− as -> head as
79− sock <- socket (addrFamily serverAddr) Stream defaultProtocol
80− setSocketOption sock ReuseAddr 1
81− N.connect sock (addrAddress serverAddr)
82− NBS.sendAll sock $ BS.pack . show $ chanId
83− return $ SourceEnd
84− { Network.Transport.send = \bss -> NBS.sendMany sock bss
85− }
88+mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
89+mkSourceEnd host service chanId = withSocketsDo $ do
90+ serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
91+ let serverAddr = case serverAddrs of
92+ [] -> error "mkSourceEnd: getAddrInfo returned []"
93+ as -> head as
94+ sock <- socket (addrFamily serverAddr) Stream defaultProtocol
95+ setSocketOption sock ReuseAddr 1
96+ N.connect sock (addrAddress serverAddr)
97+ NBS.sendAll sock $ B.encode (fromIntegral chanId :: Int64)
98+ return SourceEnd
99+ { send = {-# SCC "send" #-} \bss -> do
100+ let size = fromIntegral (sum . map BS.length $ bss) :: Int64
101+ if size < 255
102+ then
103+ NBS.sendAll sock (B.encode (fromIntegral size :: Word8))
104+ else do
105+ NBS.sendAll sock (B.encode (255 :: Word8))
106+ NBS.sendAll sock (B.encode size)
107+ mapM_ (NBS.sendAll sock) bss
108+ , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
109+ }
110+
111+mkTargetEnd :: Chans -> ChanId -> Chan [ByteString] -> TargetEnd
112+mkTargetEnd chans chanId chan = TargetEnd
113+ { -- for now we will implement this as a Chan
114+ receive = {-# SCC "receive" #-} readChan chan
115+ , closeTargetEnd = {-# SCC "closeTargetEnd" #-} do
116+ (chanId', chanMap) <- takeMVar chans
117+ case IntMap.lookup chanId chanMap of
118+ Nothing -> putMVar chans (chanId', chanMap)
119+ Just (_, socks) -> do
120+ forM_ socks $ \(threadId, sock) -> do
121+ killThread threadId
122+ sClose sock
123+ let chanMap' = IntMap.delete chanId chanMap
124+ putMVar chans (chanId', chanMap')
125+ }
126+
127+-- | This function waits for inbound connections. If a connection fails
128+-- for some reason, an error is raised.
129+procConnections :: Chans -> Socket -> IO ()
130+procConnections chans sock = forever $ do
131+ (clientSock, _clientAddr) <- accept sock
132+ -- decode the first message to find the correct chanId
133+ bss <- recvExact clientSock 8
134+ case bss of
135+ [] -> error "procConnections: inbound chanId aborted"
136+ bss -> do
137+ let bs = BS.concat bss
138+ let chanId = fromIntegral (B.decode bs :: Int64)
139+ (chanId', chanMap) <- takeMVar chans
140+ case IntMap.lookup chanId chanMap of
141+ Nothing -> do
142+ putMVar chans (chanId', chanMap)
143+ error "procConnections: cannot find chanId"
144+ Just (chan, socks) -> do
145+ threadId <- forkIO $ procMessages chans chanId chan clientSock
146+ let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
147+ putMVar chans (chanId', chanMap')
86148
87− mkTargetEnd :: Chan ByteString -> TargetEnd
88− mkTargetEnd chan = TargetEnd
89− { -- for now we will implement this as a Chan
90− receive = do
91− bs <- readChan chan
92− return [bs]
93− }
149+-- | This function first extracts a header of type Word8, which determines
150+-- the size of the ByteString that follows. If this size is 0, this indicates
151+-- that the ByteString is large, so the next value is an Int64, which
152+-- determines the size of the ByteString that follows. The ByteString is then
153+-- extracted from the socket, and then written to the Chan only when
154+-- complete. If either of the first header size is null this indicates the
155+-- socket has closed.
156+procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
157+procMessages chans chanId chan sock = do
158+ sizeBSs <- recvExact sock 1
159+ case sizeBSs of
160+ [] -> closeSocket
161+ _ -> do
162+ let size = fromIntegral (B.decode . BS.concat $ sizeBSs :: Word8)
163+ if size == 255
164+ then do
165+ sizeBSs' <- recvExact sock 8
166+ case sizeBSs' of
167+ [] -> closeSocket
168+ _ -> procMessage (B.decode . BS.concat $ sizeBSs' :: Int64)
169+ else procMessage size
170+ where
171+ closeSocket :: IO ()
172+ closeSocket = do
173+ (chanId', chanMap) <- takeMVar chans
174+ case IntMap.lookup chanId chanMap of
175+ Nothing -> do
176+ putMVar chans (chanId', chanMap)
177+ error "procMessages: chanId not found."
178+ Just (chan, socks) -> do
179+ let socks' = filter ((/= sock) . snd) socks
180+ let chanMap' = IntMap.insert chanId (chan, socks') chanMap
181+ putMVar chans (chanId', chanMap')
182+ sClose sock
183+ procMessage :: Int64 -> IO ()
184+ procMessage size = do
185+ bss <- recvExact sock size
186+ case bss of
187+ [] -> closeSocket
188+ _ -> do
189+ writeChan chan bss
190+ procMessages chans chanId chan sock
94191
95− procConnections :: Chans -> Socket -> IO ()
96− procConnections chans sock = forever $ do
97− (clientSock, _clientAddr) <- accept sock
98− -- decode the first message to find the correct chanId
99− bs <- NBS.recv clientSock 4096
100− case BS.readInt bs of
101− Nothing -> error "procConnections: cannot parse chanId"
102− Just (chanId, bs') -> do
103− -- lookup the channel
104− (_, chanMap) <- readMVar chans
105− case IntMap.lookup chanId chanMap of
106− Nothing -> error "procConnections: cannot find chanId"
107− Just chan -> forkIO $ do
108− unless (BS.null bs') $ writeChan chan bs'
109− procMessages chan clientSock
192+-- | The result of `recvExact sock n` is a `[ByteString]` whose concatenation
193+-- is of length `n`, received from `sock`. No more bytes than necessary are
194+-- read from the socket.
195+-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
196+-- superfluous input depending on the socket type. Also note that
197+-- if `recv` returns an empty `ByteString` then this means that the socket
198+-- was closed: in this case, we return an empty list.
199+-- NB: It may be appropriate to change the return type to
200+-- IO (Either ByteString ByteString), where `return Left bs` indicates
201+-- a partial retrieval since the socket was closed, and `return Right bs`
202+-- indicates success. This hasn't been implemented, since a Transport `receive`
203+-- represents an atomic receipt of a message.
204+recvExact :: Socket -> Int64 -> IO [ByteString]
205+recvExact sock n = go [] sock n
206+ where
207+ go :: [ByteString] -> Socket -> Int64 -> IO [ByteString]
208+ go bss _ 0 = return (reverse bss)
209+ go bss sock n = do
210+ bs <- NBS.recv sock n
211+ if BS.null bs
212+ then return []
213+ else go (bs:bss) sock (n - BS.length bs)
110214
111− procMessages :: Chan ByteString -> Socket -> IO ()
112− procMessages chan sock = forever $ do
113− bs <- NBS.recv sock 4096
114− writeChan chan bs