diff --git a/CHANGELOG b/CHANGELOG --- a/CHANGELOG +++ b/CHANGELOG @@ -1,3 +1,6 @@ +- 0.9.5.0 + * Bugfixes wrt closing sockets and streams + - 0.9.4.0 * Add `makePendingConnectionFromStream` function * Bump `attoparsec` dependency diff --git a/src/Network/WebSockets/Client.hs b/src/Network/WebSockets/Client.hs --- a/src/Network/WebSockets/Client.hs +++ b/src/Network/WebSockets/Client.hs @@ -12,8 +12,7 @@ -------------------------------------------------------------------------------- import qualified Blaze.ByteString.Builder as Builder -import Control.Concurrent.MVar (newMVar) -import Control.Exception (finally, throwIO) +import Control.Exception (bracket, finally, throwIO) import Data.IORef (newIORef) import qualified Data.Text as T import qualified Data.Text.Encoding as T @@ -100,18 +99,14 @@ Response _ _ <- return $ finishResponse protocol request response parse <- decodeMessages protocol stream write <- encodeMessages protocol ClientConnection stream + sentRef <- newIORef False - parseLock <- newMVar () - writeLock <- newMVar () - stream' <- newIORef (DecoderEncoder parse write) - sentRef <- newIORef False app Connection { connectionOptions = opts , connectionType = ClientConnection , connectionProtocol = protocol - , connectionParseLock = parseLock - , connectionWriteLock = writeLock - , connectionStream = stream' + , connectionParse = parse + , connectionWrite = write , connectionSentClose = sentRef } where @@ -128,6 +123,8 @@ -> Headers -- ^ Custom headers to send -> ClientApp a -- ^ Client application -> IO a -runClientWithSocket sock host path opts customHeaders app = do - stream <- Stream.makeSocketStream sock - runClientWithStream stream host path opts customHeaders app +runClientWithSocket sock host path opts customHeaders app = bracket + (Stream.makeSocketStream sock) + Stream.close + (\stream -> + runClientWithStream stream host path opts customHeaders app) diff --git a/src/Network/WebSockets/Connection.hs b/src/Network/WebSockets/Connection.hs --- a/src/Network/WebSockets/Connection.hs +++ b/src/Network/WebSockets/Connection.hs @@ -9,7 +9,6 @@ , acceptRequestWith , rejectRequest - , DecoderEncoder (..) , Connection (..) , ConnectionOptions (..) @@ -33,9 +32,8 @@ -------------------------------------------------------------------------------- import qualified Blaze.ByteString.Builder as Builder import Control.Concurrent (forkIO, threadDelay) -import Control.Concurrent.MVar (MVar, newMVar, withMVar) import Control.Exception (AsyncException, fromException, - handle, onException, throwIO) + handle, throwIO) import Control.Monad (unless) import qualified Data.ByteString as B import Data.IORef (IORef, newIORef, readIORef, @@ -103,17 +101,13 @@ parse <- decodeMessages protocol (pendingStream pc) write <- encodeMessages protocol ServerConnection (pendingStream pc) - parseLock <- newMVar () - writeLock <- newMVar () - stream <- newIORef (DecoderEncoder parse write) sentRef <- newIORef False let connection = Connection { connectionOptions = pendingOptions pc , connectionType = ServerConnection , connectionProtocol = protocol - , connectionParseLock = parseLock - , connectionWriteLock = writeLock - , connectionStream = stream + , connectionParse = parse + , connectionWrite = write , connectionSentClose = sentRef } @@ -131,20 +125,12 @@ -------------------------------------------------------------------------------- --- | Type representing an available or unavailable stream. -data DecoderEncoder a - = DecoderEncoder !(IO (Maybe a)) !(a -> IO ()) - | Closed - - --------------------------------------------------------------------------------- data Connection = Connection { connectionOptions :: !ConnectionOptions , connectionType :: !ConnectionType , connectionProtocol :: !Protocol - , connectionParseLock :: !(MVar ()) - , connectionWriteLock :: !(MVar ()) - , connectionStream :: !(IORef (DecoderEncoder Message)) + , connectionParse :: !(IO (Maybe Message)) + , connectionWrite :: !(Message -> IO ()) , connectionSentClose :: !(IORef Bool) -- ^ According to the RFC, both the client and the server MUST send -- a close control message to each other. Either party can initiate @@ -172,18 +158,11 @@ -------------------------------------------------------------------------------- receive :: Connection -> IO Message -receive conn = withMVar (connectionParseLock conn) $ \() -> - tryParse `onException` (writeIORef (connectionStream conn) Closed) - where - tryParse = do - stream <- readIORef (connectionStream conn) - case stream of - Closed -> throwIO ConnectionClosed - DecoderEncoder parse _ -> do - mbMsg <- parse - case mbMsg of - Nothing -> throwIO ConnectionClosed - Just msg -> return msg +receive conn = do + mbMsg <- connectionParse conn + case mbMsg of + Nothing -> throwIO ConnectionClosed + Just msg -> return msg -------------------------------------------------------------------------------- @@ -227,18 +206,12 @@ -------------------------------------------------------------------------------- send :: Connection -> Message -> IO () -send conn msg = withMVar (connectionWriteLock conn) $ \() -> - trySend `onException` writeIORef (connectionStream conn) Closed - where - trySend = do - stream <- readIORef (connectionStream conn) - case msg of - (ControlMessage (Close _ _)) -> - writeIORef (connectionSentClose conn) True - _ -> return () - case stream of - Closed -> throwIO ConnectionClosed - DecoderEncoder _ write -> write msg +send conn msg = do + case msg of + (ControlMessage (Close _ _)) -> + writeIORef (connectionSentClose conn) True + _ -> return () + connectionWrite conn msg -------------------------------------------------------------------------------- diff --git a/src/Network/WebSockets/Server.hs b/src/Network/WebSockets/Server.hs --- a/src/Network/WebSockets/Server.hs +++ b/src/Network/WebSockets/Server.hs @@ -84,18 +84,22 @@ -> ConnectionOptions -> ServerApp -> IO () -runApp socket opts app = do - pending <- makePendingConnection socket opts - app pending +runApp socket opts app = + bracket + (makePendingConnection socket opts) + (Stream.close . pendingStream) + app -------------------------------------------------------------------------------- --- | Turns a socket, connected to some client, into a 'PendingConnection'. +-- | Turns a socket, connected to some client, into a 'PendingConnection'. The +-- 'PendingConnection' should be closed using 'Stream.close' later. makePendingConnection :: Socket -> ConnectionOptions -> IO PendingConnection makePendingConnection socket opts = do - stream <- Stream.makeSocketStream socket + stream <- Stream.makeSocketStream socket makePendingConnectionFromStream stream opts + -- | More general version of 'makePendingConnection' for 'Stream.Stream' -- instead of a 'Socket'. diff --git a/src/Network/WebSockets/Stream.hs b/src/Network/WebSockets/Stream.hs --- a/src/Network/WebSockets/Stream.hs +++ b/src/Network/WebSockets/Stream.hs @@ -11,14 +11,15 @@ , close ) where -import Control.Applicative ((<$>)) -import qualified Control.Concurrent.Chan as Chan -import Control.Exception (throwIO) -import Control.Monad (forM_) +import Control.Concurrent.MVar (MVar, newEmptyMVar, newMVar, + putMVar, takeMVar, withMVar) +import Control.Exception (onException, throwIO) +import Control.Monad (forM_, when) import qualified Data.Attoparsec.ByteString as Atto import qualified Data.ByteString as B import qualified Data.ByteString.Lazy as BL -import Data.IORef (IORef, newIORef, readIORef, +import Data.IORef (IORef, atomicModifyIORef, + newIORef, readIORef, writeIORef) import qualified Network.Socket as S import qualified Network.Socket.ByteString as SB (recv) @@ -49,13 +50,53 @@ -------------------------------------------------------------------------------- +-- | Create a stream from a "receive" and "send" action. The following +-- properties apply: +-- +-- - Regardless of the provided "receive" and "send" functions, reading and +-- writing from the stream will be thread-safe, i.e. this function will create +-- a receive and write lock to be used internally. +-- +-- - Reading from or writing or to a closed 'Stream' will always throw an +-- exception, even if the underlying "receive" and "send" functions do not +-- (we do the bookkeeping). +-- +-- - Streams should always be closed. makeStream :: IO (Maybe B.ByteString) -- ^ Reading -> (Maybe BL.ByteString -> IO ()) -- ^ Writing -> IO Stream -- ^ Resulting stream -makeStream i o = Stream i o <$> newIORef (Open B.empty) +makeStream receive send = do + ref <- newIORef (Open B.empty) + receiveLock <- newMVar () + sendLock <- newMVar () + return $ Stream (receive' ref receiveLock) (send' ref sendLock) ref + where + closeRef :: IORef StreamState -> IO () + closeRef ref = atomicModifyIORef ref $ \state -> case state of + Open buf -> (Closed buf, ()) + Closed buf -> (Closed buf, ()) + assertNotClosed :: IORef StreamState -> IO a -> IO a + assertNotClosed ref io = do + state <- readIORef ref + case state of + Closed _ -> throwIO ConnectionClosed + Open _ -> io + receive' :: IORef StreamState -> MVar () -> IO (Maybe B.ByteString) + receive' ref lock = withMVar lock $ \() -> assertNotClosed ref $ do + mbBs <- onException receive (closeRef ref) + case mbBs of + Nothing -> closeRef ref >> return Nothing + Just bs -> return (Just bs) + + send' :: IORef StreamState -> MVar () -> (Maybe BL.ByteString -> IO ()) + send' ref lock mbBs = withMVar lock $ \() -> assertNotClosed ref $ do + when (mbBs == Nothing) (closeRef ref) + onException (send mbBs) (closeRef ref) + + -------------------------------------------------------------------------------- makeSocketStream :: S.Socket -> IO Stream makeSocketStream socket = makeStream receive send @@ -65,7 +106,7 @@ return $ if B.null bs then Nothing else Just bs send Nothing = return () - send (Just bs) = + send (Just bs) = do #if !defined(mingw32_HOST_OS) SBL.sendAll socket bs #else @@ -76,10 +117,10 @@ -------------------------------------------------------------------------------- makeEchoStream :: IO Stream makeEchoStream = do - chan <- Chan.newChan - makeStream (Chan.readChan chan) $ \mbBs -> case mbBs of - Nothing -> Chan.writeChan chan Nothing - Just bs -> forM_ (BL.toChunks bs) $ \c -> Chan.writeChan chan (Just c) + mvar <- newEmptyMVar + makeStream (takeMVar mvar) $ \mbBs -> case mbBs of + Nothing -> putMVar mvar Nothing + Just bs -> forM_ (BL.toChunks bs) $ \c -> putMVar mvar (Just c) -------------------------------------------------------------------------------- diff --git a/tests/haskell/Network/WebSockets/Handshake/Tests.hs b/tests/haskell/Network/WebSockets/Handshake/Tests.hs --- a/tests/haskell/Network/WebSockets/Handshake/Tests.hs +++ b/tests/haskell/Network/WebSockets/Handshake/Tests.hs @@ -22,7 +22,6 @@ import Network.WebSockets.Connection import Network.WebSockets.Http import qualified Network.WebSockets.Stream as Stream -import Network.WebSockets.Tests.Util -------------------------------------------------------------------------------- @@ -43,6 +42,7 @@ _ <- app (PendingConnection defaultConnectionOptions rq nullify echo) return () mbRh <- Stream.parse echo decodeResponseHead + Stream.close echo case mbRh of Nothing -> fail "testHandshake: No response" Just rh -> return rh diff --git a/tests/haskell/Network/WebSockets/Tests.hs b/tests/haskell/Network/WebSockets/Tests.hs --- a/tests/haskell/Network/WebSockets/Tests.hs +++ b/tests/haskell/Network/WebSockets/Tests.hs @@ -9,7 +9,8 @@ -------------------------------------------------------------------------------- import qualified Blaze.ByteString.Builder as Builder import Control.Applicative ((<$>)) -import Control.Monad (replicateM, forM_) +import Control.Concurrent (forkIO) +import Control.Monad (forM_, replicateM) import qualified Data.ByteString.Lazy as BL import Data.List (intersperse) import Data.Maybe (catMaybes) @@ -47,8 +48,9 @@ echo <- Stream.makeEchoStream parse <- decodeMessages protocol echo write <- encodeMessages protocol ClientConnection echo - forM_ msgs write + _ <- forkIO $ forM_ msgs write msgs' <- catMaybes <$> replicateM (length msgs) parse + Stream.close echo msgs @=? msgs' @@ -61,12 +63,13 @@ -- is' <- Streams.filter isDataMessage =<< Hybi13.decodeMessages is -- Simple hacky encoding of all frames - mapM_ (Stream.write echo) - [ Builder.toLazyByteString (Hybi13.encodeFrame Nothing f) - | FragmentedMessage _ frames <- fragmented - , f <- frames - ] - Stream.close echo + _ <- forkIO $ do + mapM_ (Stream.write echo) + [ Builder.toLazyByteString (Hybi13.encodeFrame Nothing f) + | FragmentedMessage _ frames <- fragmented + , f <- frames + ] + Stream.close echo -- Check if we got all data msgs <- filter isDataMessage <$> parseAll parse diff --git a/websockets.cabal b/websockets.cabal --- a/websockets.cabal +++ b/websockets.cabal @@ -1,5 +1,5 @@ Name: websockets -Version: 0.9.4.0 +Version: 0.9.5.0 Synopsis: A sensible and clean way to write WebSocket-capable servers in Haskell.