@@ -80,90 +80,89 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
80 sClose sock'))
80 sClose sock'))
84− mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
83+mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
85− mkSourceAddr host service chanId = SourceAddr
84+mkSourceAddr host service chanId = SourceAddr
86− { connectWith = \_ -> mkSourceEnd host service chanId
85+ { connectWith = \_ -> mkSourceEnd host service chanId
87− , serialize = B.encode (host, service, chanId)
86+ , serialize = B.encode (host, service, chanId)
90− mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
89+mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
91− mkSourceEnd host service chanId = withSocketsDo $ do
90+mkSourceEnd host service chanId = withSocketsDo $ do
92− serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
91+ serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
93− let serverAddr = case serverAddrs of
92+ let serverAddr = case serverAddrs of
94− [] -> error "mkSourceEnd: getAddrInfo returned []"
93+ [] -> error "mkSourceEnd: getAddrInfo returned []"
95− as -> head as
94+ as -> head as
96− sock <- socket (addrFamily serverAddr) Stream defaultProtocol
95+ sock <- socket (addrFamily serverAddr) Stream defaultProtocol
97− setSocketOption sock ReuseAddr 1
96+ setSocketOption sock ReuseAddr 1
98− N.connect sock (addrAddress serverAddr)
97+ N.connect sock (addrAddress serverAddr)
99− NBS.sendAll sock $ B.encode (fromIntegral chanId :: Int64)
98+ NBS.sendAll sock $ B.encode (fromIntegral chanId :: Int64)
100− return SourceEnd
99+ return SourceEnd
101− { send = \bss -> do
100+ { send = \bss -> do
102− let size = fromIntegral (sum . map BS.length $ bss) :: Int64
101+ let size = fromIntegral (sum . map BS.length $ bss) :: Int64
103− NBS.sendAll sock (B.encode size)
102+ NBS.sendAll sock (B.encode size)
104− mapM_ (NBS.sendAll sock) bss
103+ mapM_ (NBS.sendAll sock) bss
105− , closeSourceEnd = sClose sock
104+ , closeSourceEnd = sClose sock
108− mkTargetEnd :: Chans -> ChanId -> Chan ByteString -> TargetEnd
107+mkTargetEnd :: Chans -> ChanId -> Chan ByteString -> TargetEnd
109− mkTargetEnd chans chanId chan = TargetEnd
108+mkTargetEnd chans chanId chan = TargetEnd
110− { -- for now we will implement this as a Chan
109+ { -- for now we will implement this as a Chan
111− receive = do
110+ receive = do
112− bs <- readChan chan
111+ bs <- readChan chan
113− return [bs]
112+ return [bs]
114− , closeTargetEnd = do
113+ , closeTargetEnd = do
115− (chanId', chanMap) <- takeMVar chans
114+ (chanId', chanMap) <- takeMVar chans
116− case IntMap.lookup chanId chanMap of
115+ case IntMap.lookup chanId chanMap of
117− Nothing -> putMVar chans (chanId', chanMap)
116+ Nothing -> putMVar chans (chanId', chanMap)
118− Just (_, socks) -> do
117+ Just (_, socks) -> do
119− forM_ socks $ \(threadId, sock) -> do
118+ forM_ socks $ \(threadId, sock) -> do
120− killThread threadId
119+ killThread threadId
121− sClose sock
120+ sClose sock
122− let chanMap' = IntMap.delete chanId chanMap
121+ let chanMap' = IntMap.delete chanId chanMap
123− putMVar chans (chanId', chanMap')
122+ putMVar chans (chanId', chanMap')
126− procConnections :: Chans -> Socket -> IO ()
125+procConnections :: Chans -> Socket -> IO ()
127− procConnections chans sock = forever $ do
126+procConnections chans sock = forever $ do
128− (clientSock, _clientAddr) <- accept sock
127+ (clientSock, _clientAddr) <- accept sock
129− -- decode the first message to find the correct chanId
128+ -- decode the first message to find the correct chanId
130− bs <- recvExact clientSock 8
129+ bs <- recvExact clientSock 8
131− let chanId = fromIntegral (B.decode bs :: Int64)
130+ let chanId = fromIntegral (B.decode bs :: Int64)
131+ (chanId', chanMap) <- takeMVar chans
132+ case IntMap.lookup chanId chanMap of
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')
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. A length of -1 indicates that the socket has finished
146+procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
147+procMessages chans chanId chan sock = do
148+ sizeBS <- recvExact sock 8
149+ let size = B.decode sizeBS
132 (chanId', chanMap) <- takeMVar chans
152 (chanId', chanMap) <- takeMVar chans
133 case IntMap.lookup chanId chanMap of
153 case IntMap.lookup chanId chanMap of
134− Nothing -> do
154+ Nothing -> do
135 putMVar chans (chanId', chanMap)
155 putMVar chans (chanId', chanMap)
136− error "procConnections: cannot find chanId"
156+ error "procMessages: chanId not found."
137 Just (chan, socks) -> do
157 Just (chan, socks) -> do
138− threadId <- forkIO $ procMessages chans chanId chan clientSock
158+ let socks' = filter ((/= sock) . snd) socks
139− let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
159+ let chanMap' = IntMap.insert chanId (chan, socks') chanMap
140 putMVar chans (chanId', chanMap')
160 putMVar chans (chanId', chanMap')
142− -- This function first extracts a header of type Int64, which determines
143− -- the size of the ByteString that follows. The ByteString is then
144− -- extracted from the socket, and then written to the Chan only when
145− -- complete. A length of -1 indicates that the socket has finished
147− procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
148− procMessages chans chanId chan sock = do
149− sizeBS <- recvExact sock 8
150− let size = B.decode sizeBS
153− (chanId', chanMap) <- takeMVar chans
154− case IntMap.lookup chanId chanMap of
156− putMVar chans (chanId', chanMap)
157− error "procMessages: chanId not found."
158− Just (chan, socks) -> do
159− let socks' = filter ((/= sock) . snd) socks
160− let chanMap' = IntMap.insert chanId (chan, socks') chanMap
161− putMVar chans (chanId', chanMap')
162− sClose sock
161+ sClose sock
164− bs <- recvExact sock size
163+ bs <- recvExact sock size
165− writeChan chan bs
164+ writeChan chan bs
166− procMessages chans chanId chan sock
165+ procMessages chans chanId chan sock
168-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
167-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
169-- from `sock`. No more bytes than necessary are read from the socket.
168-- from `sock`. No more bytes than necessary are read from the socket.