@@ -26,20 +26,22 @@ import qualified Network.Socket as N ( HostName
29−import qualified Network.Socket.ByteString as NSB (sendMany, recv)
29+import qualified Network.Socket.ByteString as NBS (sendMany, recv)
30import Control.Concurrent (forkIO, ThreadId)
30import Control.Concurrent (forkIO, ThreadId)
31import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
31import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
32import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, readMVar)
32import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, readMVar)
33import Control.Category ((>>>))
33import Control.Category ((>>>))
34import Control.Applicative ((<*>), (*>), (<$>), pure)
34import Control.Applicative ((<*>), (*>), (<$>), pure)
35−import Control.Monad (forever)
35+import Control.Monad (forever, mzero)
36import Control.Monad.Error (ErrorT, liftIO, runErrorT)
36import Control.Monad.Error (ErrorT, liftIO, runErrorT)
37+import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT)
37import Data.ByteString (ByteString)
38import Data.ByteString (ByteString)
38−import qualified Data.ByteString as BS (length, concat)
39+import qualified Data.ByteString as BS (length, concat, null)
39import qualified Data.ByteString.Char8 as BSC (pack, unpack)
40import qualified Data.ByteString.Char8 as BSC (pack, unpack)
40import Data.IntMap (IntMap)
41import Data.IntMap (IntMap)
41import qualified Data.IntMap as IntMap (empty, size)
42import qualified Data.IntMap as IntMap (empty, size)
42import Data.Lens.Lazy (Lens, lens, intMapLens, (^.), (^=), (^+=))
43import Data.Lens.Lazy (Lens, lens, intMapLens, (^.), (^=), (^+=))
44+import Data.Int (Int32)
43import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
45import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
45data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState) }
47data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState) }
@@ -91,28 +93,31 @@ forkServer host port server = do
91transportServer :: MVar TransportState -> N.Socket -> IO ()
93transportServer :: MVar TransportState -> N.Socket -> IO ()
92transportServer state sock = forever $ do
94transportServer state sock = forever $ do
93 (clientSock, _) <- N.accept sock
95 (clientSock, _) <- N.accept sock
94− forkIO $ do
96+ forkIO $ handleClient state clientSock >> N.sClose clientSock
98+-- | Handle a new incoming connection
100+-- Returns when the client closes the socket and on invalid input.
101+handleClient :: MVar TransportState -> N.Socket -> IO ()
102+handleClient state sock = do
95 -- Our endpoint the other side is interested in
104 -- Our endpoint the other side is interested in
96− endPointBS <- NSB.recv clientSock 2
105+ endPointBS <- recvExact sock 2
97− endPointIx <- decodeInt16 endPointBS
106+ endPointIx <- decodeInt16 (BS.concat $ endPointBS)
98 -- Full endpoint address of the other side
107 -- Full endpoint address of the other side
99− theirAddress <- BS.concat <$> recvWithLength clientSock
108+ theirAddress <- BS.concat <$> recvWithLength sock
100 -- Start handling incoming messages
109 -- Start handling incoming messages
101− endPointMb <- (^. endPointAt (fromIntegral endPointIx)) <$> readMVar state
110+ (endPointCh, endPointSt) <- MaybeT $ (^. endPointAt (fromIntegral endPointIx)) <$> readMVar state
102− case endPointMb of
111+ endpointServer sock (EndPointAddress theirAddress) endPointCh endPointSt
103− Just (endPointCh, endPointSt) -> do
112+ return ()
104− endpointServer clientSock (EndPointAddress theirAddress) endPointCh endPointSt
113+
105− Nothing -> -- TODO: should we do more than just close the socket?
108−-- | Handle requests for one endpoint
114+-- | Handle requests for one endpoint. Returns only if an error occurs
109−endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> IO ()
115+endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> MaybeT IO ()
110−endpointServer sock theirAddress chan state = forever $ do
116+endpointServer sock theirAddress chan state = forever $ do
111− -- Get connection number
112− connBS <- NSB.recv sock 2
117+ connBS <- recvExact sock 2
113− connIx <- fromIntegral <$> decodeInt16 connBS
118+ connIx <- fromIntegral <$> decodeInt16 (BS.concat connBS)
114 if connIx == 0
119 if connIx == 0
115− then createNewConnection
120+ then liftIO $ createNewConnection
116 else readMessage connIx
121 else readMessage connIx
118 -- Create a new connection
123 -- Create a new connection
@@ -120,14 +125,14 @@ endpointServer sock theirAddress chan state = forever $ do
120 createNewConnection = do
125 createNewConnection = do
121 newIx <- modifyMVar state $ \st -> return (nextConnectionIx ^+= 1 $ st, st ^. nextConnectionIx)
126 newIx <- modifyMVar state $ \st -> return (nextConnectionIx ^+= 1 $ st, st ^. nextConnectionIx)
122 newBs <- encodeInt16 (fromIntegral newIx)
127 newBs <- encodeInt16 (fromIntegral newIx)
123− NSB.sendMany sock [newBs]
128+ NBS.sendMany sock [newBs]
124 writeChan chan (ConnectionOpened newIx ReliableOrdered theirAddress)
129 writeChan chan (ConnectionOpened newIx ReliableOrdered theirAddress)
126 -- Read a message and output it on the endPoint's channel
131 -- Read a message and output it on the endPoint's channel
127− readMessage :: Int -> IO ()
132+ readMessage :: Int -> MaybeT IO ()
128 readMessage connIx = do
133 readMessage connIx = do
129 payload <- recvWithLength sock
134 payload <- recvWithLength sock
130− writeChan chan (Received connIx payload)
135+ liftIO $ writeChan chan (Received connIx payload)
132-- | Create a new endpoint
137-- | Create a new endpoint
133tcpNewEndPoint :: MVar TransportState -> N.HostName -> N.ServiceName -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
138tcpNewEndPoint :: MVar TransportState -> N.HostName -> N.ServiceName -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
@@ -156,8 +161,8 @@ tcpConnect myAddress theirAddress _ = runErrorT $ do
156-- | Request a new connection
161-- | Request a new connection
157requestNewConnection :: N.Socket -> IO ByteString
162requestNewConnection :: N.Socket -> IO ByteString
158requestNewConnection sock = do
163requestNewConnection sock = do
159− pure <$> encodeInt16 0 >>= NSB.sendMany sock
164+ pure <$> encodeInt16 0 >>= NBS.sendMany sock
160− NSB.recv sock 2
165+ NBS.recv sock 2
162-- | Find a socket to the specified endpoint
167-- | Find a socket to the specified endpoint
@@ -191,15 +196,14 @@ sendWithLength :: N.Socket -- ^ Socket to send on
191sendWithLength sock header payload = do
196sendWithLength sock header payload = do
192 lengthBs <- encodeInt32 (fromIntegral . sum . map BS.length $ payload)
197 lengthBs <- encodeInt32 (fromIntegral . sum . map BS.length $ payload)
193 let msg = maybe id (:) header $ lengthBs : payload
198 let msg = maybe id (:) header $ lengthBs : payload
194− NSB.sendMany sock msg
199+ NBS.sendMany sock msg
196-- | Read a length and then a payload of that length
201-- | Read a length and then a payload of that length
197-- TODO: error handling
202-- TODO: error handling
198−recvWithLength :: N.Socket -> IO [ByteString]
203+recvWithLength :: N.Socket -> MaybeT IO [ByteString]
199recvWithLength sock = do
204recvWithLength sock = do
200− msgLengthBs <- NSB.recv sock 4
205+ msgLengthBs <- recvExact sock 4
201− msgLength <- fromIntegral <$> decodeInt32 msgLengthBs
206+ decodeInt32 (BS.concat msgLengthBs) >>= recvExact sock
202− pure <$> NSB.recv sock msgLength
204-- | Decode end point address
208-- | Decode end point address
205-- TODO: This uses regular expression parsing, which is nice, but unnecessary
209-- TODO: This uses regular expression parsing, which is nice, but unnecessary
@@ -208,3 +212,22 @@ decodeEndPointAddress bs = match endPointAddressRE (BSC.unpack bs)
209 endPointAddressRE :: RE Char (N.HostName, N.ServiceName, Int)
213 endPointAddressRE :: RE Char (N.HostName, N.ServiceName, Int)
210 endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
214 endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
216+-- | Read an exact number of bytes from a socket
218+-- Returns 'Nothing' if the socket closes prematurely.
219+recvExact :: N.Socket -- ^ Socket to read from
220+ -> Int32 -- ^ Number of bytes to read
221+ -> MaybeT IO [ByteString]
222+recvExact sock len = do
223+ (socketClosed, input) <- liftIO $ go [] len
224+ if socketClosed then mzero else return input
226+ -- Returns input read and whether the socket closed prematurely
227+ go :: [ByteString] -> Int32 -> IO (Bool, [ByteString])
228+ go acc 0 = return (False, reverse acc)
230+ bs <- NBS.recv sock (fromIntegral l `min` 4096)
232+ then return (True, reverse acc)
233+ else go (bs : acc) (l - (fromIntegral $ BS.length bs))