@@ -1,309 +1,1130 @@
1−{-# LANGUAGE CPP #-}
2−
3−module Network.Transport.TCP
4− ( mkTransport
5− , TCPConfig (..)
6− ) where
1+-- | TCP implementation of the transport layer.
2+--
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.
9+--
10+-- Applications that use the TCP transport should use
11+-- 'Network.Socket.withSocketsDo' in their main function for Windows
12+-- compatibility (see "Network.Socket").
13+--
14+-- TODOs:
15+-- * Output exception on channel after endpoint is closed
16+module Network.Transport.TCP ( -- * Main API
17+ createTransport
18+ , -- * TCP specific functionality (exported mostly for testing purposes_
19+ EndPointId
20+ , encodeEndPointAddress
21+ , decodeEndPointAddress
22+ , ControlHeader(..)
23+ , ConnectionRequestResponse(..)
24+ , firstNonReservedConnectionId
25+ , socketToEndPoint
26+ -- * Design notes
27+ -- $design
28+ ) where
729
30+import Prelude hiding (catch)
831import Network.Transport
9−
10−import Control.Applicative
32+import Network.Transport.Internal.TCP ( forkServer
33+ , recvWithLength
34+ , recvInt32
35+ )
36+import Network.Transport.Internal ( encodeInt32
37+ , decodeInt32
38+ , prependLength
39+ , mapExceptionIO
40+ , tryIO
41+ , tryToEnum
42+ , void
43+ )
44+import qualified Network.Socket as N ( HostName
45+ , ServiceName
46+ , Socket
47+ , sClose
48+ , getAddrInfo
49+ , socket
50+ , addrFamily
51+ , addrAddress
52+ , SocketType(Stream)
53+ , defaultProtocol
54+ , setSocketOption
55+ , SocketOption(ReuseAddr)
56+ , connect
57+ , sOMAXCONN
58+ , AddrInfo
59+ )
60+import Network.Socket.ByteString (sendMany)
1161import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
12−import Control.Concurrent.MVar
13−import Control.Exception (SomeException, IOException, AsyncException(ThreadKilled),
14− fromException, throwTo, throw, catch, handle)
15−import Control.Monad (forever, forM_)
62+import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
63+import Control.Concurrent.MVar ( MVar
64+ , newMVar
65+ , modifyMVar
66+ , modifyMVar_
67+ , readMVar
68+ , takeMVar
69+ , putMVar
70+ , newEmptyMVar
71+ , withMVar
72+ )
73+import Control.Category ((>>>))
74+import Control.Applicative ((<$>))
75+import Control.Monad (forM_, when, unless)
76+import Control.Exception (IOException, SomeException, handle, throw, try, bracketOnError)
77+import Data.IORef (IORef, newIORef, writeIORef, readIORef)
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)
86+import Data.Map (Map)
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)
1691
17−import qualified Data.ByteString.Internal as BSI
92+-- $design
93+--
94+-- [Goals]
95+--
96+-- The TCP transport maps multiple logical connections between /A/ and /B/ (in
97+-- either direction) to a single TCP connection:
98+--
99+-- > +-------+ +-------+
100+-- > | A |==========================| B |
101+-- > | |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~\ |
102+-- > | Q |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~Q |
103+-- > | \~~~|~~~~~~~~~~~~~~~~~~~~~~~~~<| |
104+-- > | |==========================| |
105+-- > +-------+ +-------+
106+--
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
110+-- closed.
111+--
112+-- [Connecting]
113+--
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).
118+--
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/.
123+--
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
130+-- tie.)
131+--
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.
135+--
136+-- [Disconnecting]
137+--
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/:
141+--
142+-- > +-------+ +-------+
143+-- > | A |==========================| B |
144+-- > | |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~\ |
145+-- > | | | Q |
146+-- > | #| | |
147+-- > | |==========================| |
148+-- > +-------+ +-------+
149+--
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
153+-- timeout).
154+--
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.)
161+--
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.
167+--
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.)
176+--
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.
18181
19−import Data.IntMap (IntMap)
20−import qualified Data.IntMap as IntMap
182+--------------------------------------------------------------------------------
183+-- Internal datatypes --
184+--------------------------------------------------------------------------------
21185
22−import Data.Int (Int32)
23−import qualified Data.Serialize as Ser
24−import Data.Word (Word8)
25−
26−import Foreign.ForeignPtr (withForeignPtr)
27−import Foreign.Ptr (Ptr, plusPtr)
28−import Foreign.Storable (sizeOf)
29−
30−import Network.Socket
31− ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
32− , SocketType (Stream), SocketOption (ReuseAddr)
33− , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
34− , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
35−import qualified Network.Socket as N
36−
37−import Text.Printf
38−
39−import Safe
40−import System.IO (stderr, hPutStrLn)
41−import Foreign.Storable (pokeByteOff, peekByteOff)
42−import Foreign.C (CInt(..))
43−import Foreign.ForeignPtr (withForeignPtr)
44−
45−#ifndef LAZY
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
50−encode = Ser.encode
51−decode = Ser.decode
52−#else
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
59−#endif
60−{-# INLINE encode #-}
61−{-# INLINE decode #-}
62−encode :: Ser.Serialize a => a -> ByteString
63−decode :: Ser.Serialize a => ByteString -> Either String a
64−
65−foreign import ccall unsafe "htonl" htonl :: CInt -> CInt
66−foreign import ccall unsafe "ntohl" ntohl :: CInt -> CInt
67−
68−type ChanId = Int
69−type Chans = MVar (ChanId, IntMap (MVar [ByteString], [(ThreadId, Socket)]))
70−
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
75−-- a port number
76−data TCPConfig = TCPConfig Hints HostName ServiceName
77−
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] }))
88− Nothing
89− (Just service)
90− let serverAddr = case serverAddrs of
91− [] -> error "mkTransport: getAddrInfo returned []"
92− as -> head as
93− sock <- socket (addrFamily serverAddr) Stream defaultProtocol
94− setSocketOption sock ReuseAddr 1
95− bindSocket sock (addrAddress serverAddr)
96− listen sock 5
97− threadId <- forkWithExceptions forkIO "Connection Listener" $
98− procConnections chans sock
99−
100− return Transport
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)
110− id $ decode bs in
111− Just $ mkSourceAddr host service chanId
112− , closeTransport = {-# SCC "closeTransport" #-} do
113− -- Kill the transport channel process
114− killThread threadId
115− sClose sock
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'
121− sClose sock'))
122− }
186+-- We use underscores for fields that we might update (using accessores)
187+--
188+-- All data types follow the same structure:
189+--
190+-- * A top-level data type describing static properties (TCPTransport,
191+-- LocalEndPoint, RemoteEndPoint)
192+-- * The 'static' properties include an MVar containing a data structure for
193+-- the dynamic properties (TransportState, LocalEndPointState,
194+-- RemoteEndPointState). The state could be invalid/valid/closed,/etc.
195+-- * For the case of "valid" we use third data structure to give more details
196+-- about the state (ValidTransportState, ValidLocalEndPointState,
197+-- ValidRemoteEndPointState).
123198
124−mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
125−mkSourceAddr host service chanId = SourceAddr
126− { connectWith = {-# SCC "connectWith" #-} \_ -> mkSourceEnd host service chanId
127− , serialize = {-# SCC "serialize" #-} encode (host, service, chanId)
199+data TCPTransport = TCPTransport
200+ { transportHost :: N.HostName
201+ , transportPort :: N.ServiceName
202+ , transportState :: MVar TransportState
128203 }
129204
130−mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
131−mkSourceEnd host service chanId = withSocketsDo $ do
132− let err m = error $ printf "mkSourceEnd: %s" m
133− serverAddrs <- getAddrInfo Nothing (Just host) (Just service)
134− let serverAddr = case serverAddrs of
135− [] -> err "getAddrInfo returned []"
136− as -> head as
137− sock <- socket (addrFamily serverAddr) Stream defaultProtocol
138− setSocketOption sock ReuseAddr 1
139− N.connect sock (addrAddress serverAddr)
140− NBS.sendAll sock $ encode (fromIntegral chanId :: Int32)
141− return SourceEnd
142− { send = {-# SCC "send" #-} \bss ->
143− let size = fromIntegral (sum . map BS.length $ bss) :: Int32 in
144− if size < 255
145− then sendNBS sock (BS.singleton (toEnum $ fromEnum size) : bss)
146− else do size' <- encodeLength size
147− sendNBS sock (BS.singleton (toEnum 255) : size' : bss)
148− , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
149− }
205+data TransportState =
206+ TransportValid ValidTransportState
207+ | TransportClosed
208+
209+data ValidTransportState = ValidTransportState
210+ { _localEndPoints :: Map EndPointAddress LocalEndPoint
211+ }
212+
213+data LocalEndPoint = LocalEndPoint
214+ { localAddress :: EndPointAddress
215+ , localChannel :: Chan Event
216+ , localState :: MVar LocalEndPointState
217+ }
218+
219+data LocalEndPointState =
220+ LocalEndPointValid ValidLocalEndPointState
221+ | LocalEndPointClosed
222+
223+data ValidLocalEndPointState = ValidLocalEndPointState
224+ { _nextConnectionId :: ConnectionId
225+ , _pendingCtrlRequests :: IntMap (MVar [ByteString])
226+ , _nextCtrlRequestId :: ControlRequestId
227+ , _localConnections :: Map EndPointAddress RemoteEndPoint
228+ , _internalThreads :: [ThreadId]
229+ }
230+
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.
234+--
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.
239+--
240+-- We therefore insist that:
241+--
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).
255+--
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.
266+--
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.
269+--
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.
275+
276+data RemoteEndPoint = RemoteEndPoint
277+ { remoteAddress :: EndPointAddress
278+ , remoteState :: MVar RemoteState
279+ }
280+
281+data RemoteState =
282+ RemoteEndPointInvalid (TransportError ConnectErrorCode)
283+ | RemoteEndPointValid ValidRemoteEndPointState
284+ | RemoteEndPointClosing (MVar ()) ValidRemoteEndPointState
285+ | RemoteEndPointClosed
150286
151−mkTargetEnd :: Chans -> ChanId -> MVar [ByteString] -> TargetEnd
152−mkTargetEnd chans chanId chan = TargetEnd
153− { -- for now we will implement this as an MVar
154− receive = {-# SCC "receive" #-} takeMVar chan
155− , closeTargetEnd = {-# SCC "closeTargetEnd" #-} do
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
161− killThread threadId
162− sClose sock
163− let chanMap' = IntMap.delete chanId chanMap
164− putMVar chans (chanId', chanMap')
287+data ValidRemoteEndPointState = ValidRemoteEndPointState
288+ { _remoteOutgoing :: !Int
289+ , _remoteIncoming :: IntSet
290+ , remoteSocket :: N.Socket
291+ , sendOn :: [ByteString] -> IO ()
165292 }
166293
167−-- | This function waits for inbound connections. If a connection fails
168−-- for some reason, an error is raised.
169−procConnections :: Chans -> Socket -> IO ()
170−procConnections chans sock = forever $ do
171− let err m = error $ printf "procConnections: %s" m
172− (clientSock, _clientAddr) <- accept sock
173− -- decode the first message to find the correct chanId
174− bss <- recvExact clientSock (fromIntegral $ sizeOf (0 :: Int32))
175− case bss of
176− Left _ -> err "inbound chanId aborted"
177− Right bs -> do
178− let chanId = either err fromIntegral (decode bs :: Either String Int32)
179− (chanId', chanMap) <- takeMVar chans
180− case IntMap.lookup chanId chanMap of
181− Nothing -> do
182− putMVar chans (chanId', chanMap)
183− err "cannot find chanId"
184− Just (chan, socks) -> do
185− threadId <- forkWithExceptions forkIO "Message Listener" $
186− procMessages chans chanId chan clientSock
187− let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
188− putMVar chans (chanId', chanMap')
189−
190−-- | This function first extracts a header of type Word8, which determines
191−-- the size of the ByteString that follows. If this size is 0, this indicates
192−-- that the ByteString is large, so the next value is an Int32, which
193−-- determines the size of the ByteString that follows. The ByteString is then
194−-- extracted from the socket, and then written to the Chan only when
195−-- complete. If either of the first header size is null this indicates the
196−-- socket has closed.
197−procMessages :: Chans -> ChanId -> MVar [ByteString] -> Socket -> IO ()
198−procMessages chans chanId chan sock = do
199− esizeBS <- recvExact sock 1
200− case esizeBS of
201− Left _ -> closeSocket
202− Right sizeBS ->
203− let size = toEnum . fromEnum $ BS.index sizeBS 0 in
204− if size < 255
205− then procMessage size
206− else do esizeBS' <- recvExact sock (fromIntegral $ sizeOf (0 :: Int32))
207− case esizeBS' of
208− Left _ -> closeSocket
209− Right sizeBS' -> do sizeDec <- decodeLength sizeBS'
210− procMessage sizeDec
211− where
212− err m = error $ printf "procMessages: %s" m
213− closeSocket :: IO ()
214− closeSocket = do
215− (chanId', chanMap) <- takeMVar chans
216− case IntMap.lookup chanId chanMap of
217− Nothing -> do
218− putMVar chans (chanId', chanMap)
219− err "chanId not found."
220− Just (chan, socks) -> do
221− let socks' = filter ((/= sock) . snd) socks
222− let chanMap' = IntMap.insert chanId (chan, socks') chanMap
223− putMVar chans (chanId', chanMap')
224− sClose sock
225− procMessage :: Int32 -> IO ()
226− procMessage size = do
227− ebs <- recvExact sock size
228− case ebs of
229− Left _ -> closeSocket
230− Right bs -> do
231− putMVar chan [bs]
232− procMessages chans chanId chan sock
233−
234−-- | The normal result of `recvExact sock n` is `Right ByteString`
235−-- whose string of length `n`, received from `sock`. If fewer than `n`
236−-- bytes are read from `sock` before it closes, the result is `Left
237−-- ByteString` whose string is those bytes that were received. No more
238−-- bytes than necessary are read from the socket. NB: This uses
239−-- Network.Socket.ByteString.recv, which may *discard* superfluous
240−-- input depending on the socket type.
241−#ifndef LAZY
242−recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
243−recvExact sock l = interceptAllExn "recvExact" $ do
244− res <- BSI.createAndTrim (fromIntegral l) (go 0)
245− case BS.length res of
246− n | n == (fromIntegral l) -> return $ Right res
247− n -> return $ Left res
248− where
249− go :: Int -> Ptr Word8 -> IO Int
250− go n ptr | n == (fromIntegral l) = return n
251− go n ptr = do
252− (p, off, len) <- BSI.toForeignPtr <$> NBS.recv sock (min (fromIntegral l-n) 4096)
253− if len == 0
254− then return n
255− else withForeignPtr p $ \p -> do
256− BSI.memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)
257− go (n+len) ptr
258−#else
259−recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
260−recvExact sock n =
261− interceptAllExn "recvExact" $
262− go [] sock n
263− where
264− go :: [ByteString] -> Socket -> Int32 -> IO (Either ByteString ByteString)
265− go bss _ 0 = return (Right $ BS.concat (reverse bss))
266− go bss sock n = do
267− bs <- NBS.recv sock (min (fromIntegral n) 4096)
268− if BS.null bs
269− then return (Left $ BS.concat (reverse bss))
270− else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
271−#endif
272−
273−interceptAllExn msg =
274− Control.Exception.handle $ \ e ->
275− case fromException e of
276− Just ThreadKilled -> throw e
277− Nothing -> do
278− BS.hPutStrLn stderr $ BS.pack$ "Exception inside "++msg++": "++show e
279− throw e
280−-- throw (e :: SomeException)
281−
282−forkWithExceptions :: (IO () -> IO ThreadId) -> String -> IO () -> IO ThreadId
283−forkWithExceptions forkit descr action = do
284− parent <- myThreadId
285− forkit $
286− Control.Exception.catch action
287− (\ e -> do
288− case fromException e of
289− Just ThreadKilled ->
290−-- BSS.hPutStrLn stderr $ BSS.pack$ "Note: Child thread killed: "++descr
291− return ()
292− _ -> do
293− BS.hPutStrLn stderr $ BS.pack$ "Exception inside child thread "++descr++": "++show e
294− throwTo parent (e::SomeException)
295− )
296−
297−-- | Encode length (manual for now)
298−encodeLength :: Int32 -> IO ByteString
299−encodeLength i32 =
300− BSI.create 4 $ \p ->
301− pokeByteOff p 0 (htonl (fromIntegral i32))
302−
303−-- | Decode length (manual for now)
304−decodeLength :: ByteString -> IO Int32
305−decodeLength bs =
306− let (fp, _, _) = BSI.toForeignPtr bs in
307− withForeignPtr fp $ \p -> do
308− w32 <- peekByteOff p 0
309− return (fromIntegral (ntohl w32))
294+type EndPointId = Int32
295+type ControlRequestId = Int32
296+type EndPointPair = (LocalEndPoint, RemoteEndPoint)
297+
298+-- | Control headers
299+data ControlHeader =
300+ RequestConnectionId -- ^ Request a new connection ID from the remote endpoint
301+ | CloseConnection -- ^ Tell the remote endpoint we will no longer be using a connection
302+ | ControlResponse -- ^ Respond to a control request /from/ the remote endpoint
303+ | CloseSocket -- ^ Request to close the connection (see module description)
304+ deriving (Enum, Bounded, Show)
305+
306+-- Response sent by /B/ to /A/ when /A/ tries to connect
307+data ConnectionRequestResponse =
308+ ConnectionRequestAccepted -- ^ /B/ accepts the connection
309+ | ConnectionRequestEndPointInvalid -- ^ /A/ requested an invalid endpoint
310+ | ConnectionRequestCrossed -- ^ /A/s request crossed with a request from /B/ (see protocols)
311+ deriving (Enum, Bounded, Show)
312+
313+--------------------------------------------------------------------------------
314+-- Top-level functionality --
315+--------------------------------------------------------------------------------
316+
317+-- | Create a TCP transport
318+--
319+-- TODOs: deal with hints
320+createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)
321+createTransport host port = do
322+ state <- newMVar . TransportValid $ ValidTransportState { _localEndPoints = Map.empty }
323+ let transport = TCPTransport { transportState = state
324+ , transportHost = host
325+ , transportPort = port
326+ }
327+ tryIO $ do
328+ -- For a discussion of the use of N.sOMAXCONN, see
329+ -- http://tangentsoft.net/wskfaq/advanced.html#backlog
330+ -- http://www.linuxjournal.com/files/linuxjournal.com/linuxjournal/articles/023/2333/2333s2.html
331+ transportThread <- forkServer host port N.sOMAXCONN (terminationHandler transport) (handleConnectionRequest transport)
332+ return Transport { newEndPoint = apiNewEndPoint transport
333+ , closeTransport = killThread transportThread -- This will invoke the termination handler
334+ }
335+ where
336+ terminationHandler :: TCPTransport -> SomeException -> IO ()
337+ terminationHandler transport _ = do
338+ -- TODO: we currently don't make a distinction between endpoint failure and manual closure
339+ mTSt <- modifyMVar (transportState transport) $ \st -> case st of
340+ TransportValid vst -> return (TransportClosed, Just vst)
341+ TransportClosed -> return (TransportClosed, Nothing)
342+ case mTSt of
343+ Nothing ->
344+ -- Transport already closed
345+ return ()
346+ Just tSt ->
347+ mapM_ (apiCloseEndPoint transport) (Map.elems $ tSt ^. localEndPoints)
348+
349+--------------------------------------------------------------------------------
350+-- API functions --
351+--------------------------------------------------------------------------------
352+
353+-- | Create a new endpoint
354+apiNewEndPoint :: TCPTransport -> IO (Either (TransportError NewEndPointErrorCode) EndPoint)
355+apiNewEndPoint transport = try $ do
356+ ourEndPoint <- createLocalEndPoint transport
357+ return EndPoint
358+ { receive = readChan (localChannel ourEndPoint)
359+ , address = localAddress ourEndPoint
360+ , connect = apiConnect ourEndPoint
361+ , closeEndPoint = apiCloseEndPoint transport ourEndPoint
362+ , newMulticastGroup = return . Left $ newMulticastGroupError
363+ , resolveMulticastGroup = return . Left . const resolveMulticastGroupError
364+ }
365+ where
366+ newMulticastGroupError = TransportError NewMulticastGroupUnsupported "TCP does not support multicast"
367+ resolveMulticastGroupError = TransportError ResolveMulticastGroupUnsupported "TCP does not support multicast"
368+
369+-- | Connnect to an endpoint
370+apiConnect :: LocalEndPoint -- ^ Local end point
371+ -> EndPointAddress -- ^ Remote address
372+ -> Reliability -- ^ Reliability (ignored)
373+ -> IO (Either (TransportError ConnectErrorCode) Connection)
374+apiConnect ourEndPoint theirAddress _ | localAddress ourEndPoint == theirAddress =
375+ connectToSelf ourEndPoint
376+apiConnect ourEndPoint theirAddress _ = try $ do
377+ (theirEndPoint, connId) <- requestConnectionTo ourEndPoint theirAddress
378+ -- connAlive can be an IORef rather than an MVar because it is protected by
379+ -- the remoteState MVar. We don't need the overhead of locking twice.
380+ connAlive <- newIORef True
381+ return $ Connection { send = apiSend theirEndPoint connId connAlive
382+ , close = apiClose theirEndPoint connId connAlive
383+ }
384+
385+-- | Close a connection
386+--
387+-- RELY: Remote endpoint must be in 'RemoteEndPointValid' or 'RemoteEndPointClosed'
388+-- GUARANTEE: If the connection is alive on entry then the remote endpoint will
389+-- either be RemoteEndPointValid or RemoteEndPointClosing. Otherwise, the state
390+-- of the remote endpoint will not be changed.
391+--
392+-- TODO: We ignore errors during a close. Is that right?
393+apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()
394+apiClose theirEndPoint connId connAlive = void . tryIO $ do
395+ modifyMVar_ (remoteState theirEndPoint) $ \st -> case st of
396+ RemoteEndPointValid vst -> do
397+ alive <- readIORef connAlive
398+ if alive
399+ then do
400+ writeIORef connAlive False
401+ sendOn vst [encodeInt32 CloseConnection, encodeInt32 connId]
402+ return (RemoteEndPointValid . (remoteOutgoing ^: (+ 1)) $ vst)
403+ else
404+ return (RemoteEndPointValid vst)
405+ RemoteEndPointClosed -> do
406+ return st
407+ RemoteEndPointClosing _ _ ->
408+ fail "apiClose RELY violation"
409+ RemoteEndPointInvalid _ ->
410+ fail "apiClose RELY violation"
411+ closeIfUnused theirEndPoint
412+
413+-- | Send data across a connection
414+--
415+-- RELY: The remote endpoint must be in state 'RemoteEndPointValid' or 'RemoteEndPointClosed'
416+-- GUARANTEE: The state of the remote endpoint will not be changed.
417+apiSend :: RemoteEndPoint -> ConnectionId -> IORef Bool -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
418+apiSend theirEndPoint connId connAlive payload = do
419+ withMVar (remoteState theirEndPoint) $ \st -> case st of
420+ RemoteEndPointValid vst -> do
421+ alive <- readIORef connAlive
422+ try $ if alive
423+ then mapExceptionIO sendFailed $ sendOn vst (encodeInt32 connId : prependLength payload)
424+ else throw $ TransportError SendFailed "Connection closed"
425+ RemoteEndPointClosed -> do
426+ return (Left $ TransportError SendFailed "Endpoint closed")
427+ RemoteEndPointClosing _ _ ->
428+ error "apiSend RELY violation"
429+ RemoteEndPointInvalid _ ->
430+ error "apiSend RELY violation"
431+ where
432+ sendFailed :: IOException -> TransportError SendErrorCode
433+ sendFailed = TransportError SendFailed . show
434+
435+-- | Force-close the endpoint
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 ->
442+ case st of
443+ LocalEndPointValid vst ->
444+ return (LocalEndPointClosed, Just vst)
445+ LocalEndPointClosed ->
446+ return (LocalEndPointClosed, Nothing)
447+ case mOurState of
448+ Nothing -> -- Already closed
449+ return ()
450+ Just vst -> do
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
459+ where
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 ->
465+ case st of
466+ RemoteEndPointInvalid _ ->
467+ return st
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
480+
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
491+ }
492+ where
493+ ourChan :: Chan Event
494+ ourChan = localChannel ourEndPoint
495+
496+ selfSend :: MVar Bool -> ConnectionId -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
497+ selfSend connAlive connId msg = do
498+ modifyMVar connAlive $ \alive ->
499+ if alive
500+ then do
501+ writeChan ourChan (Received connId msg)
502+ return (alive, Right ())
503+ else
504+ return (alive, Left (TransportError SendFailed "Connection closed"))
505+
506+ selfClose :: MVar Bool -> ConnectionId -> IO ()
507+ selfClose connAlive connId = do
508+ modifyMVar_ connAlive $ \alive -> do
509+ when alive $ writeChan ourChan (ConnectionClosed connId)
510+ return False
511+
512+--------------------------------------------------------------------------------
513+-- Lower level functionality --
514+--------------------------------------------------------------------------------
515+
516+-- | Create a new local endpoint
517+--
518+-- May throw a TransportError NewEndPointErrorCode exception if the transport is closed.
519+createLocalEndPoint :: TCPTransport -> IO LocalEndPoint
520+createLocalEndPoint transport = do
521+ chan <- newChan
522+ state <- newMVar . LocalEndPointValid $ ValidLocalEndPointState
523+ { _nextConnectionId = firstNonReservedConnectionId
524+ , _pendingCtrlRequests = IntMap.empty
525+ , _nextCtrlRequestId = 0
526+ , _localConnections = Map.empty
527+ , _internalThreads = []
528+ }
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
536+ }
537+ return (TransportValid . (localEndPointAt addr ^= Just localEndPoint) $ vst, localEndPoint)
538+ TransportClosed ->
539+ throw (TransportError NewEndPointFailed "Transport closed")
540+
541+-- | Request a connection to a remote endpoint
542+--
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.
546+--
547+-- May throw a TransportError ConnectErrorCode exception.
548+requestConnectionTo :: LocalEndPoint
549+ -> EndPointAddress
550+ -> IO (RemoteEndPoint, ConnectionId)
551+requestConnectionTo ourEndPoint theirAddress = go
552+ where
553+ go = do
554+ -- Find the remote endpoint (create it if it doesn't yet exist)
555+ theirEndPoint <- findTheirEndPoint
556+ let theirState = remoteState theirEndPoint
557+
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 ->
563+ case st of
564+ RemoteEndPointValid ep ->
565+ return (RemoteEndPointValid . (remoteOutgoing ^: (+ 1)) $ ep, st)
566+ _ ->
567+ return (st, st)
568+
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
572+ -- endpoint lock.
573+ --
574+ -- Although 'endPointStateSnapshot' here refers to a snapshot of the
575+ -- endpoint state, and might have changed in the meantime, these changes
576+ -- won't matter.
577+ case endPointStateSnapshot of
578+ RemoteEndPointInvalid err -> do
579+ throw err
580+
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
585+
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.
591+ go
592+
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)
604+
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
607+ -- point)
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)
612+
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
616+ -- ID)
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))
626+ Nothing -> do
627+ theirState <- newEmptyMVar
628+ let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
629+ , remoteState = theirState
630+ }
631+ return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) $ vst, (theirEndPoint, True))
632+
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
635+ -- here.
636+ when isNew . void . forkEndPointThread ourEndPoint $
637+ setupRemoteEndPoint (ourEndPoint, theirEndPoint)
638+ return theirEndPoint
639+
640+ ourState :: MVar LocalEndPointState
641+ ourState = localState ourEndPoint
642+
643+-- | Set up a remote endpoint
644+--
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)
651+ case result of
652+ Right (sock, ConnectionRequestAccepted) -> do
653+ let vst = ValidRemoteEndPointState
654+ { remoteSocket = sock
655+ , _remoteOutgoing = 0
656+ , _remoteIncoming = IntSet.empty
657+ , sendOn = sendMany sock
658+ }
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"))
669+ N.sClose sock
670+ Right (sock, ConnectionRequestCrossed) -> do
671+ N.sClose sock
672+ Left err -> do
673+ -- See comment above
674+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
675+ putMVar theirState (RemoteEndPointInvalid err)
676+ where
677+ theirState = remoteState theirEndPoint
678+ invalidAddress = TransportError ConnectNotFound
679+
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])
694+ recvInt32 sock
695+ case tryToEnum response of
696+ Nothing -> throw (failed . userError $ "Unexpected response")
697+ Just r -> return (sock, r)
698+ where
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
703+ return sock
704+
705+ invalidAddress, insufficientResources, failed :: IOException -> TransportError ConnectErrorCode
706+ invalidAddress = TransportError ConnectNotFound . show
707+ insufficientResources = TransportError ConnectInsufficientResources . show
708+ failed = TransportError ConnectFailed . show
709+
710+-- | Remove reference to a remote endpoint from a local endpoint
711+--
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
720+
721+-- | Remove reference to a local endpoint from the transport state
722+--
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)
729+ TransportClosed ->
730+ return TransportClosed
731+
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
736+
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)
743+ _ -> Nothing
744+ _ ->
745+ Nothing
746+
747+-- | Do a (blocking) remote request
748+--
749+-- RELY: Remote endpoint must be in valid state.
750+-- GUARANTEE: Will not change the state of the remote endpoint.
751+--
752+-- May throw IO (user) exception if the local endpoint is closed or if the send
753+-- fails.
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]
765+ takeMVar reply
766+
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 ->
770+ case st of
771+ RemoteEndPointValid vst -> do
772+ if vst ^. remoteOutgoing == 0 && IntSet.null (vst ^. remoteIncoming)
773+ then do
774+ resolved <- newEmptyMVar
775+ sendOn vst [encodeInt32 CloseSocket]
776+ return (RemoteEndPointClosing resolved vst)
777+ else
778+ return st
779+ _ -> do
780+ return st
781+
782+-- | Fork a new thread and store its ID as part of the transport state
783+--
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)
794+ where
795+ removeThread :: IO ()
796+ removeThread = do
797+ tid <- myThreadId
798+ modifyMVar_ ourState $ \st -> case st of
799+ LocalEndPointValid vst ->
800+ return (LocalEndPointValid . (internalThreads ^: filter (/= tid)) $ vst)
801+ LocalEndPointClosed ->
802+ return LocalEndPointClosed
803+
804+ ourState :: MVar LocalEndPointState
805+ ourState = localState ourEndPoint
806+
807+--------------------------------------------------------------------------------
808+-- Incoming requests --
809+--------------------------------------------------------------------------------
810+
811+-- | Handle a connection request (that is, a remote endpoint that is trying to
812+-- establish a TCP connection with us)
813+--
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
830+ Nothing -> do
831+ sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
832+ fail "Invalid endpoint"
833+ Just ourEndPoint ->
834+ return ourEndPoint
835+ TransportClosed -> do
836+ fail "Transport closed"
837+ void . forkEndPointThread ourEndPoint $ go ourEndPoint theirAddress
838+ where
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
845+ Nothing -> do
846+ theirState <- newEmptyMVar
847+ let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
848+ , remoteState = theirState
849+ }
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"
855+ if crossed
856+ then do
857+ tryIO $ sendMany sock [encodeInt32 ConnectionRequestCrossed]
858+ tryIO $ N.sClose sock
859+ return Nothing
860+ else do
861+ let vst = ValidRemoteEndPointState
862+ { remoteSocket = sock
863+ , _remoteOutgoing = 0
864+ , _remoteIncoming = IntSet.empty
865+ , sendOn = sendMany sock
866+ }
867+ -- TODO: this putMVar might block if the remote endpoint sends a
868+ -- connection request for a local endpoint that it is already
869+ -- connected to
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.
878+ case mEndPoint of
879+ Nothing ->
880+ return ()
881+ Just theirEndPoint ->
882+ handleIncomingMessages (ourEndPoint, theirEndPoint)
883+
884+ tryCloseSocket :: IOException -> IO ()
885+ tryCloseSocket _ = void . tryIO $ N.sClose sock
886+
887+ invalidEndPoint :: IOException -> IO ()
888+ invalidEndPoint ex = do
889+ tryIO $ sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
890+ tryCloseSocket ex
891+
892+-- | Handle requests from a remote endpoint.
893+--
894+-- Returns only if the remote party closes the socket or if an error occurs.
895+--
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 ->
904+ case st of
905+ RemoteEndPointValid ep ->
906+ return (remoteSocket ep)
907+ RemoteEndPointClosing _ ep ->
908+ return (remoteSocket ep)
909+ _ ->
910+ error "handleIncomingMessages RELY violation"
911+
912+ mCleanExit <- tryIO $ go sock
913+ case mCleanExit of
914+ Right () -> return ()
915+ Left err -> prematureExit sock err
916+ where
917+ -- Dispatch
918+ go :: N.Socket -> IO ()
919+ go sock = do
920+ connId <- recvInt32 sock
921+ if connId >= firstNonReservedConnectionId
922+ then do
923+ readMessage sock connId
924+ go sock
925+ else do
926+ case tryToEnum (fromIntegral connId) of
927+ Just RequestConnectionId -> do
928+ recvInt32 sock >>= createNewConnection
929+ go sock
930+ Just ControlResponse -> do
931+ recvInt32 sock >>= readControlResponse sock
932+ go sock
933+ Just CloseConnection -> do
934+ recvInt32 sock >>= closeConnection
935+ go sock
936+ Just CloseSocket -> do
937+ didClose <- closeSocket sock
938+ unless didClose $ go sock
939+ Nothing ->
940+ -- Invalid control request, exit
941+ hPutStrLn stderr "Warning: invalid control request"
942+
943+ -- Create a new connection
944+ createNewConnection :: ControlRequestId -> IO ()
945+ createNewConnection reqId = do
946+ newId <- getNextConnectionId ourEndPoint
947+ modifyMVar_ theirState $ \st -> do
948+ vst <- case st of
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
958+ putMVar resolved ()
959+ return (remoteIncoming ^= IntSet.singleton newId $ vst)
960+ _ ->
961+ error "handleIncomingMessages RELY violation"
962+ sendOn vst ( encodeInt32 ControlResponse
963+ : encodeInt32 reqId
964+ : prependLength [encodeInt32 newId]
965+ )
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)
970+
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"
980+ case mmvar of
981+ Nothing -> do
982+ hPutStrLn stderr $ "Warning: Invalid request ID"
983+ return () -- Invalid request ID. TODO: We just ignore it?
984+ Just mvar ->
985+ putMVar mvar response
986+
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
995+
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 ->
1009+ case st of
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
1019+ then do
1020+ -- Attempt to reply (but don't insist)
1021+ tryIO $ sendOn vst' [encodeInt32 CloseSocket]
1022+ resolved <- newEmptyMVar
1023+ return (RemoteEndPointClosing resolved vst', Just resolved)
1024+ else
1025+ return (RemoteEndPointValid vst', Nothing)
1026+ RemoteEndPointClosing resolved _ ->
1027+ return (st, Just resolved)
1028+ _ ->
1029+ error "handleIncomingConnections RELY violation"
1030+
1031+ case canClose of
1032+ Nothing ->
1033+ return False
1034+ Just resolved -> do
1035+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
1036+ modifyMVar_ theirState $ return . const RemoteEndPointClosed
1037+ N.sClose sock
1038+ putMVar resolved ()
1039+ return True
1040+
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
1044+
1045+ -- Arguments
1046+ ourChannel = localChannel ourEndPoint
1047+ ourState = localState ourEndPoint
1048+ theirState = remoteState theirEndPoint
1049+ theirAddr = remoteAddress theirEndPoint
1050+
1051+ -- Deal with a premature exit
1052+ prematureExit :: N.Socket -> IOException -> IO ()
1053+ prematureExit sock err = do
1054+ N.sClose sock
1055+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
1056+ mUnclosedConnections <- modifyMVar theirState $ \st ->
1057+ case st of
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)
1066+
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
1070+ Nothing ->
1071+ return ()
1072+ Just unclosedConnections -> writeChan ourChannel . ErrorEvent $
1073+ TransportError (EventConnectionLost (remoteAddress theirEndPoint) (IntSet.elems unclosedConnections)) (show err)
1074+
1075+-- | Get the next connection ID
1076+--
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"
1086+
1087+--------------------------------------------------------------------------------
1088+-- Constants --
1089+--------------------------------------------------------------------------------
1090+
1091+-- | We reserve a bunch of connection IDs for control messages
1092+firstNonReservedConnectionId :: ConnectionId
1093+firstNonReservedConnectionId = 1024
1094+
1095+--------------------------------------------------------------------------------
1096+-- Accessor definitions --
1097+--------------------------------------------------------------------------------
1098+
1099+localEndPoints :: Accessor ValidTransportState (Map EndPointAddress LocalEndPoint)
1100+localEndPoints = accessor _localEndPoints (\es st -> st { _localEndPoints = es })
1101+
1102+pendingCtrlRequests :: Accessor ValidLocalEndPointState (IntMap (MVar [ByteString]))
1103+pendingCtrlRequests = accessor _pendingCtrlRequests (\rep st -> st { _pendingCtrlRequests = rep })
1104+
1105+nextCtrlRequestId :: Accessor ValidLocalEndPointState ControlRequestId
1106+nextCtrlRequestId = accessor _nextCtrlRequestId (\cid st -> st { _nextCtrlRequestId = cid })
1107+
1108+nextConnectionId :: Accessor ValidLocalEndPointState ConnectionId
1109+nextConnectionId = accessor _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
1110+
1111+localConnections :: Accessor ValidLocalEndPointState (Map EndPointAddress RemoteEndPoint)
1112+localConnections = accessor _localConnections (\es st -> st { _localConnections = es })
1113+
1114+internalThreads :: Accessor ValidLocalEndPointState [ThreadId]
1115+internalThreads = accessor _internalThreads (\ts st -> st { _internalThreads = ts })
1116+
1117+localEndPointAt :: EndPointAddress -> Accessor ValidTransportState (Maybe LocalEndPoint)
1118+localEndPointAt addr = localEndPoints >>> DAC.mapMaybe addr
1119+
1120+pendingCtrlRequestsAt :: ControlRequestId -> Accessor ValidLocalEndPointState (Maybe (MVar [ByteString]))
1121+pendingCtrlRequestsAt ix = pendingCtrlRequests >>> DAC.intMapMaybe (fromIntegral ix)
1122+
1123+localConnectionTo :: EndPointAddress -> Accessor ValidLocalEndPointState (Maybe RemoteEndPoint)
1124+localConnectionTo addr = localConnections >>> DAC.mapMaybe addr
1125+
1126+remoteOutgoing :: Accessor ValidRemoteEndPointState Int
1127+remoteOutgoing = accessor _remoteOutgoing (\cs conn -> conn { _remoteOutgoing = cs })
1128+
1129+remoteIncoming :: Accessor ValidRemoteEndPointState IntSet
1130+remoteIncoming = accessor _remoteIncoming (\cs conn -> conn { _remoteIncoming = cs })