snap-server-1.1.2.1: test/Test/Common/TestHandler.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
module Test.Common.TestHandler (testHandler) where
------------------------------------------------------------------------------
import Control.Concurrent (threadDelay)
import Control.Exception (throwIO)
import Control.Monad (liftM)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.ByteString.Builder (Builder, byteString)
import Data.ByteString.Builder.Extra (flush)
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Lazy.Char8 as L
import Data.List (sort)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import Data.Monoid (Monoid (mappend, mconcat, mempty))
------------------------------------------------------------------------------
import Snap.Core (Request (rqParams, rqURI), Snap, getParam, getRequest, logError, modifyResponse, redirect, route, rqClientAddr, rqClientPort, setContentLength, setContentType, setHeader, setResponseBody, setResponseCode, setTimeout, transformRequestBody, writeBS, writeBuilder, writeLBS)
import Snap.Internal.Debug ()
import Snap.Util.FileServe (serveDirectory)
import Snap.Util.FileUploads (PartInfo (partContentType, partFileName), allowWithMaximumSize, defaultUploadPolicy, disallow, handleFileUploads)
import Snap.Util.GZip (noCompression, withCompression)
import System.Directory (createDirectoryIfMissing)
import System.IO.Streams (OutputStream)
import qualified System.IO.Streams as Streams
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 10)
. setContentType "text/plain"
setTimeout 2
badTimeoutTickleHandler :: Snap ()
badTimeoutTickleHandler = do
noCompression -- FIXME: remove this when zlib-bindings and
-- zlib-enumerator support gzip stream flushing
modifyResponse $ setResponseBody (trickleOutput 10)
. setContentType "text/plain"
setTimeout 2
trickleOutput :: Int -> OutputStream Builder -> IO (OutputStream Builder)
trickleOutput n os = do
Streams.fromList dots >>= Streams.mapM f >>= Streams.supplyTo os
return os
where
dots = replicate n ".\n"
f x = threadDelay 1000000 >> return (byteString x `mappend` flush)
------------------------------------------------------------------------------
pongHandler :: Snap ()
pongHandler = modifyResponse $
setResponseBody body .
setContentType "text/plain" .
setContentLength 4
where
body os = do Streams.write (Just $ byteString "PONG") os
return os
echoUriHandler :: Snap ()
echoUriHandler = do
req <- getRequest
writeBS $ rqURI req
echoHandler :: Snap ()
echoHandler = transformRequestBody return
rot13Handler :: Snap ()
rot13Handler = transformRequestBody (Streams.map rot13)
bigResponseHandler :: Snap ()
bigResponseHandler = do
let sz = 4000000
let s = L.take sz $ L.cycle $ L.fromChunks [S.replicate 400000 '.']
modifyResponse $ setContentLength $ fromIntegral 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 [ "<html><head><title>Upload a file</title></head><body>\n"
, "<p>Upload some <code>text/plain</code> files:</p>\n"
, "<form method='post' "
, "enctype='multipart/form-data' "
, "action='/upload/handle'>\n"
, "<input type='file' "
, "accept='text/plain' "
, "multiple='true' "
, "name='file'></input>\n"
, "<input type='submit' name='Submit'></input>\n"
, "</form></body></html>" ]
uploadHandler :: Snap ()
uploadHandler = do
logError "uploadHandler"
liftIO $ createDirectoryIfMissing True tmpdir
files <- handleFileUploads tmpdir defaultUploadPolicy partPolicy hndl
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
where
f p = fromMaybe "-" $ partFileName p
hndl _ (Left e) = throwIO e
hndl partInfo (Right fp) = do
!c <- liftIO $ S.readFile fp
return $! (f partInfo, c)
builder _ [] = mempty
builder ty ((k,v):xs) =
mconcat [ byteString ty
, byteString ":\n"
, byteString k
, byteString "\nValue:\n"
, byteString v
, byteString "\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
serverHeaderHandler :: Snap ()
serverHeaderHandler = modifyResponse $ setHeader "Server" "foo"
chunkedResponse :: Snap ()
chunkedResponse = writeBS "chunked"
remoteAddrPort :: Snap ()
remoteAddrPort = do
rq <- getRequest
let addr = rqClientAddr rq
let port = rqClientPort rq
let out = S.concat [ addr, ":", S.pack (show port) ]
modifyResponse $ setContentLength $ fromIntegral $ S.length out
writeBS out
testHandler :: Snap ()
testHandler = withCompression $
route [ ("pong" , pongHandler )
, ("redirect" , redirect "/pong" )
, ("echo" , echoHandler )
, ("rot13" , rot13Handler )
, ("echoUri" , echoUriHandler )
, ("remoteAddrPort" , remoteAddrPort )
, ("fileserve" , noCompression >>
serveDirectory "testserver/static")
, ("bigresponse" , bigResponseHandler )
, ("respcode/:code" , responseHandler )
, ("upload/form" , uploadForm )
, ("upload/handle" , uploadHandler )
, ("timeout/tickle" , timeoutTickleHandler )
, ("timeout/badtickle" , badTimeoutTickleHandler )
, ("server-header" , serverHeaderHandler )
, ("chunked" , chunkedResponse )
]