diff --git a/src/Network/WebSockets/Snap.hs b/src/Network/WebSockets/Snap.hs --- a/src/Network/WebSockets/Snap.hs +++ b/src/Network/WebSockets/Snap.hs @@ -1,41 +1,168 @@ +-------------------------------------------------------------------------------- -- | Snap integration for the WebSockets library +{-# LANGUAGE DeriveDataTypeable #-} module Network.WebSockets.Snap ( runWebSocketsSnap , runWebSocketsSnapWith ) where -import qualified Network.WebSockets as WS -import qualified Snap.Core as Snap -import qualified Snap.Internal.Http.Types as Snap -import qualified Snap.Types.Headers as Headers +-------------------------------------------------------------------------------- +import Blaze.ByteString.Builder (Builder) +import qualified Blaze.ByteString.Builder as Builder +import Control.Concurrent (forkIO, myThreadId, threadDelay) +import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, + takeMVar) +import Control.Exception (Exception (..), + SomeException (..), handle, + throw, throwTo) +import Control.Monad (forever) +import Control.Monad.Trans (lift) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Char8 as BC +import qualified Data.ByteString.Lazy as BL +import qualified Data.Enumerator as E +import qualified Data.Enumerator.List as EL +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Typeable (Typeable, cast) +import qualified Network.WebSockets as WS +import qualified Network.WebSockets.Connection as WS +import qualified Snap.Core as Snap +import qualified Snap.Internal.Http.Types as Snap +import qualified Snap.Types.Headers as Headers +import System.IO.Streams (InputStream, OutputStream) +import qualified System.IO.Streams as Streams + + +-------------------------------------------------------------------------------- +data Chunk + = Chunk ByteString + | Eof + | Error SomeException + deriving (Show) + + +-------------------------------------------------------------------------------- +data ServerAppDone = ServerAppDone + deriving (Eq, Ord, Show, Typeable) + + +-------------------------------------------------------------------------------- +instance Exception ServerAppDone where + toException ServerAppDone = SomeException ServerAppDone + fromException (SomeException e) = cast e + + +-------------------------------------------------------------------------------- +copyIterateeToMVar :: MVar Chunk -> E.Iteratee ByteString IO () +copyIterateeToMVar mvar = E.catchError go handler + where + go = do + mbs <- EL.head + case mbs of + Just x -> lift (putMVar mvar (Chunk x)) >> go + Nothing -> lift (putMVar mvar Eof) + + handler se@(SomeException e) = case cast e of + -- Clean exit + Just ServerAppDone -> return () + -- Actual error + Nothing -> lift $ putMVar mvar $ Error se + + +-------------------------------------------------------------------------------- +copyMVarToInputStream :: MVar Chunk -> IO (InputStream ByteString) +copyMVarToInputStream mvar = Streams.makeInputStream go + where + go = do + chunk <- takeMVar mvar + case chunk of + Chunk x -> return (Just x) + Eof -> return Nothing + Error (SomeException e) -> throw e + + +-------------------------------------------------------------------------------- +copyOutputStreamToIteratee :: E.Iteratee ByteString IO () + -> IO (OutputStream Builder) +copyOutputStreamToIteratee iteratee0 = do + ref <- newIORef =<< E.runIteratee iteratee0 + Streams.makeOutputStream (go ref) + where + go _ Nothing = return () + go ref (Just bld) = do + step <- readIORef ref + case step of + E.Continue f -> do + let chunks = BL.toChunks $ Builder.toLazyByteString bld + step' <- E.runIteratee $ f $ E.Chunks chunks + writeIORef ref step' + E.Yield () _ -> throw WS.ConnectionClosed + E.Error (SomeException e) -> throw e + + +-------------------------------------------------------------------------------- -- | The following function escapes from the current 'Snap.Snap' handler, and -- continues processing the 'WS.WebSockets' action. The action to be executed -- takes the 'WS.Request' as a parameter, because snap has already read this -- from the socket. -runWebSocketsSnap :: WS.Protocol p - => (WS.Request -> WS.WebSockets p ()) +runWebSocketsSnap :: WS.ServerApp -> Snap.Snap () -runWebSocketsSnap = runWebSocketsSnapWith WS.defaultWebSocketsOptions +runWebSocketsSnap = runWebSocketsSnapWith WS.defaultConnectionOptions + +-------------------------------------------------------------------------------- -- | Variant of 'runWebSocketsSnap' which allows custom options -runWebSocketsSnapWith :: WS.Protocol p - => WS.WebSocketsOptions - -> (WS.Request -> WS.WebSockets p ()) +runWebSocketsSnapWith :: WS.ConnectionOptions + -> WS.ServerApp -> Snap.Snap () -runWebSocketsSnapWith options ws = do +runWebSocketsSnapWith options app = do rq <- Snap.getRequest - Snap.escapeHttp $ \tickle writeEnd -> + Snap.escapeHttp $ \tickle writeEnd -> do + + thisThread <- lift myThreadId + mvar <- lift newEmptyMVar + is <- lift $ copyMVarToInputStream mvar + os <- lift $ copyOutputStreamToIteratee writeEnd + let options' = options - { WS.onPong = tickle (max 30) >> WS.onPong options - } + { WS.connectionOnPong = do + tickle (max 60) + WS.connectionOnPong options + } - in WS.runWebSocketsWith options' (fromSnapRequest rq) ws writeEnd + pc = WS.PendingConnection + { WS.pendingOptions = options' + , WS.pendingRequest = fromSnapRequest rq + , WS.pendingOnAccept = forkPingThread + , WS.pendingIn = is + , WS.pendingOut = os + } + _ <- lift $ forkIO $ app pc >> throwTo thisThread ServerAppDone + copyIterateeToMVar mvar + + +-------------------------------------------------------------------------------- +-- | Start a ping thread in the background +forkPingThread :: WS.Connection -> IO () +forkPingThread conn = do + _ <- forkIO pingThread + return () + where + pingThread = handle ignore $ forever $ do + WS.sendPing conn (BC.pack "ping") + threadDelay $ 30 * 1000 * 1000 + + ignore :: SomeException -> IO () + ignore _ = return () + + +-------------------------------------------------------------------------------- -- | Convert a snap request to a websockets request -fromSnapRequest :: Snap.Request -> WS.RequestHttpPart -fromSnapRequest rq = WS.RequestHttpPart - { WS.requestHttpPath = Snap.rqURI rq - , WS.requestHttpHeaders = Headers.toList (Snap.rqHeaders rq) - , WS.requestHttpSecure = Snap.rqIsSecure rq +fromSnapRequest :: Snap.Request -> WS.RequestHead +fromSnapRequest rq = WS.RequestHead + { WS.requestPath = Snap.rqURI rq + , WS.requestHeaders = Headers.toList (Snap.rqHeaders rq) + , WS.requestSecure = Snap.rqIsSecure rq } diff --git a/websockets-snap.cabal b/websockets-snap.cabal --- a/websockets-snap.cabal +++ b/websockets-snap.cabal @@ -1,5 +1,5 @@ Name: websockets-snap -Version: 0.7.0.0 +Version: 0.8.0.0 Synopsis: Snap integration for the websockets library Description: Snap integration for the websockets library License: BSD3 @@ -17,10 +17,15 @@ Network.WebSockets.Snap Build-depends: - base >= 4 && < 5, - snap-core >= 0.8 && < 0.10, - snap-server >= 0.8 && < 0.10, - websockets >= 0.7 && < 0.8 + base >= 4 && < 5, + blaze-builder >= 0.3 && < 0.4, + bytestring >= 0.10 && < 0.11, + enumerator >= 0.4 && < 0.5, + io-streams >= 1.1 && < 1.2, + mtl >= 2.1 && < 2.2, + snap-core >= 0.8 && < 0.10, + snap-server >= 0.8 && < 0.10, + websockets >= 0.8 && < 0.9 Source-repository head Type: git