module Network.Wai.Middleware.Cache
( cache
, cacheNoBody
, CacheBackend(..)
, responseToLBS
) where
import Blaze.ByteString.Builder (Builder, toLazyByteString)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as S8
import qualified Data.ByteString.Lazy as LZ
import Data.IORef
import Network.Wai (Middleware, Request, Response,
requestBody, responseToStream,
responseStatus, mapResponseHeaders)
import Network.HTTP.Types.Status (statusCode)
data CacheBackend cacheContainer cacheKey cacheVal =
CacheBackend {
keyFromReq :: Request -> ByteString -> IO cacheKey
, toCache :: Request -> ByteString -> IO Bool
, addToCache :: cacheContainer -> cacheKey -> cacheVal -> IO ()
, actionOnCache :: Request -> Response -> IO ()
, actionOnCacheMiss :: Request -> Response -> IO ()
, responseToCacheVal :: Response -> IO cacheVal
, cacheValToResponse :: cacheVal -> Response
, lookupCache :: cacheContainer -> cacheKey -> IO (Maybe cacheVal)
, cacheContainer :: cacheContainer
}
cache :: CacheBackend cc ck cv
-> Middleware
cache cb app req sendResponse = do
(req',body) <- getRequestBody req
caching <- toCache cb req' body
if not caching
then app req' sendResponse
else do
(req'',_) <- getRequestBody req'
cacheKey <- keyFromReq cb req'' body
found <- lookupCache cb (cacheContainer cb) cacheKey
maybe (app req'' (addToCacheAndRespond cb sendResponse req cacheKey))
(respondFromCache cb sendResponse req'')
found
cacheNoBody :: CacheBackend cc ck cv
-> Middleware
cacheNoBody cb app req sendResponse = do
caching <- toCache cb req S8.empty
if not caching
then app req sendResponse
else do
cacheKey <- keyFromReq cb req S8.empty
found <- lookupCache cb (cacheContainer cb) cacheKey
maybe (app req (addToCacheAndRespond cb sendResponse req cacheKey))
(respondFromCache cb sendResponse req)
found
addXCacheHeader :: Response -> Response
addXCacheHeader = mapResponseHeaders (("X-Cached","true"):)
respondFromCache :: CacheBackend cc ck cv
-> (Response -> IO b)
-> Request
-> cv
-> IO b
respondFromCache cb sendResponse r cachedVal = do
let response = cacheValToResponse cb cachedVal
actionOnCache cb r response
sendResponse (addXCacheHeader response)
addToCacheAndRespond :: CacheBackend cc ck cv
-> (Response -> IO b)
-> Request
-> ck
-> Response
-> IO b
addToCacheAndRespond cb sendResponse req key r = do
let code = statusCode (responseStatus r)
if (code >= 200) && (code < 400)
then do
cacheVal <- responseToCacheVal cb r
addToCache cb (cacheContainer cb) key cacheVal
actionOnCacheMiss cb req r
sendResponse (cacheValToResponse cb cacheVal)
else
sendResponse r
getRequestBody :: Request -> IO (Request, S8.ByteString)
getRequestBody req = do
let loop front = do
bs <- requestBody req
if S8.null bs
then return $ front []
else loop $ front . (bs:)
body <- loop id
ichunks <- newIORef body
let rbody = atomicModifyIORef ichunks $ \chunks ->
case chunks of
[] -> ([], S8.empty)
x:y -> (y, x)
let req' = req { requestBody = rbody }
return (req', S8.concat body)
responseToLBS :: Response -> IO LZ.ByteString
responseToLBS response = do
let (_,_,f) = responseToStream response
f $ \streamingBody -> do
builderRef <- newIORef mempty
let add :: Builder -> IO ()
add b = atomicModifyIORef builderRef $ \builder -> (builder `mappend` b,())
flush :: IO ()
flush = return ()
streamingBody add flush
fmap toLazyByteString (readIORef builderRef)