@@ -0,0 +1,210 @@
1+-- | TCP transport
2+module Network.Transport.TCP where
3+
4+import Network.Transport
5+import Network.Transport.Internal ( encodeInt16
6+ , decodeInt16
7+ , encodeInt32
8+ , decodeInt32
9+ , maybeToErrorT
10+ )
11+import qualified Network.Socket as N ( HostName
12+ , ServiceName
13+ , Socket
14+ , SocketType(Stream)
15+ , SocketOption(ReuseAddr)
16+ , getAddrInfo
17+ , defaultHints
18+ , socket
19+ , bindSocket
20+ , listen
21+ , addrFamily
22+ , addrAddress
23+ , defaultProtocol
24+ , setSocketOption
25+ , connect
26+ , sClose
27+ , accept
28+ )
29+import qualified Network.Socket.ByteString as NSB (sendMany, recv)
30+import Control.Concurrent (forkIO, ThreadId)
31+import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
32+import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, readMVar)
33+import Control.Category ((>>>))
34+import Control.Applicative ((<*>), (*>), (<$>), pure)
35+import Control.Monad (forever)
36+import Control.Monad.Error (ErrorT, liftIO, runErrorT)
37+import Data.ByteString (ByteString)
38+import qualified Data.ByteString as BS (length, concat)
39+import qualified Data.ByteString.Char8 as BSC (pack, unpack)
40+import Data.IntMap (IntMap)
41+import qualified Data.IntMap as IntMap (empty, size)
42+import Data.Lens.Lazy (Lens, lens, intMapLens, (^.), (^=), (^+=))
43+import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
44+
45+data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState) }
46+data EndPointState = EndPointState { _nextConnectionIx :: ConnectionIx }
47+type EndPointIx = Int
48+type ConnectionIx = Int
49+
50+endPoints :: Lens TransportState (IntMap (Chan Event, MVar EndPointState))
51+endPoints = lens _endPoints (\es st -> st { _endPoints = es })
52+
53+nextConnectionIx :: Lens EndPointState ConnectionIx
54+nextConnectionIx = lens _nextConnectionIx (\cix st -> st { _nextConnectionIx = cix })
55+
56+endPointAt :: Int -> Lens TransportState (Maybe (Chan Event, MVar EndPointState))
57+endPointAt ix = endPoints >>> intMapLens ix
58+
59+-- | Create a TCP transport
60+--
61+-- Hostname and service name should be "canonical" for the transport layer to
62+-- use TCP channels bidirectionally (otherwise when A connects to B and B
63+-- connects to A two socket pairs will be created that will both be used
64+-- unidirectionally, at degraded performance).
65+--
66+-- TODOs:
67+-- * Perhaps we should allow to use something different other than 'defaultHints'
68+-- * Deal with all exceptions that may occur
69+createTransport :: N.HostName -> N.ServiceName -> IO Transport
70+createTransport host port = do
71+ state <- newMVar $ TransportState { _endPoints = IntMap.empty }
72+ forkServer host port (transportServer state)
73+ return Transport { newEndPoint = tcpNewEndPoint state host port }
74+
75+-- | Start a server at the specified address
76+forkServer :: N.HostName -> N.ServiceName -> (N.Socket -> IO ()) -> IO ThreadId
77+forkServer host port server = do
78+ -- Resolve the specified address. By specification, getAddrInfo will never
79+ -- return an empty list (but will throw an exception instead) and will return
80+ -- the "best" address first, whatever that means
81+ addr:_ <- N.getAddrInfo (Just N.defaultHints) (Just host) (Just port)
82+ sock <- N.socket (N.addrFamily addr) N.Stream N.defaultProtocol
83+ N.setSocketOption sock N.ReuseAddr 1
84+ N.bindSocket sock (N.addrAddress addr)
85+ N.listen sock 5
86+ forkIO $ server sock
87+
88+-- | The transport server handles incoming connections for all endpoints
89+--
90+-- TODO: lots of error checking
91+transportServer :: MVar TransportState -> N.Socket -> IO ()
92+transportServer state sock = forever $ do
93+ (clientSock, _) <- N.accept sock
94+ forkIO $ do
95+ -- Our endpoint the other side is interested in
96+ endPointBS <- NSB.recv clientSock 2
97+ endPointIx <- decodeInt16 endPointBS
98+ -- Full endpoint address of the other side
99+ theirAddress <- BS.concat <$> recvWithLength clientSock
100+ -- Start handling incoming messages
101+ endPointMb <- (^. endPointAt (fromIntegral endPointIx)) <$> readMVar state
102+ case endPointMb of
103+ Just (endPointCh, endPointSt) -> do
104+ endpointServer clientSock (EndPointAddress theirAddress) endPointCh endPointSt
105+ Nothing -> -- TODO: should we do more than just close the socket?
106+ N.sClose clientSock
107+
108+-- | Handle requests for one endpoint
109+endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> IO ()
110+endpointServer sock theirAddress chan state = forever $ do
111+ -- Get connection number
112+ connBS <- NSB.recv sock 2
113+ connIx <- fromIntegral <$> decodeInt16 connBS
114+ if connIx == 0
115+ then createNewConnection
116+ else readMessage connIx
117+ where
118+ -- Create a new connection
119+ createNewConnection :: IO ()
120+ createNewConnection = do
121+ newIx <- modifyMVar state $ \st -> return (nextConnectionIx ^+= 1 $ st, st ^. nextConnectionIx)
122+ newBs <- encodeInt16 (fromIntegral newIx)
123+ NSB.sendMany sock [newBs]
124+ writeChan chan (ConnectionOpened newIx ReliableOrdered theirAddress)
125+
126+ -- Read a message and output it on the endPoint's channel
127+ readMessage :: Int -> IO ()
128+ readMessage connIx = do
129+ payload <- recvWithLength sock
130+ writeChan chan (Received connIx payload)
131+
132+-- | Create a new endpoint
133+tcpNewEndPoint :: MVar TransportState -> N.HostName -> N.ServiceName -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
134+tcpNewEndPoint transportState host port = do
135+ chan <- newChan
136+ endPointState <- newMVar $ EndPointState { _nextConnectionIx = 1 } -- Connection id 0 is reserved for requesting new connections
137+ endPointIx <- modifyMVar transportState $ \st -> do let ix = IntMap.size (st ^. endPoints)
138+ return ((endPointAt ix ^= Just (chan, endPointState)) $ st, ix)
139+ let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show endPointIx
140+ return . Right $ EndPoint { receive = readChan chan
141+ , address = addr
142+ , connect = tcpConnect addr
143+ , newMulticastGroup = undefined
144+ , resolveMulticastGroup = undefined
145+ }
146+
147+-- | Connnect to an endpoint
148+tcpConnect :: EndPointAddress -> EndPointAddress -> Reliability -> IO (Either (FailedWith ConnectErrorCode) Connection)
149+tcpConnect myAddress theirAddress _ = runErrorT $ do
150+ sock <- socketForEndPoint myAddress theirAddress
151+ connIx <- liftIO $ requestNewConnection sock
152+ return $ Connection { send = sendWithLength sock (Just connIx)
153+ , close = N.sClose sock
154+ }
155+
156+-- | Request a new connection
157+requestNewConnection :: N.Socket -> IO ByteString
158+requestNewConnection sock = do
159+ pure <$> encodeInt16 0 >>= NSB.sendMany sock
160+ NSB.recv sock 2
161+
162+-- | Find a socket to the specified endpoint
163+--
164+-- TODOs:
165+-- * Reuse connections
166+-- * Hints?
167+socketForEndPoint :: EndPointAddress -- ^ Our address
168+ -> EndPointAddress -- ^ Their address
169+ -> ErrorT (FailedWith ConnectErrorCode) IO N.Socket
170+socketForEndPoint (EndPointAddress myAddress) (EndPointAddress theirAddress) = do
171+ (host, port, endPointIx) <- maybeToErrorT invalidAddress (decodeEndPointAddress theirAddress)
172+ liftIO $ do
173+ -- Connect to the destination transport
174+ addr:_ <- N.getAddrInfo Nothing (Just host) (Just port)
175+ sock <- N.socket (N.addrFamily addr) N.Stream N.defaultProtocol
176+ N.setSocketOption sock N.ReuseAddr 1
177+ N.connect sock (N.addrAddress addr)
178+ -- Tell it what endpoint we're interested in and our own address
179+ endPointBS <- encodeInt16 (fromIntegral endPointIx)
180+ sendWithLength sock (Just endPointBS) [myAddress]
181+ -- Connection is now established
182+ return sock
183+ where
184+ invalidAddress = FailedWith ConnectInvalidAddress "Invalid address"
185+
186+-- | Send a bunch of bytestrings prepended with their length
187+sendWithLength :: N.Socket -- ^ Socket to send on
188+ -> Maybe ByteString -- ^ Optional header to send before the length
189+ -> [ByteString] -- ^ Payload
190+ -> IO ()
191+sendWithLength sock header payload = do
192+ lengthBs <- encodeInt32 (fromIntegral . sum . map BS.length $ payload)
193+ let msg = maybe id (:) header $ lengthBs : payload
194+ NSB.sendMany sock msg
195+
196+-- | Read a length and then a payload of that length
197+-- TODO: error handling
198+recvWithLength :: N.Socket -> IO [ByteString]
199+recvWithLength sock = do
200+ msgLengthBs <- NSB.recv sock 4
201+ msgLength <- fromIntegral <$> decodeInt32 msgLengthBs
202+ pure <$> NSB.recv sock msgLength
203+
204+-- | Decode end point address
205+-- TODO: This uses regular expression parsing, which is nice, but unnecessary
206+decodeEndPointAddress :: ByteString -> Maybe (N.HostName, N.ServiceName, Int)
207+decodeEndPointAddress bs = match endPointAddressRE (BSC.unpack bs)
208+ where
209+ endPointAddressRE :: RE Char (N.HostName, N.ServiceName, Int)
210+ endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))