@@ -35,9 +35,11 @@ import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
3535import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, readMVar)
3636import Control.Category ((>>>))
3737import Control.Applicative ((<*>), (*>), (<$>), pure)
38−import Control.Monad (forever)
38+import Control.Monad (forever, forM_)
3939import Control.Monad.Error (ErrorT, liftIO, runErrorT)
4040import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT)
41+import Control.Monad.Trans.Writer (WriterT, execWriterT)
42+import Control.Monad.Writer.Class (tell)
4143import Data.ByteString (ByteString)
4244import qualified Data.ByteString as BS (concat)
4345import qualified Data.ByteString.Char8 as BSC (pack, unpack)
@@ -72,7 +74,7 @@ endPointAt ix = endPoints >>> intMapLens ix
7274-- * Deal with all exceptions that may occur
7375createTransport :: N.HostName -> N.ServiceName -> IO Transport
7476createTransport host port = do
75− state <- newMVar $ TransportState { _endPoints = IntMap.empty }
77+ state <- newMVar TransportState { _endPoints = IntMap.empty }
7678 forkServer host port (transportServer state)
7779 return Transport { newEndPoint = tcpNewEndPoint state host port }
7880
@@ -89,38 +91,49 @@ transportServer state sock = forever $ do
8991-- Returns when the client closes the socket and on invalid input.
9092handleClient :: MVar TransportState -> N.Socket -> IO ()
9193handleClient state sock = do
92− runMaybeT $ do
94+ mendpoint <- runMaybeT $ do
9395 -- Our endpoint the other side is interested in
9496 endPointBS <- recvExact sock 2
95− endPointIx <- decodeInt16 (BS.concat $ endPointBS)
97+ endPointIx <- decodeInt16 (BS.concat endPointBS)
9698 -- Full endpoint address of the other side
9799 theirAddress <- BS.concat <$> recvWithLength sock
98− -- Start handling incoming messages
100+ -- Find our endpoint
99101 (endPointCh, endPointSt) <- MaybeT $ (^. endPointAt (fromIntegral endPointIx)) <$> readMVar state
100− endpointServer sock (EndPointAddress theirAddress) endPointCh endPointSt
101− return ()
102+ return (endPointCh, endPointSt, theirAddress)
103+ case mendpoint of
104+ Just (endPointCh, endPointSt, theirAddress) -> do
105+ openConnections <- endpointServer sock (EndPointAddress theirAddress) endPointCh endPointSt
106+ forM_ openConnections $ \cix -> writeChan endPointCh (ConnectionClosed cix)
107+ Nothing ->
108+ -- Invalid request
109+ return ()
102110
103−-- | Handle requests for one endpoint. Returns only if an error occurs
104−endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> MaybeT IO ()
105−endpointServer sock theirAddress chan state = forever $ do
111+-- | Handle requests for one endpoint. Returns only if an error occurs,
112+-- in which case it returns the connections that are still open (normally,
113+-- this should be an empty list).
114+endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> IO [ConnectionIx]
115+endpointServer sock theirAddress chan state = execWriterT . runMaybeT . forever $ do
106116 connBS <- recvExact sock 2
107117 connIx <- fromIntegral <$> decodeInt16 (BS.concat connBS)
108118 if connIx == 0
109119 then createNewConnection
110120 else readMessage connIx
111121 where
122+ -- A brief note on the types:
123+ -- These functions may fail (connection error): hence MaybeT
124+ -- We also need to keep track of which connections are still open: hence WriterT
125+
112126 -- Create a new connection
113− createNewConnection :: MaybeT IO ()
127+ createNewConnection :: MaybeT (WriterT [ConnectionIx] IO) ()
114128 createNewConnection = do
115− (newIx, newBs) <- liftIO $ do
116− newIx <- modifyMVar state $ \st -> return (nextConnectionIx ^+= 1 $ st, st ^. nextConnectionIx)
117− newBs <- encodeInt16 (fromIntegral newIx)
118− return (newIx, newBs)
129+ newIx <- liftIO $ modifyMVar state $ \st -> return (nextConnectionIx ^+= 1 $ st, st ^. nextConnectionIx)
130+ newBs <- encodeInt16 (fromIntegral newIx)
119131 sendMany sock [newBs]
120132 liftIO $ writeChan chan (ConnectionOpened newIx ReliableOrdered theirAddress)
133+ tell [newIx]
121134
122135 -- Read a message and output it on the endPoint's channel
123− readMessage :: ConnectionIx -> MaybeT IO ()
136+ readMessage :: ConnectionIx -> MaybeT (WriterT [ConnectionIx] IO) ()
124137 readMessage connIx = do
125138 payload <- recvWithLength sock
126139 liftIO $ writeChan chan (Received connIx payload)
@@ -129,7 +142,7 @@ endpointServer sock theirAddress chan state = forever $ do
129142tcpNewEndPoint :: MVar TransportState -> N.HostName -> N.ServiceName -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
130143tcpNewEndPoint transportState host port = do
131144 chan <- newChan
132− endPointState <- newMVar $ EndPointState { _nextConnectionIx = 1 } -- Connection id 0 is reserved for requesting new connections
145+ endPointState <- newMVar EndPointState { _nextConnectionIx = 1 } -- Connection id 0 is reserved for requesting new connections
133146 endPointIx <- modifyMVar transportState $ \st -> do let ix = IntMap.size (st ^. endPoints)
134147 return ((endPointAt ix ^= Just (chan, endPointState)) $ st, ix)
135148 let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show endPointIx
@@ -149,15 +162,16 @@ tcpConnect :: EndPointAddress -> EndPointAddress -> Reliability -> IO (Either (F
149162tcpConnect myAddress theirAddress _ = runErrorT $ do
150163 sock <- socketForEndPoint myAddress theirAddress
151164 connIx <- maybeTToErrorT undefined $ requestNewConnection sock
152− return $ Connection { send = \msg -> runErrorT $ maybeTToErrorT (FailedWith SendFailed "Send failed") $
153− sendWithLength sock (Just connIx) msg
154− , close = N.sClose sock
155− }
165+ return Connection { send = runErrorT .
166+ maybeTToErrorT (FailedWith SendFailed "Send failed") .
167+ sendWithLength sock (Just connIx)
168+ , close = N.sClose sock
169+ }
156170
157171-- | Request a new connection
158172requestNewConnection :: N.Socket -> MaybeT IO ByteString
159173requestNewConnection sock = do
160− (liftIO $ pure <$> encodeInt16 0) >>= sendMany sock
174+ liftIO (pure <$> encodeInt16 0) >>= sendMany sock
161175 BS.concat <$> recvExact sock 2
162176
163177-- | Find a socket to the specified endpoint