@@ -7,6 +7,7 @@ module Network.Wai.Handler.WebSockets
77 , runWebSockets
88 ) where
99
10+import Control.Exception (bracket)
1011import Data.ByteString (ByteString)
1112import qualified Data.ByteString.Char8 as BC
1213import qualified Data.ByteString.Lazy as BL
@@ -94,20 +95,20 @@ runWebSockets :: WS.ConnectionOptions
9495 -> IO ByteString
9596 -> (ByteString -> IO ())
9697 -> IO a
97−runWebSockets opts req app src sink = do
98− stream <- WS.makeStream
99− (do
100− bs <- src
101− return $ if BC.null bs then Nothing else Just bs)
102− (\mbBl -> case mbBl of
103− Nothing -> return ()
104− Just bl -> mapM_ sink (BL.toChunks bl))
105−
106− let pc = WS.PendingConnection
107− { WS.pendingOptions = opts
108− , WS.pendingRequest = req
109− , WS.pendingOnAccept = \_ -> return ()
110− , WS.pendingStream = stream
111− }
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))
112108
113− app pc
109+ pc stream = WS.PendingConnection
110+ { WS.pendingOptions = opts
111+ , WS.pendingRequest = req
112+ , WS.pendingOnAccept = \_ -> return ()
113+ , WS.pendingStream = stream
114+ }