1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
|
{-# 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 )
]
|