@@ -7,7 +7,7 @@ module Network.Wai.Handler.WebSockets
77 , runWebSockets
88 ) where
99
10−import Control.Exception (finally)
10+import Control.Exception (bracket)
1111import Data.ByteString (ByteString)
1212import qualified Data.ByteString.Char8 as BC
1313import qualified Data.ByteString.Lazy as BL
@@ -95,20 +95,20 @@ runWebSockets :: WS.ConnectionOptions
9595 -> IO ByteString
9696 -> (ByteString -> IO ())
9797 -> IO a
98−runWebSockets opts req app src sink = do
99− stream <- WS.makeStream
100− (do
101− bs <- src
102− return $ if BC.null bs then Nothing else Just bs)
103− (\mbBl -> case mbBl of
104− Nothing -> return ()
105− Just bl -> mapM_ sink (BL.toChunks bl))
106−
107− let pc = WS.PendingConnection
108− { WS.pendingOptions = opts
109− , WS.pendingRequest = req
110− , WS.pendingOnAccept = \_ -> return ()
111− , WS.pendingStream = stream
112− }
98+runWebSockets opts req app src sink = bracket mkStream WS.close (app . pc)
99+ where
100+ mkStream =
101+ WS.makeStream
102+ (do
103+ bs <- src
104+ return $ if BC.null bs then Nothing else Just bs)
105+ (\mbBl -> case mbBl of
106+ Nothing -> return ()
107+ Just bl -> mapM_ sink (BL.toChunks bl))
113108
114− app pc `finally` WS.close stream
109+ pc stream = WS.PendingConnection
110+ { WS.pendingOptions = opts
111+ , WS.pendingRequest = req
112+ , WS.pendingOnAccept = \_ -> return ()
113+ , WS.pendingStream = stream
114+ }