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