Permalink
Browse files

fix the websockets interceptor, stop processing when due

  • Loading branch information...
1 parent 8149617 commit bd63158ca53a6dec6b04316800236aa6c521b907 @bitonic committed May 11, 2012
Showing with 8 additions and 1 deletion.
  1. +8 −1 wai-websockets/Network/Wai/Handler/WebSockets.hs
@@ -52,7 +52,10 @@ runWebSockets opts req app source conn = do
push iter bs = do
step <- liftIO $ E.runIteratee $ E.enumList 1 [bs] E.$$ iter
- return $ C.StateProcessing $ E.returnI step
+ case step of
+ E.Continue _ -> return $ C.StateProcessing $ E.returnI step
+ E.Yield out inp -> return $ C.StateDone (streamToMaybe inp) out
+ E.Error e -> C.monadThrow e
close iter = do
_ <- liftIO $ E.runIteratee $ E.enumEOF E.$$ iter
return ()
@@ -63,3 +66,7 @@ iterConnection c = E.continue go
go (E.Chunks []) = E.continue go
go (E.Chunks cs) = E.tryIO (Warp.connSendMany c cs) >> E.continue go
go E.EOF = E.continue go
+
+streamToMaybe :: E.Stream S.ByteString -> Maybe S.ByteString
+streamToMaybe E.EOF = Nothing
+streamToMaybe (E.Chunks bs) = Just $ S.concat bs

0 comments on commit bd63158

Please sign in to comment.