@@ -12,7 +12,7 @@ import Foreign
1212import Foreign.C.String
1313import Control.Monad.IO.Class (liftIO)
1414import qualified Data.ByteString as S
15−import Data.Enumerator (($$), enumList, joinI)
15+import Data.Enumerator (($$), enumList, joinI, Enumeratee, Stream (..), Iteratee (..), Step (..))
1616import Blaze.ByteString.Builder (fromByteString)
1717#if !WINDOWS
1818import System.Cmd (rawSystem)
@@ -20,6 +20,7 @@ import System.Cmd (rawSystem)
2020import Codec.Zlib.Enum (ungzip)
2121import Blaze.ByteString.Builder.Enumerator (builderToByteString)
2222import qualified Data.Enumerator.List as EL
23+import Control.Monad.Trans.Class (lift)
2324
2425ping :: IORef Bool -> Middleware
2526ping var app req
@@ -37,11 +38,68 @@ ping var app req
3738 let (isEnc, headers') = fixHeaders id headers
3839 let fixEnc x =
3940 if isEnc
40− then joinI $ builderToByteString $$ joinI $ ungzip $$ joinI $ EL.map fromByteString $$ x
41+ then joinI $ ungzip $$ x
4142 else x
42− fixEnc $ enumList 1 [fromByteString "<script>setInterval(function(){var x;if(window.XMLHttpRequest){x=new XMLHttpRequest();}else{x=new ActiveXObject(\"Microsoft.XMLHTTP\");}x.open(\"GET\",\"/_ping\",false);x.send();},60000)</script>"] $$ f status headers'
43+ joinI $ builderToByteString $$ fixEnc $ joinI $ insideHead "<script>setInterval(function(){var x;if(window.XMLHttpRequest){x=new XMLHttpRequest();}else{x=new ActiveXObject(\"Microsoft.XMLHTTP\");}x.open(\"GET\",\"/_ping\",false);x.send();},60000)</script>" $$ joinI $ EL.map fromByteString $$ f status headers'
4344 _ -> f status headers
4445
46+insideHead :: S.ByteString -> Enumeratee S.ByteString S.ByteString IO a
47+insideHead toInsert =
48+ go "" whole
49+ where
50+ whole = "<head>"
51+ go :: S.ByteString -> S.ByteString -> Step S.ByteString IO a -> Iteratee S.ByteString IO (Step S.ByteString IO a)
52+ go held atFront step = do
53+ mx <- EL.head
54+ case mx of
55+ Nothing -> feedDone $ Chunks [held, toInsert]
56+ Just x
57+ | atFront `S.isPrefixOf` x -> do
58+ let y = S.drop (S.length atFront) x
59+ let stream = Chunks [held, atFront, toInsert, y]
60+ feedDone stream
61+ | whole `S.isInfixOf` x -> do
62+ let (before, rest) = S.breakSubstring whole x
63+ let after = S.drop (S.length whole) rest
64+ feedDone $ Chunks [held, before, whole, toInsert, after]
65+ | x `S.isPrefixOf` atFront -> go
66+ (held `S.append` x)
67+ (S.drop (S.length x) atFront)
68+ step
69+ | otherwise -> do
70+ let (held', atFront', x') = getOverlap whole x
71+ feedCont held' atFront' $ Chunks [held, x']
72+ where
73+ --feedDone :: Stream S.ByteString -> Iteratee S.ByteString IO (Step S.ByteString IO a)
74+ feedDone stream =
75+ case step of
76+ Continue k -> do
77+ step' <- lift $ runIteratee $ k stream
78+ return step'
79+ Yield b s -> return $ Yield b s
80+ Error e -> return $ Error e
81+
82+ --feedCont :: Monad m => S.ByteString -> S.ByteString -> Stream S.ByteString -> Iteratee S.ByteString m (Step S.ByteString m a)
83+ feedCont held atFront stream = do
84+ case step of
85+ Continue k -> do
86+ step' <- lift $ runIteratee $ k stream
87+ go held atFront step'
88+ Yield b s -> return $ Yield b s
89+ Error e -> return $ Error e
90+
91+getOverlap :: S.ByteString -> S.ByteString -> (S.ByteString, S.ByteString, S.ByteString)
92+getOverlap whole x =
93+ go whole
94+ where
95+ go piece
96+ | S.null piece = ("", whole, x)
97+ | piece `S.isSuffixOf` x =
98+ let x' = S.take (S.length x - S.length piece) x
99+ atFront = S.drop (S.length piece) whole
100+ in (piece, atFront, x')
101+ | otherwise = go $ S.init piece
102+
45103fixHeaders front [] = (False, front [])
46104fixHeaders front (("content-encoding", "gzip"):rest) = (True, front rest)
47105fixHeaders front (x:xs) = fixHeaders (front . (:) x) xs