@@ -26,20 +26,22 @@ import qualified Network.Socket as N ( HostName
2626 , sClose
2727 , accept
2828 )
29−import qualified Network.Socket.ByteString as NSB (sendMany, recv)
29+import qualified Network.Socket.ByteString as NBS (sendMany, recv)
3030import Control.Concurrent (forkIO, ThreadId)
3131import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
3232import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, readMVar)
3333import Control.Category ((>>>))
3434import Control.Applicative ((<*>), (*>), (<$>), pure)
35−import Control.Monad (forever)
35+import Control.Monad (forever, mzero)
3636import Control.Monad.Error (ErrorT, liftIO, runErrorT)
37+import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT)
3738import Data.ByteString (ByteString)
38−import qualified Data.ByteString as BS (length, concat)
39+import qualified Data.ByteString as BS (length, concat, null)
3940import qualified Data.ByteString.Char8 as BSC (pack, unpack)
4041import Data.IntMap (IntMap)
4142import qualified Data.IntMap as IntMap (empty, size)
4243import Data.Lens.Lazy (Lens, lens, intMapLens, (^.), (^=), (^+=))
44+import Data.Int (Int32)
4345import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
4446
4547data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState) }
@@ -91,28 +93,31 @@ forkServer host port server = do
9193transportServer :: MVar TransportState -> N.Socket -> IO ()
9294transportServer state sock = forever $ do
9395 (clientSock, _) <- N.accept sock
94− forkIO $ do
96+ forkIO $ handleClient state clientSock >> N.sClose clientSock
97+
98+-- | Handle a new incoming connection
99+--
100+-- Returns when the client closes the socket and on invalid input.
101+handleClient :: MVar TransportState -> N.Socket -> IO ()
102+handleClient state sock = do
103+ runMaybeT $ do
95104 -- Our endpoint the other side is interested in
96− endPointBS <- NSB.recv clientSock 2
97− endPointIx <- decodeInt16 endPointBS
105+ endPointBS <- recvExact sock 2
106+ endPointIx <- decodeInt16 (BS.concat $ endPointBS)
98107 -- Full endpoint address of the other side
99− theirAddress <- BS.concat <$> recvWithLength clientSock
108+ theirAddress <- BS.concat <$> recvWithLength sock
100109 -- 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
110+ (endPointCh, endPointSt) <- MaybeT $ (^. endPointAt (fromIntegral endPointIx)) <$> readMVar state
111+ endpointServer sock (EndPointAddress theirAddress) endPointCh endPointSt
112+ return ()
113+
114+-- | Handle requests for one endpoint. Returns only if an error occurs
115+endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> MaybeT IO ()
116+endpointServer sock theirAddress chan state = forever $ do
117+ connBS <- recvExact sock 2
118+ connIx <- fromIntegral <$> decodeInt16 (BS.concat connBS)
114119 if connIx == 0
115− then createNewConnection
120+ then liftIO $ createNewConnection
116121 else readMessage connIx
117122 where
118123 -- Create a new connection
@@ -120,14 +125,14 @@ endpointServer sock theirAddress chan state = forever $ do
120125 createNewConnection = do
121126 newIx <- modifyMVar state $ \st -> return (nextConnectionIx ^+= 1 $ st, st ^. nextConnectionIx)
122127 newBs <- encodeInt16 (fromIntegral newIx)
123− NSB.sendMany sock [newBs]
128+ NBS.sendMany sock [newBs]
124129 writeChan chan (ConnectionOpened newIx ReliableOrdered theirAddress)
125130
126131 -- Read a message and output it on the endPoint's channel
127− readMessage :: Int -> IO ()
132+ readMessage :: Int -> MaybeT IO ()
128133 readMessage connIx = do
129134 payload <- recvWithLength sock
130− writeChan chan (Received connIx payload)
135+ liftIO $ writeChan chan (Received connIx payload)
131136
132137-- | Create a new endpoint
133138tcpNewEndPoint :: MVar TransportState -> N.HostName -> N.ServiceName -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
@@ -156,8 +161,8 @@ tcpConnect myAddress theirAddress _ = runErrorT $ do
156161-- | Request a new connection
157162requestNewConnection :: N.Socket -> IO ByteString
158163requestNewConnection sock = do
159− pure <$> encodeInt16 0 >>= NSB.sendMany sock
160− NSB.recv sock 2
164+ pure <$> encodeInt16 0 >>= NBS.sendMany sock
165+ NBS.recv sock 2
161166
162167-- | Find a socket to the specified endpoint
163168--
@@ -191,15 +196,14 @@ sendWithLength :: N.Socket -- ^ Socket to send on
191196sendWithLength sock header payload = do
192197 lengthBs <- encodeInt32 (fromIntegral . sum . map BS.length $ payload)
193198 let msg = maybe id (:) header $ lengthBs : payload
194− NSB.sendMany sock msg
199+ NBS.sendMany sock msg
195200
196201-- | Read a length and then a payload of that length
197202-- TODO: error handling
198−recvWithLength :: N.Socket -> IO [ByteString]
203+recvWithLength :: N.Socket -> MaybeT IO [ByteString]
199204recvWithLength sock = do
200− msgLengthBs <- NSB.recv sock 4
201− msgLength <- fromIntegral <$> decodeInt32 msgLengthBs
202− pure <$> NSB.recv sock msgLength
205+ msgLengthBs <- recvExact sock 4
206+ decodeInt32 (BS.concat msgLengthBs) >>= recvExact sock
203207
204208-- | Decode end point address
205209-- TODO: This uses regular expression parsing, which is nice, but unnecessary
@@ -208,3 +212,22 @@ decodeEndPointAddress bs = match endPointAddressRE (BSC.unpack bs)
208212 where
209213 endPointAddressRE :: RE Char (N.HostName, N.ServiceName, Int)
210214 endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
215+
216+-- | Read an exact number of bytes from a socket
217+--
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
225+ where
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)
229+ go acc l = do
230+ bs <- NBS.recv sock (fromIntegral l `min` 4096)
231+ if BS.null bs
232+ then return (True, reverse acc)
233+ else go (bs : acc) (l - (fromIntegral $ BS.length bs))