@@ -28,13 +28,13 @@ module Network.Transport.TCP.Mock.Socket
2828 -- * Internal API
2929 , writeSocket
3030 , readSocket
31+ , Message(..)
3132 ) where
3233
3334import Data.Word (Word8)
3435import Data.Map (Map)
3536import qualified Data.Map as Map
3637import Control.Category ((>>>))
37−import Control.Applicative ((<$>))
3838import Control.Concurrent.MVar
3939import Control.Concurrent.Chan
4040import System.IO.Unsafe (unsafePerformIO)
@@ -99,12 +99,16 @@ data Socket = Socket {
9999data SocketState =
100100 Uninit
101101 | BoundSocket { socketBacklog :: Chan (Socket, SockAddr, MVar Socket) }
102− | Connected { socketPeer :: Socket,socketBuff :: Chan Word8 }
102+ | Connected { socketPeer :: Socket, socketBuff :: Chan Message }
103103 | Closed
104104
105+data Message =
106+ Payload Word8
107+ | CloseSocket
108+
105109data AddrInfo = AddrInfo {
106− addrFamily :: Family
107− , addrAddress :: !SockAddr
110+ addrFamily :: Family
111+ , addrAddress :: SockAddr
108112 }
109113
110114data SockAddr = SockAddrInet PortNumber HostAddress
@@ -165,29 +169,23 @@ accept serverSock = do
165169 return (socketBacklog st)
166170 _ ->
167171 error "accept: socket not bound"
168− (them, theirAddress, reply) <- readChan backlog
169− buff <- newChan
172+ (theirSocket, theirAddress, reply) <- readChan backlog
173+ ourBuff <- newChan
170174 ourState <- newMVar Connected {
171− socketPeer = them
172− , socketBuff = buff
175+ socketPeer = theirSocket
176+ , socketBuff = ourBuff
173177 }
174− let us = Socket {
178+ let ourSocket = Socket {
175179 socketState = ourState
176180 , socketDescription = ""
177181 }
178− putMVar reply us
179− return (us, theirAddress)
182+ putMVar reply ourSocket
183+ return (ourSocket, theirAddress)
180184
181185sClose :: Socket -> IO ()
182186sClose sock = do
183− mPeer <- modifyMVar (socketState sock) $ \st -> case st of
184− Connected {} ->
185− return (Closed, Just $ socketPeer st)
186− _ ->
187− return (Closed, Nothing)
188− case mPeer of
189− Just peer -> modifyMVar_ (socketState peer) $ const (return Closed)
190− Nothing -> return ()
187+ writeSocket sock CloseSocket
188+ modifyMVar_ (socketState sock) $ const (return Closed)
191189
192190connect :: Socket -> SockAddr -> IO ()
193191connect us serverAddr = do
@@ -223,19 +221,29 @@ shutdown = error "shutdown not implemented"
223221-- Functions with no direct public counterpart --
224222--------------------------------------------------------------------------------
225223
226−writeSocket :: Socket -> Word8 -> IO ()
227−writeSocket sock w = do
228− peer <- withMVar (socketState sock) $ \st -> case st of
224+peerBuffer :: Socket -> IO (Either String (Chan Message))
225+peerBuffer sock = do
226+ mPeer <- withMVar (socketState sock) $ \st -> case st of
229227 Connected {} ->
230− return (socketPeer st)
231− _ ->
232− error "writeSocket: not connected"
233− theirBuff <- withMVar (socketState peer) $ \st -> case st of
234− Connected {} ->
235− return (socketBuff st)
228+ return (Just (socketPeer st))
236229 _ ->
237− error "writeSocket: peer socket closed"
238− writeChan theirBuff w
230+ return Nothing
231+ case mPeer of
232+ Just peer -> withMVar (socketState peer) $ \st -> case st of
233+ Connected {} ->
234+ return (Right (socketBuff st))
235+ _ ->
236+ return (Left "Peer socket closed")
237+ Nothing ->
238+ return (Left "Socket closed")
239+
240+writeSocket :: Socket -> Message -> IO ()
241+writeSocket sock msg = do
242+ theirBuff <- peerBuffer sock
243+ case theirBuff of
244+ Right buff -> writeChan buff msg
245+ Left err -> case msg of Payload _ -> error $ "writeSocket: " ++ err
246+ CloseSocket -> return ()
239247
240248readSocket :: Socket -> IO (Maybe Word8)
241249readSocket sock = do
@@ -245,5 +253,14 @@ readSocket sock = do
245253 _ ->
246254 return Nothing
247255 case mBuff of
248− Just buff -> Just <$> readChan buff
249− Nothing -> return Nothing
256+ Just buff -> do
257+ msg <- readChan buff
258+ case msg of
259+ Payload w -> return (Just w)
260+ CloseSocket -> modifyMVar (socketState sock) $ \st -> case st of
261+ Connected {} ->
262+ return (Closed, Nothing)
263+ _ ->
264+ error "readSocket: socket in unexpected state"
265+ Nothing ->
266+ return Nothing