@@ -1,309 +1,1130 @@
1−{-# LANGUAGE CPP #-}
1+-- | TCP implementation of the transport layer.
3+-- The TCP implementation guarantees that only a single TCP connection (socket)
4+-- will be used between endpoints, provided that the addresses specified are
5+-- canonical. If /A/ connects to /B/ and reports its address as
6+-- @192.168.0.1:8080@ and /B/ subsequently connects tries to connect to /A/ as
7+-- @client1.local:http-alt@ then the transport layer will not realize that the
8+-- TCP connection can be reused.
10+-- Applications that use the TCP transport should use
11+-- 'Network.Socket.withSocketsDo' in their main function for Windows
12+-- compatibility (see "Network.Socket").
15+-- * Output exception on channel after endpoint is closed
3−module Network.Transport.TCP
16+module Network.Transport.TCP ( -- * Main API
4− ( mkTransport
17+ createTransport
5− , TCPConfig (..)
18+ , -- * TCP specific functionality (exported mostly for testing purposes_
20+ , encodeEndPointAddress
21+ , decodeEndPointAddress
23+ , ConnectionRequestResponse(..)
24+ , firstNonReservedConnectionId
30+import Prelude hiding (catch)
8import Network.Transport
31import Network.Transport
9−
32+import Network.Transport.Internal.TCP ( forkServer
10−import Control.Applicative
33+ , recvWithLength
36+import Network.Transport.Internal ( encodeInt32
44+import qualified Network.Socket as N ( HostName
55+ , SocketOption(ReuseAddr)
60+import Network.Socket.ByteString (sendMany)
11import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
61import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
12−import Control.Concurrent.MVar
62+import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
63+import Control.Concurrent.MVar ( MVar
73+import Control.Category ((>>>))
74+import Control.Applicative ((<$>))
75+import Control.Monad (forM_, when, unless)
13−import Control.Exception (SomeException, IOException, AsyncException(ThreadKilled),
76+import Control.Exception (IOException, SomeException, handle, throw, try, bracketOnError)
14− fromException, throwTo, throw, catch, handle)
77+import Data.IORef (IORef, newIORef, writeIORef, readIORef)
15−import Control.Monad (forever, forM_)
78+import Data.ByteString (ByteString)
79+import qualified Data.ByteString as BS (concat)
80+import qualified Data.ByteString.Char8 as BSC (pack, unpack, split)
81+import Data.Int (Int32)
82+import Data.IntMap (IntMap)
83+import qualified Data.IntMap as IntMap (empty)
84+import Data.IntSet (IntSet)
85+import qualified Data.IntSet as IntSet (empty, insert, elems, singleton, null, delete)
87+import qualified Data.Map as Map (empty, elems, size)
88+import Data.Accessor (Accessor, accessor, (^.), (^=), (^:))
89+import qualified Data.Accessor.Container as DAC (mapMaybe, intMapMaybe)
90+import System.IO (hPutStrLn, stderr)
17−import qualified Data.ByteString.Internal as BSI
92+-- $design
96+-- The TCP transport maps multiple logical connections between /A/ and /B/ (in
97+-- either direction) to a single TCP connection:
99+-- > +-------+ +-------+
100+-- > | A |==========================| B |
101+-- > | |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~\ |
102+-- > | Q |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~Q |
103+-- > | \~~~|~~~~~~~~~~~~~~~~~~~~~~~~~<| |
104+-- > | |==========================| |
105+-- > +-------+ +-------+
107+-- Ignoring the complications detailed below, the TCP connection is set up is
108+-- when the first lightweight connection is created (in either direction), and
109+-- torn down when the last lightweight connection (in either direction) is
114+-- Let /A/, /B/ be two endpoints without any connections. When /A/ wants to
115+-- connect to /B/, it locally records that it is trying to connect to /B/ and
116+-- sends a request to /B/. As part of the request /A/ sends its own endpoint
117+-- address to /B/ (so that /B/ can reuse the connection in the other direction).
119+-- When /B/ receives the connection request it first checks if it did not
120+-- already initiate a connection request to /A/. If not it will acknowledge the
121+-- connection request by sending 'ConnectionRequestAccepted' to /A/ and record
122+-- that it has a TCP connection to /A/.
124+-- The tricky case arises when /A/ sends a connection request to /B/ and /B/
125+-- finds that it had already sent a connection request to /A/. In this case /B/
126+-- will accept the connection request from /A/ if /A/s endpoint address is
127+-- smaller (lexicographically) than /B/s, and reject it otherwise. If it rejects
128+-- it, it sends a 'ConnectionRequestCrossed' message to /A/. (The
129+-- lexicographical ordering is an arbitrary but convenient way to break the
132+-- When it receives a 'ConnectionRequestCrossed' message the /A/ thread that
133+-- initiated the request just needs to wait until the /A/ thread that is dealing
134+-- with /B/'s connection request completes.
138+-- The TCP connection is created as soon as the first logical connection from
139+-- /A/ to /B/ (or /B/ to /A/) is established. At this point a thread (@#@) is
140+-- spawned that listens for incoming connections from /B/:
142+-- > +-------+ +-------+
143+-- > | A |==========================| B |
144+-- > | |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~\ |
147+-- > | |==========================| |
148+-- > +-------+ +-------+
150+-- The question is when the TCP connection can be closed again. Conceptually,
151+-- we want to do reference counting: when there are no logical connections left
152+-- between /A/ and /B/ we want to close the socket (possibly after some
155+-- However, /A/ and /B/ need to agree that the refcount has reached zero. It
156+-- might happen that /B/ sends a connection request over the existing socket at
157+-- the same time that /A/ closes its logical connection to /B/ and closes the
158+-- socket. This will cause a failure in /B/ (which will have to retry) which is
159+-- not caused by a network failure, which is unfortunate. (Note that the
160+-- connection request from /B/ might succeed even if /A/ closes the socket.)
162+-- Instead, when /A/ is ready to close the socket it sends a 'CloseSocket'
163+-- request to /B/ and records that its connection to /B/ is closing. If /A/
164+-- receives a new connection request from /B/ after having sent the
165+-- 'CloseSocket' request it simply forgets that it sent a 'CloseSocket' request
166+-- and increments the reference count of the connection again.
168+-- When /B/ receives a 'CloseSocket' message and it too is ready to close the
169+-- connection, it will respond with a reciprocal 'CloseSocket' request to /A/
170+-- and then actually close the socket. /A/ meanwhile will not send any more
171+-- requests to /B/ after having sent a 'CloseSocket' request, and will actually
172+-- close its end of the socket only when receiving the 'CloseSocket' message
173+-- from /B/. (Since /A/ recorded that its connection to /B/ is in closing state
174+-- after sending a 'CloseSocket' request to /B/, it knows not to reciprocate /B/
175+-- reciprocal 'CloseSocket' message.)
177+-- If there is a concurrent thread in /A/ waiting to connect to /B/ after /A/
178+-- has sent a 'CloseSocket' request then this thread will block until /A/ knows
179+-- whether to reuse the old socket (if /B/ sends a new connection request
180+-- instead of acknowledging the 'CloseSocket') or to set up a new socket.
19−import Data.IntMap (IntMap)
182+--------------------------------------------------------------------------------
20−import qualified Data.IntMap as IntMap
183+-- Internal datatypes --
184+--------------------------------------------------------------------------------
22−import Data.Int (Int32)
186+-- We use underscores for fields that we might update (using accessores)
23−import qualified Data.Serialize as Ser
187+--
24−import Data.Word (Word8)
188+-- All data types follow the same structure:
26−import Foreign.ForeignPtr (withForeignPtr)
190+-- * A top-level data type describing static properties (TCPTransport,
27−import Foreign.Ptr (Ptr, plusPtr)
191+-- LocalEndPoint, RemoteEndPoint)
28−import Foreign.Storable (sizeOf)
192+-- * The 'static' properties include an MVar containing a data structure for
29−
193+-- the dynamic properties (TransportState, LocalEndPointState,
30−import Network.Socket
194+-- RemoteEndPointState). The state could be invalid/valid/closed,/etc.
31− ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
195+-- * For the case of "valid" we use third data structure to give more details
32− , SocketType (Stream), SocketOption (ReuseAddr)
196+-- about the state (ValidTransportState, ValidLocalEndPointState,
33− , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
197+-- ValidRemoteEndPointState).
34− , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
35−import qualified Network.Socket as N
40−import System.IO (stderr, hPutStrLn)
41−import Foreign.Storable (pokeByteOff, peekByteOff)
42−import Foreign.C (CInt(..))
43−import Foreign.ForeignPtr (withForeignPtr)
46−import Data.ByteString.Char8 (ByteString)
47−import qualified Data.ByteString.Char8 as BS
48−import qualified Network.Socket.ByteString as NBS
49−sendNBS = NBS.sendMany
53−import Data.ByteString.Lazy.Char8 (ByteString)
54−import qualified Data.ByteString.Lazy.Char8 as BS
55−import qualified Network.Socket.ByteString.Lazy as NBS
56−sendNBS sock = NBS.sendAll sock . BS.concat
57−encode = Ser.encodeLazy
58−decode = Ser.decodeLazy
62−encode :: Ser.Serialize a => a -> ByteString
63−decode :: Ser.Serialize a => ByteString -> Either String a
65−foreign import ccall unsafe "htonl" htonl :: CInt -> CInt
66−foreign import ccall unsafe "ntohl" ntohl :: CInt -> CInt
69−type Chans = MVar (ChanId, IntMap (MVar [ByteString], [(ThreadId, Socket)]))
71−-- | This deals with several different configuration properties:
72−-- * Buffer size, specified in Hints
73−-- * LAN/WAN, since we can inspect the addresses
74−-- Note that HostName could be an IP address, and ServiceName could be
76−data TCPConfig = TCPConfig Hints HostName ServiceName
78−-- | This creates a TCP connection between a server and a number of
79−-- clients. Behind the scenes, the server hostname is passed as the SourceAddr
80−-- and when a connection is made, messages sent down the SourceEnd go
81−-- via a socket made for the client that connected.
82−-- Messages are all queued using an unbounded Chan.
83−mkTransport :: TCPConfig -> IO Transport
84−mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
85− chans <- newMVar (0, IntMap.empty)
86− serverAddrs <- getAddrInfo
87− (Just (N.defaultHints { addrFlags = [AI_PASSIVE] }))
90− let serverAddr = case serverAddrs of
91− [] -> error "mkTransport: getAddrInfo returned []"
93− sock <- socket (addrFamily serverAddr) Stream defaultProtocol
94− setSocketOption sock ReuseAddr 1
95− bindSocket sock (addrAddress serverAddr)
97− threadId <- forkWithExceptions forkIO "Connection Listener" $
98− procConnections chans sock
101− { newConnectionWith = {-# SCC "newConnectionWith" #-}\_ -> do
102− (chanId, chanMap) <- takeMVar chans
103− chan <- newEmptyMVar
104− putMVar chans (chanId + 1, IntMap.insert chanId (chan, []) chanMap)
105− return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
106− , newMulticastWith = error "newMulticastWith: not defined"
107− , deserialize = {-# SCC "deserialize" #-} \bs ->
108− let (host, service, chanId) =
109− either (\m -> error $ printf "deserialize: %s" m)
111− Just $ mkSourceAddr host service chanId
112− , closeTransport = {-# SCC "closeTransport" #-} do
113− -- Kill the transport channel process
116− -- Kill all target end processes
117− (chanId, chanMap) <- takeMVar chans
118− forM_ (IntMap.elems chanMap) (\(chan, socks) ->
119− forM_ socks (\(threadId', sock') -> do
120− killThread threadId'
124−mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
199+data TCPTransport = TCPTransport
125−mkSourceAddr host service chanId = SourceAddr
200+ { transportHost :: N.HostName
126− { connectWith = {-# SCC "connectWith" #-} \_ -> mkSourceEnd host service chanId
201+ , transportPort :: N.ServiceName
127− , serialize = {-# SCC "serialize" #-} encode (host, service, chanId)
202+ , transportState :: MVar TransportState
130−mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
205+data TransportState =
131−mkSourceEnd host service chanId = withSocketsDo $ do
206+ TransportValid ValidTransportState
132− let err m = error $ printf "mkSourceEnd: %s" m
207+ | TransportClosed
133− serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
208+
134− let serverAddr = case serverAddrs of
209+data ValidTransportState = ValidTransportState
135− [] -> err "getAddrInfo returned []"
210+ { _localEndPoints :: Map EndPointAddress LocalEndPoint
137− sock <- socket (addrFamily serverAddr) Stream defaultProtocol
212+
138− setSocketOption sock ReuseAddr 1
213+data LocalEndPoint = LocalEndPoint
139− N.connect sock (addrAddress serverAddr)
214+ { localAddress :: EndPointAddress
140− NBS.sendAll sock $ encode (fromIntegral chanId :: Int32)
215+ , localChannel :: Chan Event
141− return SourceEnd
216+ , localState :: MVar LocalEndPointState
142− { send = {-# SCC "send" #-} \bss ->
217+ }
143− let size = fromIntegral (sum . map BS.length $ bss) :: Int32 in
218+
144− if size < 255
219+data LocalEndPointState =
145− then sendNBS sock (BS.singleton (toEnum $ fromEnum size) : bss)
220+ LocalEndPointValid ValidLocalEndPointState
146− else do size' <- encodeLength size
221+ | LocalEndPointClosed
147− sendNBS sock (BS.singleton (toEnum 255) : size' : bss)
222+
148− , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
223+data ValidLocalEndPointState = ValidLocalEndPointState
149− }
224+ { _nextConnectionId :: ConnectionId
225+ , _pendingCtrlRequests :: IntMap (MVar [ByteString])
226+ , _nextCtrlRequestId :: ControlRequestId
227+ , _localConnections :: Map EndPointAddress RemoteEndPoint
228+ , _internalThreads :: [ThreadId]
231+-- A remote endpoint has incoming and outgoing connections, and when the total
232+-- number of connections (that is, the 'remoteRefCount') drops to zero we want
233+-- to close the TCP connection to the endpoint.
235+-- What we need to avoid, however, is a situation with two concurrent threads
236+-- where one closes the last (incoming or outgoing) connection, initiating the
237+-- process of closing the connection, while another requests (but does not yet
238+-- have) a new connection.
240+-- We therefore insist that:
242+-- 1. All operations that change the state of the endpoint (ask for a new
243+-- connection, close a connection, close the endpoint completely) are
244+-- serialized (that is, they take the contents of the MVar containing the
245+-- endpoint state before starting and don't put the updated contents back
246+-- until they have completed).
247+-- 2. Writing to ('apiSend') or reading from (in 'handleIncomingMessages') must
248+-- maintain the invariant that the connection they are writing to/reading
249+-- from *must* be "included" in the 'remoteRefCount'.
250+-- 3. Since every endpoint is associated with a single socket, we regard writes
251+-- that endpoint a state change too (i.e., we take the MVar before the write
252+-- and put it back after). The reason is that we don't want to "scramble" the
253+-- output of multiple concurrent writes (either from an explicit 'send' or
254+-- the writes for control messages).
256+-- Of course, "serialize" does not mean that we want for the remote endpoint to
257+-- reply. "Send" takes the mvar, sends to the endpoint (asynchronously), and
258+-- then puts the mvar back, without waiting for the endpoint to receive the
259+-- message. Similarly, when requesting a new connection, we take the mvar,
260+-- tentatively increment the reference count, send the control request, and
261+-- then put the mvar back. When the remote host responds to the new connection
262+-- request we might have to do another state change (reduce the refcount) if
263+-- the connection request was refused but we don't want to increment the ref
264+-- count only after the remote host acknowledges the request because then a
265+-- concurrent 'close' might actually attempt to close the socket.
267+-- Since we don't do concurrent reads from the same socket we don't need to
268+-- take the lock when reading from the socket.
270+-- Moreover, we insist on the invariant (INV-CLOSE) that whenever we put an
271+-- endpoint in closed state we remove that endpoint from localConnections
272+-- first, so that if a concurrent thread reads the mvar, finds EndPointClosed,
273+-- and then looks up the endpoint in localConnections it is guaranteed to
274+-- either find a different remote endpoint, or else none at all.
276+data RemoteEndPoint = RemoteEndPoint
277+ { remoteAddress :: EndPointAddress
278+ , remoteState :: MVar RemoteState
282+ RemoteEndPointInvalid (TransportError ConnectErrorCode)
283+ | RemoteEndPointValid ValidRemoteEndPointState
284+ | RemoteEndPointClosing (MVar ()) ValidRemoteEndPointState
285+ | RemoteEndPointClosed
151−mkTargetEnd :: Chans -> ChanId -> MVar [ByteString] -> TargetEnd
287+data ValidRemoteEndPointState = ValidRemoteEndPointState
152−mkTargetEnd chans chanId chan = TargetEnd
288+ { _remoteOutgoing :: !Int
153− { -- for now we will implement this as an MVar
289+ , _remoteIncoming :: IntSet
154− receive = {-# SCC "receive" #-} takeMVar chan
290+ , remoteSocket :: N.Socket
155− , closeTargetEnd = {-# SCC "closeTargetEnd" #-} do
291+ , sendOn :: [ByteString] -> IO ()
156− (chanId', chanMap) <- takeMVar chans
157− case IntMap.lookup chanId chanMap of
158− Nothing -> putMVar chans (chanId', chanMap)
159− Just (_, socks) -> do
160− forM_ socks $ \(threadId, sock) -> do
163− let chanMap' = IntMap.delete chanId chanMap
164− putMVar chans (chanId', chanMap')
167−-- | This function waits for inbound connections. If a connection fails
294+type EndPointId = Int32
168−-- for some reason, an error is raised.
295+type ControlRequestId = Int32
169−procConnections :: Chans -> Socket -> IO ()
296+type EndPointPair = (LocalEndPoint, RemoteEndPoint)
170−procConnections chans sock = forever $ do
297+
171− let err m = error $ printf "procConnections: %s" m
298+-- | Control headers
172− (clientSock, _clientAddr) <- accept sock
299+data ControlHeader =
173− -- decode the first message to find the correct chanId
300+ RequestConnectionId -- ^ Request a new connection ID from the remote endpoint
174− bss <- recvExact clientSock (fromIntegral $ sizeOf (0 :: Int32))
301+ | CloseConnection -- ^ Tell the remote endpoint we will no longer be using a connection
175− case bss of
302+ | ControlResponse -- ^ Respond to a control request /from/ the remote endpoint
176− Left _ -> err "inbound chanId aborted"
303+ | CloseSocket -- ^ Request to close the connection (see module description)
177− Right bs -> do
304+ deriving (Enum, Bounded, Show)
178− let chanId = either err fromIntegral (decode bs :: Either String Int32)
305+
179− (chanId', chanMap) <- takeMVar chans
306+-- Response sent by /B/ to /A/ when /A/ tries to connect
180− case IntMap.lookup chanId chanMap of
307+data ConnectionRequestResponse =
181− Nothing -> do
308+ ConnectionRequestAccepted -- ^ /B/ accepts the connection
182− putMVar chans (chanId', chanMap)
309+ | ConnectionRequestEndPointInvalid -- ^ /A/ requested an invalid endpoint
183− err "cannot find chanId"
310+ | ConnectionRequestCrossed -- ^ /A/s request crossed with a request from /B/ (see protocols)
184− Just (chan, socks) -> do
311+ deriving (Enum, Bounded, Show)
185− threadId <- forkWithExceptions forkIO "Message Listener" $
312+
186− procMessages chans chanId chan clientSock
313+--------------------------------------------------------------------------------
187− let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
314+-- Top-level functionality --
188− putMVar chans (chanId', chanMap')
315+--------------------------------------------------------------------------------
190−-- | This function first extracts a header of type Word8, which determines
317+-- | Create a TCP transport
191−-- the size of the ByteString that follows. If this size is 0, this indicates
318+--
192−-- that the ByteString is large, so the next value is an Int32, which
319+-- TODOs: deal with hints
193−-- determines the size of the ByteString that follows. The ByteString is then
320+createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)
194−-- extracted from the socket, and then written to the Chan only when
321+createTransport host port = do
195−-- complete. If either of the first header size is null this indicates the
322+ state <- newMVar . TransportValid $ ValidTransportState { _localEndPoints = Map.empty }
196−-- socket has closed.
323+ let transport = TCPTransport { transportState = state
197−procMessages :: Chans -> ChanId -> MVar [ByteString] -> Socket -> IO ()
324+ , transportHost = host
198−procMessages chans chanId chan sock = do
325+ , transportPort = port
199− esizeBS <- recvExact sock 1
326+ }
200− case esizeBS of
327+ tryIO $ do
201− Left _ -> closeSocket
328+ -- For a discussion of the use of N.sOMAXCONN, see
202− Right sizeBS ->
329+ -- http://tangentsoft.net/wskfaq/advanced.html#backlog
203− let size = toEnum . fromEnum $ BS.index sizeBS 0 in
330+ -- http://www.linuxjournal.com/files/linuxjournal.com/linuxjournal/articles/023/2333/2333s2.html
204− if size < 255
331+ transportThread <- forkServer host port N.sOMAXCONN (terminationHandler transport) (handleConnectionRequest transport)
205− then procMessage size
332+ return Transport { newEndPoint = apiNewEndPoint transport
206− else do esizeBS' <- recvExact sock (fromIntegral $ sizeOf (0 :: Int32))
333+ , closeTransport = killThread transportThread -- This will invoke the termination handler
207− case esizeBS' of
334+ }
208− Left _ -> closeSocket
335+ where
209− Right sizeBS' -> do sizeDec <- decodeLength sizeBS'
336+ terminationHandler :: TCPTransport -> SomeException -> IO ()
210− procMessage sizeDec
337+ terminationHandler transport _ = do
211− where
338+ -- TODO: we currently don't make a distinction between endpoint failure and manual closure
212− err m = error $ printf "procMessages: %s" m
339+ mTSt <- modifyMVar (transportState transport) $ \st -> case st of
213− closeSocket :: IO ()
340+ TransportValid vst -> return (TransportClosed, Just vst)
214− closeSocket = do
341+ TransportClosed -> return (TransportClosed, Nothing)
215− (chanId', chanMap) <- takeMVar chans
342+ case mTSt of
216− case IntMap.lookup chanId chanMap of
343+ Nothing ->
217− Nothing -> do
344+ -- Transport already closed
218− putMVar chans (chanId', chanMap)
345+ return ()
219− err "chanId not found."
346+ Just tSt ->
220− Just (chan, socks) -> do
347+ mapM_ (apiCloseEndPoint transport) (Map.elems $ tSt ^. localEndPoints)
221− let socks' = filter ((/= sock) . snd) socks
348+
222− let chanMap' = IntMap.insert chanId (chan, socks') chanMap
349+--------------------------------------------------------------------------------
223− putMVar chans (chanId', chanMap')
350+-- API functions --
224− sClose sock
351+--------------------------------------------------------------------------------
225− procMessage :: Int32 -> IO ()
352+
226− procMessage size = do
353+-- | Create a new endpoint
227− ebs <- recvExact sock size
354+apiNewEndPoint :: TCPTransport -> IO (Either (TransportError NewEndPointErrorCode) EndPoint)
228− case ebs of
355+apiNewEndPoint transport = try $ do
229− Left _ -> closeSocket
356+ ourEndPoint <- createLocalEndPoint transport
230− Right bs -> do
357+ return EndPoint
231− putMVar chan [bs]
358+ { receive = readChan (localChannel ourEndPoint)
232− procMessages chans chanId chan sock
359+ , address = localAddress ourEndPoint
233−
360+ , connect = apiConnect ourEndPoint
234−-- | The normal result of `recvExact sock n` is `Right ByteString`
361+ , closeEndPoint = apiCloseEndPoint transport ourEndPoint
235−-- whose string of length `n`, received from `sock`. If fewer than `n`
362+ , newMulticastGroup = return . Left $ newMulticastGroupError
236−-- bytes are read from `sock` before it closes, the result is `Left
363+ , resolveMulticastGroup = return . Left . const resolveMulticastGroupError
237−-- ByteString` whose string is those bytes that were received. No more
364+ }
238−-- bytes than necessary are read from the socket. NB: This uses
365+ where
239−-- Network.Socket.ByteString.recv, which may *discard* superfluous
366+ newMulticastGroupError = TransportError NewMulticastGroupUnsupported "TCP does not support multicast"
240−-- input depending on the socket type.
367+ resolveMulticastGroupError = TransportError ResolveMulticastGroupUnsupported "TCP does not support multicast"
242−recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
369+-- | Connnect to an endpoint
243−recvExact sock l = interceptAllExn "recvExact" $ do
370+apiConnect :: LocalEndPoint -- ^ Local end point
244− res <- BSI.createAndTrim (fromIntegral l) (go 0)
371+ -> EndPointAddress -- ^ Remote address
245− case BS.length res of
372+ -> Reliability -- ^ Reliability (ignored)
246− n | n == (fromIntegral l) -> return $ Right res
373+ -> IO (Either (TransportError ConnectErrorCode) Connection)
247− n -> return $ Left res
374+apiConnect ourEndPoint theirAddress _ | localAddress ourEndPoint == theirAddress =
248− where
375+ connectToSelf ourEndPoint
249− go :: Int -> Ptr Word8 -> IO Int
376+apiConnect ourEndPoint theirAddress _ = try $ do
250− go n ptr | n == (fromIntegral l) = return n
377+ (theirEndPoint, connId) <- requestConnectionTo ourEndPoint theirAddress
251− go n ptr = do
378+ -- connAlive can be an IORef rather than an MVar because it is protected by
252− (p, off, len) <- BSI.toForeignPtr <$> NBS.recv sock (min (fromIntegral l-n) 4096)
379+ -- the remoteState MVar. We don't need the overhead of locking twice.
253− if len == 0
380+ connAlive <- newIORef True
254− then return n
381+ return $ Connection { send = apiSend theirEndPoint connId connAlive
255− else withForeignPtr p $ \p -> do
382+ , close = apiClose theirEndPoint connId connAlive
256− BSI.memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)
383+ }
258−#else
385+-- | Close a connection
259−recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
386+--
260−recvExact sock n =
387+-- RELY: Remote endpoint must be in 'RemoteEndPointValid' or 'RemoteEndPointClosed'
261− interceptAllExn "recvExact" $
388+-- GUARANTEE: If the connection is alive on entry then the remote endpoint will
262− go [] sock n
389+-- either be RemoteEndPointValid or RemoteEndPointClosing. Otherwise, the state
263− where
390+-- of the remote endpoint will not be changed.
264− go :: [ByteString] -> Socket -> Int32 -> IO (Either ByteString ByteString)
391+--
265− go bss _ 0 = return (Right $ BS.concat (reverse bss))
392+-- TODO: We ignore errors during a close. Is that right?
266− go bss sock n = do
393+apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()
267− bs <- NBS.recv sock (min (fromIntegral n) 4096)
394+apiClose theirEndPoint connId connAlive = void . tryIO $ do
268− if BS.null bs
395+ modifyMVar_ (remoteState theirEndPoint) $ \st -> case st of
269− then return (Left $ BS.concat (reverse bss))
396+ RemoteEndPointValid vst -> do
270− else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
397+ alive <- readIORef connAlive
273−interceptAllExn msg =
400+ writeIORef connAlive False
274− Control.Exception.handle $ \ e ->
401+ sendOn vst [encodeInt32 CloseConnection, encodeInt32 connId]
275− case fromException e of
402+ return (RemoteEndPointValid . (remoteOutgoing ^: (+ 1)) $ vst)
276− Just ThreadKilled -> throw e
403+ else
277− Nothing -> do
404+ return (RemoteEndPointValid vst)
278− BS.hPutStrLn stderr $ BS.pack$ "Exception inside "++msg++": "++show e
405+ RemoteEndPointClosed -> do
279− throw e
406+ return st
280−-- throw (e :: SomeException)
407+ RemoteEndPointClosing _ _ ->
281−
408+ fail "apiClose RELY violation"
282−forkWithExceptions :: (IO () -> IO ThreadId) -> String -> IO () -> IO ThreadId
409+ RemoteEndPointInvalid _ ->
283−forkWithExceptions forkit descr action = do
410+ fail "apiClose RELY violation"
284− parent <- myThreadId
411+ closeIfUnused theirEndPoint
286− Control.Exception.catch action
413+-- | Send data across a connection
288− case fromException e of
415+-- RELY: The remote endpoint must be in state 'RemoteEndPointValid' or 'RemoteEndPointClosed'
289− Just ThreadKilled ->
416+-- GUARANTEE: The state of the remote endpoint will not be changed.
290−-- BSS.hPutStrLn stderr $ BSS.pack$ "Note: Child thread killed: "++descr
417+apiSend :: RemoteEndPoint -> ConnectionId -> IORef Bool -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
291− return ()
418+apiSend theirEndPoint connId connAlive payload = do
292− _ -> do
419+ withMVar (remoteState theirEndPoint) $ \st -> case st of
293− BS.hPutStrLn stderr $ BS.pack$ "Exception inside child thread "++descr++": "++show e
420+ RemoteEndPointValid vst -> do
294− throwTo parent (e::SomeException)
421+ alive <- readIORef connAlive
295− )
422+ try $ if alive
296−
423+ then mapExceptionIO sendFailed $ sendOn vst (encodeInt32 connId : prependLength payload)
297−-- | Encode length (manual for now)
424+ else throw $ TransportError SendFailed "Connection closed"
298−encodeLength :: Int32 -> IO ByteString
425+ RemoteEndPointClosed -> do
299−encodeLength i32 =
426+ return (Left $ TransportError SendFailed "Endpoint closed")
300− BSI.create 4 $ \p ->
427+ RemoteEndPointClosing _ _ ->
301− pokeByteOff p 0 (htonl (fromIntegral i32))
428+ error "apiSend RELY violation"
302−
429+ RemoteEndPointInvalid _ ->
303−-- | Decode length (manual for now)
430+ error "apiSend RELY violation"
304−decodeLength :: ByteString -> IO Int32
431+ where
305−decodeLength bs =
432+ sendFailed :: IOException -> TransportError SendErrorCode
306− let (fp, _, _) = BSI.toForeignPtr bs in
433+ sendFailed = TransportError SendFailed . show
307− withForeignPtr fp $ \p -> do
434+
308− w32 <- peekByteOff p 0
435+-- | Force-close the endpoint
309− return (fromIntegral (ntohl w32))
436+apiCloseEndPoint :: TCPTransport -> LocalEndPoint -> IO ()
437+apiCloseEndPoint transport ourEndPoint = do
438+ -- Remove the reference from the transport state
439+ removeLocalEndPoint transport ourEndPoint
440+ -- Close the local endpoint
441+ mOurState <- modifyMVar (localState ourEndPoint) $ \st ->
443+ LocalEndPointValid vst ->
444+ return (LocalEndPointClosed, Just vst)
445+ LocalEndPointClosed ->
446+ return (LocalEndPointClosed, Nothing)
448+ Nothing -> -- Already closed
451+ -- Close all endpoints and kill all threads
452+ forM_ (Map.elems $ vst ^. localConnections) $ tryCloseRemoteSocket
453+ forM_ (vst ^. internalThreads) killThread
454+ -- We send a single message to the endpoint that it is closed. Subsequent
455+ -- calls will block. We could change this so that all subsequent calls to
456+ -- receive return an error, but this would mean checking for some state on
457+ -- every call to receive, which is an unnecessary overhead.
458+ writeChan (localChannel ourEndPoint) EndPointClosed
460+ -- Close the remote socket and return the set of all incoming connections
461+ tryCloseRemoteSocket :: RemoteEndPoint -> IO ()
462+ tryCloseRemoteSocket theirEndPoint = do
463+ -- We make an attempt to close the connection nicely (by sending a CloseSocket first)
464+ modifyMVar_ (remoteState theirEndPoint) $ \st ->
466+ RemoteEndPointInvalid _ ->
468+ RemoteEndPointValid conn -> do
469+ -- Try to send a CloseSocket request
470+ tryIO $ sendOn conn [encodeInt32 CloseSocket]
471+ -- .. but even if it fails, close the socket anyway
472+ -- (hence, two separate calls to tryIO)
473+ tryIO $ N.sClose (remoteSocket conn)
474+ return RemoteEndPointClosed
475+ RemoteEndPointClosing _ conn -> do
476+ tryIO $ N.sClose (remoteSocket conn)
477+ return RemoteEndPointClosed
478+ RemoteEndPointClosed ->
479+ return RemoteEndPointClosed
481+-- | Special case of 'apiConnect': connect an endpoint to itself
482+connectToSelf :: LocalEndPoint -> IO (Either (TransportError ConnectErrorCode) Connection)
483+connectToSelf ourEndPoint = do
484+ -- Here connAlive must an MVar because it is not protected by another lock
485+ connAlive <- newMVar True
486+ -- TODO: catch exception
487+ connId <- getNextConnectionId ourEndPoint
488+ writeChan ourChan (ConnectionOpened connId ReliableOrdered (localAddress ourEndPoint))
489+ return . Right $ Connection { send = selfSend connAlive connId
490+ , close = selfClose connAlive connId
493+ ourChan :: Chan Event
494+ ourChan = localChannel ourEndPoint
496+ selfSend :: MVar Bool -> ConnectionId -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
497+ selfSend connAlive connId msg = do
498+ modifyMVar connAlive $ \alive ->
501+ writeChan ourChan (Received connId msg)
502+ return (alive, Right ())
504+ return (alive, Left (TransportError SendFailed "Connection closed"))
506+ selfClose :: MVar Bool -> ConnectionId -> IO ()
507+ selfClose connAlive connId = do
508+ modifyMVar_ connAlive $ \alive -> do
509+ when alive $ writeChan ourChan (ConnectionClosed connId)
512+--------------------------------------------------------------------------------
513+-- Lower level functionality --
514+--------------------------------------------------------------------------------
516+-- | Create a new local endpoint
518+-- May throw a TransportError NewEndPointErrorCode exception if the transport is closed.
519+createLocalEndPoint :: TCPTransport -> IO LocalEndPoint
520+createLocalEndPoint transport = do
522+ state <- newMVar . LocalEndPointValid $ ValidLocalEndPointState
523+ { _nextConnectionId = firstNonReservedConnectionId
524+ , _pendingCtrlRequests = IntMap.empty
525+ , _nextCtrlRequestId = 0
526+ , _localConnections = Map.empty
527+ , _internalThreads = []
529+ modifyMVar (transportState transport) $ \st -> case st of
530+ TransportValid vst -> do
531+ let ix = fromIntegral $ Map.size (vst ^. localEndPoints)
532+ let addr = encodeEndPointAddress (transportHost transport) (transportPort transport) ix
533+ let localEndPoint = LocalEndPoint { localAddress = addr
534+ , localChannel = chan
535+ , localState = state
537+ return (TransportValid . (localEndPointAt addr ^= Just localEndPoint) $ vst, localEndPoint)
539+ throw (TransportError NewEndPointFailed "Transport closed")
541+-- | Request a connection to a remote endpoint
543+-- This will block until we get a connection ID from the remote endpoint; if
544+-- the remote endpoint was in 'RemoteEndPointClosing' state then we will
545+-- additionally block until that is resolved.
547+-- May throw a TransportError ConnectErrorCode exception.
548+requestConnectionTo :: LocalEndPoint
550+ -> IO (RemoteEndPoint, ConnectionId)
551+requestConnectionTo ourEndPoint theirAddress = go
554+ -- Find the remote endpoint (create it if it doesn't yet exist)
555+ theirEndPoint <- findTheirEndPoint
556+ let theirState = remoteState theirEndPoint
558+ -- Before we initiate the new connection request we want to make sure that
559+ -- refcount on the endpoint is incremented so that a concurrent thread will
560+ -- not close the connection. Note that if IF we return RemoteEndPointValid
561+ -- here then we can rely on the endpoint remaining in that state.
562+ endPointStateSnapshot <- modifyMVar theirState $ \st ->
564+ RemoteEndPointValid ep ->
565+ return (RemoteEndPointValid . (remoteOutgoing ^: (+ 1)) $ ep, st)
569+ -- From this point on we are guaranteed the refcount is positive, provided
570+ -- that the endpoint was valid. We still need to deal with the case where
571+ -- it was not valid, however, which we didn't want to do while holding the
574+ -- Although 'endPointStateSnapshot' here refers to a snapshot of the
575+ -- endpoint state, and might have changed in the meantime, these changes
577+ case endPointStateSnapshot of
578+ RemoteEndPointInvalid err -> do
581+ RemoteEndPointClosing resolved _ ->
582+ -- If the remote endpoint is closing, then we need to block until
583+ -- this is resolved and we then try again
584+ readMVar resolved >> go
586+ RemoteEndPointClosed -> do
587+ -- EndPointClosed indicates that a concurrent thread was in the
588+ -- process of closing the TCP connection to the remote endpoint when
589+ -- we obtained a reference to it. The remote endpoint will now have
590+ -- been removed from ourState, so we simply try again.
593+ RemoteEndPointValid _ -> do
594+ -- On a failure we decrement the refcount again and return an error.
595+ -- The only state the remote endpoint can be in at this point is
596+ -- valid. As mentioned above, we can rely on the endpoint being in
597+ -- valid state at this point.
598+ let failureHandler :: IOException -> IO b
599+ failureHandler err = do
600+ modifyMVar_ theirState $ \(RemoteEndPointValid ep) ->
601+ return (RemoteEndPointValid . (remoteOutgoing ^: (+ 1)) $ ep)
602+ -- TODO: should we call closeIfUnused here?
603+ throw $ TransportError ConnectFailed (show err)
605+ -- Do the actual connection request. This blocks until the remote
606+ -- endpoint replies (but note that we don't hold any locks at this
608+ reply <- handle failureHandler $ doRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
609+ case decodeInt32 . BS.concat $ reply of
610+ Nothing -> failureHandler (userError "Invalid integer")
611+ Just cid -> return (theirEndPoint, cid)
613+ -- If this is a new endpoint, fork a thread to listen for incoming
614+ -- connections. We don't want to do this while we hold the lock, because
615+ -- forkEndPointThread modifies the local state too (to record the thread
617+ findTheirEndPoint :: IO RemoteEndPoint
618+ findTheirEndPoint = do
619+ (theirEndPoint, isNew) <- modifyMVar ourState $ \st -> case st of
620+ LocalEndPointClosed ->
621+ throw (TransportError ConnectFailed "Local endpoint closed")
622+ LocalEndPointValid vst ->
623+ case vst ^. localConnectionTo theirAddress of
624+ Just theirEndPoint ->
625+ return (st, (theirEndPoint, False))
627+ theirState <- newEmptyMVar
628+ let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
629+ , remoteState = theirState
631+ return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) $ vst, (theirEndPoint, True))
633+ -- The only way for forkEndPointThread to fail is if the local endpoint
634+ -- gets closed. This error will be caught elsewhere, so we ignore it
636+ when isNew . void . forkEndPointThread ourEndPoint $
637+ setupRemoteEndPoint (ourEndPoint, theirEndPoint)
638+ return theirEndPoint
640+ ourState :: MVar LocalEndPointState
641+ ourState = localState ourEndPoint
643+-- | Set up a remote endpoint
645+-- RELY: The state of the remote endpoint must be uninitialized.
646+-- GUARANTEE: Will only change the state to RemoteEndPointValid or
647+-- RemoteEndPointInvalid.
648+setupRemoteEndPoint :: EndPointPair -> IO ()
649+setupRemoteEndPoint (ourEndPoint, theirEndPoint) = do
650+ result <- socketToEndPoint (localAddress ourEndPoint) (remoteAddress theirEndPoint)
652+ Right (sock, ConnectionRequestAccepted) -> do
653+ let vst = ValidRemoteEndPointState
654+ { remoteSocket = sock
655+ , _remoteOutgoing = 0
656+ , _remoteIncoming = IntSet.empty
657+ , sendOn = sendMany sock
659+ putMVar theirState (RemoteEndPointValid vst)
660+ handleIncomingMessages (ourEndPoint, theirEndPoint)
661+ Right (sock, ConnectionRequestEndPointInvalid) -> do
662+ -- We remove the endpoint from our local state again because the next
663+ -- call to 'connect' might give a different result. Threads that were
664+ -- waiting on the result of this call to connect will get the
665+ -- RemoteEndPointInvalid; subsequent threads will initiate a new
666+ -- connection requests.
667+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
668+ putMVar theirState (RemoteEndPointInvalid (invalidAddress "Invalid endpoint"))
670+ Right (sock, ConnectionRequestCrossed) -> do
673+ -- See comment above
674+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
675+ putMVar theirState (RemoteEndPointInvalid err)
677+ theirState = remoteState theirEndPoint
678+ invalidAddress = TransportError ConnectNotFound
680+-- | Establish a connection to a remote endpoint
681+socketToEndPoint :: EndPointAddress -- ^ Our address
682+ -> EndPointAddress -- ^ Their address
683+ -> IO (Either (TransportError ConnectErrorCode) (N.Socket, ConnectionRequestResponse))
684+socketToEndPoint (EndPointAddress ourAddress) theirAddress = try $ do
685+ (host, port, theirEndPointId) <- case decodeEndPointAddress theirAddress of
686+ Nothing -> throw (failed . userError $ "Could not parse")
687+ Just dec -> return dec
688+ addr:_ <- mapExceptionIO invalidAddress $ N.getAddrInfo Nothing (Just host) (Just port)
689+ bracketOnError (createSocket addr) N.sClose $ \sock -> do
690+ mapExceptionIO failed $ N.setSocketOption sock N.ReuseAddr 1
691+ mapExceptionIO invalidAddress $ N.connect sock (N.addrAddress addr)
692+ response <- mapExceptionIO failed $ do
693+ sendMany sock (encodeInt32 theirEndPointId : prependLength [ourAddress])
695+ case tryToEnum response of
696+ Nothing -> throw (failed . userError $ "Unexpected response")
697+ Just r -> return (sock, r)
699+ createSocket :: N.AddrInfo -> IO N.Socket
700+ createSocket addr = mapExceptionIO insufficientResources $ do
701+ sock <- N.socket (N.addrFamily addr) N.Stream N.defaultProtocol
702+ -- putStrLn $ "Created client socket " ++ show sock
705+ invalidAddress, insufficientResources, failed :: IOException -> TransportError ConnectErrorCode
706+ invalidAddress = TransportError ConnectNotFound . show
707+ insufficientResources = TransportError ConnectInsufficientResources . show
708+ failed = TransportError ConnectFailed . show
710+-- | Remove reference to a remote endpoint from a local endpoint
712+-- If the local endpoint is closed, do nothing
713+removeRemoteEndPoint :: EndPointPair -> IO ()
714+removeRemoteEndPoint (ourEndPoint, theirEndPoint) = do
715+ modifyMVar_ (localState ourEndPoint) $ \st -> case st of
716+ LocalEndPointValid vst ->
717+ return (LocalEndPointValid . (localConnectionTo (remoteAddress theirEndPoint) ^= Nothing) $ vst)
718+ LocalEndPointClosed ->
719+ return LocalEndPointClosed
721+-- | Remove reference to a local endpoint from the transport state
723+-- Does nothing if the transport is closed
724+removeLocalEndPoint :: TCPTransport -> LocalEndPoint -> IO ()
725+removeLocalEndPoint transport ourEndPoint = do
726+ modifyMVar_ (transportState transport) $ \st -> case st of
727+ TransportValid vst ->
728+ return (TransportValid . (localEndPointAt (localAddress ourEndPoint) ^= Nothing) $ vst)
730+ return TransportClosed
732+-- | Encode end point address
733+encodeEndPointAddress :: N.HostName -> N.ServiceName -> EndPointId -> EndPointAddress
734+encodeEndPointAddress host port ix = EndPointAddress . BSC.pack $
735+ host ++ ":" ++ port ++ ":" ++ show ix
737+-- | Decode end point address
738+decodeEndPointAddress :: EndPointAddress -> Maybe (N.HostName, N.ServiceName, EndPointId)
739+decodeEndPointAddress (EndPointAddress bs) = case map BSC.unpack $ BSC.split ':' bs of
740+ [host, port, endPointIdStr] ->
741+ case reads endPointIdStr of
742+ [(endPointId, "")] -> Just (host, port, endPointId)
747+-- | Do a (blocking) remote request
749+-- RELY: Remote endpoint must be in valid state.
750+-- GUARANTEE: Will not change the state of the remote endpoint.
752+-- May throw IO (user) exception if the local endpoint is closed or if the send
754+doRemoteRequest :: EndPointPair -> ControlHeader -> IO [ByteString]
755+doRemoteRequest (ourEndPoint, theirEndPoint) header = do
756+ reply <- newEmptyMVar
757+ reqId <- modifyMVar (localState ourEndPoint) $ \st -> case st of
758+ LocalEndPointValid vst -> do
759+ let reqId = vst ^. nextCtrlRequestId
760+ return (LocalEndPointValid . (nextCtrlRequestId ^: (+ 1)) . (pendingCtrlRequestsAt reqId ^= Just reply) $ vst, reqId)
761+ LocalEndPointClosed ->
762+ throw (userError "Local endpoint closed")
763+ withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid vst) ->
764+ sendOn vst [encodeInt32 header, encodeInt32 reqId]
767+-- | Check if the remote endpoint is unused, and if so, send a CloseSocket request
768+closeIfUnused :: RemoteEndPoint -> IO ()
769+closeIfUnused theirEndPoint = modifyMVar_ (remoteState theirEndPoint) $ \st ->
771+ RemoteEndPointValid vst -> do
772+ if vst ^. remoteOutgoing == 0 && IntSet.null (vst ^. remoteIncoming)
774+ resolved <- newEmptyMVar
775+ sendOn vst [encodeInt32 CloseSocket]
776+ return (RemoteEndPointClosing resolved vst)
782+-- | Fork a new thread and store its ID as part of the transport state
784+-- If the local end point is closed this function does nothing (no thread is
785+-- spawned). Returns whether or not a thread was spawned.
786+forkEndPointThread :: LocalEndPoint -> IO () -> IO Bool
787+forkEndPointThread ourEndPoint p =
788+ modifyMVar ourState $ \st -> case st of
789+ LocalEndPointValid vst -> do
790+ tid <- forkIO (p >> removeThread)
791+ return (LocalEndPointValid . (internalThreads ^: (tid :)) $ vst, True)
792+ LocalEndPointClosed ->
793+ return (LocalEndPointClosed, False)
795+ removeThread :: IO ()
798+ modifyMVar_ ourState $ \st -> case st of
799+ LocalEndPointValid vst ->
800+ return (LocalEndPointValid . (internalThreads ^: filter (/= tid)) $ vst)
801+ LocalEndPointClosed ->
802+ return LocalEndPointClosed
804+ ourState :: MVar LocalEndPointState
805+ ourState = localState ourEndPoint
807+--------------------------------------------------------------------------------
808+-- Incoming requests --
809+--------------------------------------------------------------------------------
811+-- | Handle a connection request (that is, a remote endpoint that is trying to
812+-- establish a TCP connection with us)
814+-- 'handleConnectionRequest' runs in the context of the transport thread, which
815+-- can be killed asynchronously by 'closeTransport'. We fork a separate thread
816+-- as soon as we have located the lcoal endpoint that the remote endpoint is
817+-- interested in. We cannot fork any sooner because then we have no way of
818+-- storing the thread ID and hence no way of killing the thread when we take
819+-- the transport down. We must be careful to close the socket when a (possibly
820+-- asynchronous, ThreadKilled) exception occurs. (If an exception escapes from
821+-- handleConnectionRequest the transport will be shut down.)
822+handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()
823+handleConnectionRequest transport sock = handle tryCloseSocket $ do
824+ ourEndPointId <- recvInt32 sock
825+ theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock
826+ let ourAddress = encodeEndPointAddress (transportHost transport) (transportPort transport) ourEndPointId
827+ ourEndPoint <- withMVar (transportState transport) $ \st -> case st of
828+ TransportValid vst ->
829+ case vst ^. localEndPointAt ourAddress of
831+ sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
832+ fail "Invalid endpoint"
835+ TransportClosed -> do
836+ fail "Transport closed"
837+ void . forkEndPointThread ourEndPoint $ go ourEndPoint theirAddress
839+ go :: LocalEndPoint -> EndPointAddress -> IO ()
840+ go ourEndPoint theirAddress = do
841+ mEndPoint <- handle (\e -> invalidEndPoint e >> return Nothing) $ do
842+ (crossed, theirEndPoint) <- modifyMVar (localState ourEndPoint) $ \st -> case st of
843+ LocalEndPointValid vst ->
844+ case vst ^. localConnectionTo theirAddress of
846+ theirState <- newEmptyMVar
847+ let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
848+ , remoteState = theirState
850+ return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) $ vst, (False, theirEndPoint))
851+ Just theirEndPoint ->
852+ return (st, (localAddress ourEndPoint < theirAddress, theirEndPoint))
853+ LocalEndPointClosed ->
854+ fail "Local endpoint closed"
857+ tryIO $ sendMany sock [encodeInt32 ConnectionRequestCrossed]
858+ tryIO $ N.sClose sock
861+ let vst = ValidRemoteEndPointState
862+ { remoteSocket = sock
863+ , _remoteOutgoing = 0
864+ , _remoteIncoming = IntSet.empty
865+ , sendOn = sendMany sock
867+ -- TODO: this putMVar might block if the remote endpoint sends a
868+ -- connection request for a local endpoint that it is already
870+ putMVar (remoteState theirEndPoint) (RemoteEndPointValid vst)
871+ sendMany sock [encodeInt32 ConnectionRequestAccepted]
872+ return (Just theirEndPoint)
873+ -- If we left the scope of the exception handler with a return value of
874+ -- Nothing then the socket is already closed; otherwise, the socket has
875+ -- been recorded as part of the remote endpoint. Either way, we no longer
876+ -- have to worry about closing the socket on receiving an asynchronous
877+ -- exception from this point forward.
881+ Just theirEndPoint ->
882+ handleIncomingMessages (ourEndPoint, theirEndPoint)
884+ tryCloseSocket :: IOException -> IO ()
885+ tryCloseSocket _ = void . tryIO $ N.sClose sock
887+ invalidEndPoint :: IOException -> IO ()
888+ invalidEndPoint ex = do
889+ tryIO $ sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
892+-- | Handle requests from a remote endpoint.
894+-- Returns only if the remote party closes the socket or if an error occurs.
896+-- RELY: The remote endpoint must be in RemoteEndPointValid or
897+-- RemoteEndPointClosing state. If the latter, then the 'resolved' MVar
898+-- associated with the closing state must be empty.
899+-- GUARANTEE: May change the remote endpoint to RemoteEndPointClosed state.
900+handleIncomingMessages :: EndPointPair -> IO ()
901+handleIncomingMessages (ourEndPoint, theirEndPoint) = do
902+ -- For efficiency sake we get the socket once and for all
903+ sock <- withMVar theirState $ \st ->
905+ RemoteEndPointValid ep ->
906+ return (remoteSocket ep)
907+ RemoteEndPointClosing _ ep ->
908+ return (remoteSocket ep)
910+ error "handleIncomingMessages RELY violation"
912+ mCleanExit <- tryIO $ go sock
914+ Right () -> return ()
915+ Left err -> prematureExit sock err
918+ go :: N.Socket -> IO ()
920+ connId <- recvInt32 sock
921+ if connId >= firstNonReservedConnectionId
923+ readMessage sock connId
926+ case tryToEnum (fromIntegral connId) of
927+ Just RequestConnectionId -> do
928+ recvInt32 sock >>= createNewConnection
930+ Just ControlResponse -> do
931+ recvInt32 sock >>= readControlResponse sock
933+ Just CloseConnection -> do
934+ recvInt32 sock >>= closeConnection
936+ Just CloseSocket -> do
937+ didClose <- closeSocket sock
938+ unless didClose $ go sock
940+ -- Invalid control request, exit
941+ hPutStrLn stderr "Warning: invalid control request"
943+ -- Create a new connection
944+ createNewConnection :: ControlRequestId -> IO ()
945+ createNewConnection reqId = do
946+ newId <- getNextConnectionId ourEndPoint
947+ modifyMVar_ theirState $ \st -> do
949+ RemoteEndPointValid vst ->
950+ return (remoteIncoming ^: IntSet.insert newId $ vst)
951+ RemoteEndPointClosing resolved vst -> do
952+ -- If the endpoint is in closing state that means we send a
953+ -- CloseSocket request to the remote endpoint. If the remote
954+ -- endpoint replies with the request to create a new connection, it
955+ -- either ignored our request or it sent the request before it got
956+ -- ours. Either way, at this point we simply restore the endpoint
957+ -- to RemoteEndPointValid
959+ return (remoteIncoming ^= IntSet.singleton newId $ vst)
961+ error "handleIncomingMessages RELY violation"
962+ sendOn vst ( encodeInt32 ControlResponse
964+ : prependLength [encodeInt32 newId]
966+ -- We add the new connection ID to the list of open connections only once the
967+ -- endpoint has been notified of the new connection (sendOn may fail)
968+ return (RemoteEndPointValid vst)
969+ writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)
971+ -- Read a control response
972+ readControlResponse :: N.Socket -> ControlRequestId -> IO ()
973+ readControlResponse sock reqId = do
974+ response <- recvWithLength sock
975+ mmvar <- modifyMVar ourState $ \st -> case st of
976+ LocalEndPointValid vst ->
977+ return (LocalEndPointValid . (pendingCtrlRequestsAt reqId ^= Nothing) $ vst, vst ^. pendingCtrlRequestsAt reqId)
978+ LocalEndPointClosed ->
979+ fail "Local endpoint closed"
982+ hPutStrLn stderr $ "Warning: Invalid request ID"
983+ return () -- Invalid request ID. TODO: We just ignore it?
985+ putMVar mvar response
987+ -- Close a connection
988+ closeConnection :: ConnectionId -> IO ()
989+ closeConnection cid = do
990+ -- TODO: we should check that this connection is in fact open
991+ writeChan ourChannel (ConnectionClosed cid)
992+ modifyMVar_ theirState $ \(RemoteEndPointValid vst) ->
993+ return (RemoteEndPointValid . (remoteIncoming ^: IntSet.delete cid) $ vst)
994+ closeIfUnused theirEndPoint
996+ -- Close the socket (if we don't have any outgoing connections)
997+ closeSocket :: N.Socket -> IO Bool
998+ closeSocket sock = do
999+ -- We need to check if we can close the socket (that is, if we don't have
1000+ -- any outgoing connections), and once we are sure that we can put the
1001+ -- endpoint in Closed state. However, by INV-CLOSE we can only put the
1002+ -- endpoint in Closed state once we remove the endpoint from our local
1003+ -- connections. But we can do that only once we are sure that we can
1004+ -- close the endpoint. Catch-22. We can resolve the catch-22 by locking
1005+ -- /both/ our local state and the endpoint state, but at the cost of
1006+ -- introducing a double lock and all the associated perils, or by putting
1007+ -- the remote endpoint in Closing state first. We opt for the latter.
1008+ canClose <- modifyMVar theirState $ \st ->
1010+ RemoteEndPointValid vst -> do
1011+ -- We regard a CloseSocket message as an (optimized) way for the
1012+ -- remote endpoint to indicate that all its connections to us are
1013+ -- now properly closed
1014+ forM_ (IntSet.elems $ vst ^. remoteIncoming) $ \cid ->
1015+ writeChan ourChannel (ConnectionClosed cid)
1016+ let vst' = remoteIncoming ^= IntSet.empty $ vst
1017+ -- Check if we agree that the connection should be closed
1018+ if vst' ^. remoteOutgoing == 0
1020+ -- Attempt to reply (but don't insist)
1021+ tryIO $ sendOn vst' [encodeInt32 CloseSocket]
1022+ resolved <- newEmptyMVar
1023+ return (RemoteEndPointClosing resolved vst', Just resolved)
1025+ return (RemoteEndPointValid vst', Nothing)
1026+ RemoteEndPointClosing resolved _ ->
1027+ return (st, Just resolved)
1029+ error "handleIncomingConnections RELY violation"
1034+ Just resolved -> do
1035+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
1036+ modifyMVar_ theirState $ return . const RemoteEndPointClosed
1038+ putMVar resolved ()
1041+ -- Read a message and output it on the endPoint's channel
1042+ readMessage :: N.Socket -> ConnectionId -> IO ()
1043+ readMessage sock connId = recvWithLength sock >>= writeChan ourChannel . Received connId
1046+ ourChannel = localChannel ourEndPoint
1047+ ourState = localState ourEndPoint
1048+ theirState = remoteState theirEndPoint
1049+ theirAddr = remoteAddress theirEndPoint
1051+ -- Deal with a premature exit
1052+ prematureExit :: N.Socket -> IOException -> IO ()
1053+ prematureExit sock err = do
1055+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
1056+ mUnclosedConnections <- modifyMVar theirState $ \st ->
1058+ RemoteEndPointInvalid _ ->
1059+ error "handleIncomingMessages RELY violation"
1060+ RemoteEndPointValid vst ->
1061+ return (RemoteEndPointClosed, Just $ vst ^. remoteIncoming)
1062+ RemoteEndPointClosing _ _ ->
1063+ return (RemoteEndPointClosed, Nothing)
1064+ RemoteEndPointClosed ->
1065+ return (st, Nothing)
1067+ -- We send the connection lost message even if unclosedConnections is the
1068+ -- empty set, because *outgoing* connections will be broken too now
1069+ case mUnclosedConnections of
1072+ Just unclosedConnections -> writeChan ourChannel . ErrorEvent $
1073+ TransportError (EventConnectionLost (remoteAddress theirEndPoint) (IntSet.elems unclosedConnections)) (show err)
1075+-- | Get the next connection ID
1077+-- Throws an IO exception when the endpoint is closed.
1078+getNextConnectionId :: LocalEndPoint -> IO ConnectionId
1079+getNextConnectionId ourEndpoint =
1080+ modifyMVar (localState ourEndpoint) $ \st -> case st of
1081+ LocalEndPointValid vst -> do
1082+ let connId = vst ^. nextConnectionId
1083+ return (LocalEndPointValid . (nextConnectionId ^= connId + 1) $ vst, connId)
1084+ LocalEndPointClosed ->
1085+ fail "Local endpoint closed"
1087+--------------------------------------------------------------------------------
1089+--------------------------------------------------------------------------------
1091+-- | We reserve a bunch of connection IDs for control messages
1092+firstNonReservedConnectionId :: ConnectionId
1093+firstNonReservedConnectionId = 1024
1095+--------------------------------------------------------------------------------
1096+-- Accessor definitions --
1097+--------------------------------------------------------------------------------
1099+localEndPoints :: Accessor ValidTransportState (Map EndPointAddress LocalEndPoint)
1100+localEndPoints = accessor _localEndPoints (\es st -> st { _localEndPoints = es })
1102+pendingCtrlRequests :: Accessor ValidLocalEndPointState (IntMap (MVar [ByteString]))
1103+pendingCtrlRequests = accessor _pendingCtrlRequests (\rep st -> st { _pendingCtrlRequests = rep })
1105+nextCtrlRequestId :: Accessor ValidLocalEndPointState ControlRequestId
1106+nextCtrlRequestId = accessor _nextCtrlRequestId (\cid st -> st { _nextCtrlRequestId = cid })
1108+nextConnectionId :: Accessor ValidLocalEndPointState ConnectionId
1109+nextConnectionId = accessor _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
1111+localConnections :: Accessor ValidLocalEndPointState (Map EndPointAddress RemoteEndPoint)
1112+localConnections = accessor _localConnections (\es st -> st { _localConnections = es })
1114+internalThreads :: Accessor ValidLocalEndPointState [ThreadId]
1115+internalThreads = accessor _internalThreads (\ts st -> st { _internalThreads = ts })
1117+localEndPointAt :: EndPointAddress -> Accessor ValidTransportState (Maybe LocalEndPoint)
1118+localEndPointAt addr = localEndPoints >>> DAC.mapMaybe addr
1120+pendingCtrlRequestsAt :: ControlRequestId -> Accessor ValidLocalEndPointState (Maybe (MVar [ByteString]))
1121+pendingCtrlRequestsAt ix = pendingCtrlRequests >>> DAC.intMapMaybe (fromIntegral ix)
1123+localConnectionTo :: EndPointAddress -> Accessor ValidLocalEndPointState (Maybe RemoteEndPoint)
1124+localConnectionTo addr = localConnections >>> DAC.mapMaybe addr
1126+remoteOutgoing :: Accessor ValidRemoteEndPointState Int
1127+remoteOutgoing = accessor _remoteOutgoing (\cs conn -> conn { _remoteOutgoing = cs })
1129+remoteIncoming :: Accessor ValidRemoteEndPointState IntSet
1130+remoteIncoming = accessor _remoteIncoming (\cs conn -> conn { _remoteIncoming = cs })