{-# LANGUAGE BangPatterns #-} {-# LANGUAGE OverloadedStrings #-} module Test.Common.TestHandler (testHandler) where import Blaze.ByteString.Builder import Control.Concurrent (threadDelay) import Control.Monad import Control.Monad.Trans import qualified Data.ByteString.Char8 as S import qualified Data.ByteString.Lazy.Char8 as L import Data.List import qualified Data.Map as Map import Data.Maybe import Data.Monoid import Snap.Iteratee hiding (Enumerator) import qualified Snap.Iteratee as I import Snap.Core import Snap.Util.FileServe import Snap.Util.FileUploads import Snap.Util.GZip import System.Directory import Test.Common.Rot13 (rot13) ------------------------------------------------------------------------------ -- timeout handling ------------------------------------------------------------------------------ timeoutTickleHandler :: Snap () timeoutTickleHandler = do noCompression -- FIXME: remove this when zlib-bindings and -- zlib-enumerator support gzip stream flushing modifyResponse $ setResponseBody (trickleOutput 6) . setContentType "text/plain" setTimeout 2 where trickleOutput :: Int -> Enumerator Builder IO a trickleOutput n = concatEnums $ dots `interleave` delays where enumOne = enumList 1 [fromByteString ".\n", flush] delay st = do liftIO $ threadDelay 1000000 returnI st interleave x0 y0 = (go id x0 y0) [] where go !dl [] ys = dl . (++ys) go !dl xs [] = dl . (++xs) go !dl (x:xs) (y:ys) = go (dl . (x:) . (y:)) xs ys dots = replicate n enumOne delays = replicate n delay ------------------------------------------------------------------------------ pongHandler :: Snap () pongHandler = modifyResponse $ setResponseBody (enumBuilder $ fromByteString "PONG") . setContentType "text/plain" . setContentLength 4 echoUriHandler :: Snap () echoUriHandler = do req <- getRequest writeBS $ rqURI req echoHandler :: Snap () echoHandler = transformRequestBody returnI rot13Handler :: Snap () rot13Handler = transformRequestBody f where f origStep = do mbX <- I.head maybe (enumEOF origStep) (feedStep origStep) mbX feedStep origStep b = do let x = toByteString b let e = enumBuilder $ fromByteString $ rot13 x step <- lift $ runIteratee $ e origStep f step bigResponseHandler :: Snap () bigResponseHandler = do let sz = 4000000 let s = L.take sz $ L.cycle $ L.fromChunks [S.replicate 400000 '.'] modifyResponse $ setContentLength sz writeLBS s responseHandler :: Snap () responseHandler = do !code <- liftM (read . S.unpack . fromMaybe "503") $ getParam "code" modifyResponse $ setResponseCode code writeBS $ S.pack $ show code uploadForm :: Snap () uploadForm = do modifyResponse $ setContentType "text/html" writeBS form where form = S.concat [ "Upload a file\n" , "

Upload some text/plain files:

\n" , "
\n" , "\n" , "\n" , "
" ] uploadHandler :: Snap () uploadHandler = do liftIO $ createDirectoryIfMissing True tmpdir handleFileUploads tmpdir defaultUploadPolicy partPolicy hndl where isRight (Left _) = False isRight (Right _) = True f (_, Left _) = error "impossible" f (p, Right x) = (fromMaybe "-" $ partFileName p, x) hndl xs' = do let xs = [ f x | x <- xs', isRight (snd x) ] files <- mapM (\(x,fp) -> do c <- liftIO $ S.readFile fp return (x,c)) xs let m = sort files params <- liftM (Prelude.map (\(a,b) -> (a,S.concat b)) . Map.toAscList . rqParams) getRequest modifyResponse $ setContentType "text/plain" writeBuilder $ buildRqParams params `mappend` buildFiles m builder _ [] = mempty builder ty ((k,v):xs) = mconcat [ fromByteString ty , fromByteString ":\n" , fromByteString k , fromByteString "\nValue:\n" , fromByteString v , fromByteString "\n\n" , builder ty xs ] buildRqParams = builder "Param" buildFiles = builder "File" tmpdir = "dist/filetmp" partPolicy partInfo = if partContentType partInfo == "text/plain" then allowWithMaximumSize 200000 else disallow testHandler :: Snap () testHandler = withCompression $ route [ ("pong" , pongHandler ) , ("echo" , echoHandler ) , ("rot13" , rot13Handler ) , ("echoUri" , echoUriHandler ) , ("fileserve" , serveDirectory "testserver/static") , ("bigresponse" , bigResponseHandler ) , ("respcode/:code" , responseHandler ) , ("upload/form" , uploadForm ) , ("upload/handle" , uploadHandler ) , ("timeout/tickle" , timeoutTickleHandler ) ]