@@ -25,8 +25,10 @@ module Network.Wai.Test
2525
2626import Network.Wai
2727import Network.Wai.Internal (ResponseReceived (ResponseReceived))
28+import Control.Applicative ((<$>))
2829import Control.Monad.IO.Class (liftIO)
29−import Control.Monad.Trans.State (StateT, evalStateT)
30+import Control.Monad.Trans.Class (lift)
31+import qualified Control.Monad.Trans.State as ST
3032import Control.Monad.Trans.Reader (ReaderT, runReaderT, ask)
3133import Control.Monad (unless)
3234import Control.DeepSeq (deepseq)
@@ -34,9 +36,10 @@ import Control.Exception (throwIO, Exception)
3436import Data.Typeable (Typeable)
3537import Data.Map (Map)
3638import qualified Data.Map as Map
39+import qualified Web.Cookie as Cookie
3740import Data.ByteString (ByteString)
3841import qualified Data.ByteString.Char8 as S8
39−import Blaze.ByteString.Builder (toLazyByteString)
42+import Blaze.ByteString.Builder (toLazyByteString, toByteString)
4043import qualified Blaze.ByteString.Builder as B
4144import qualified Data.ByteString.Lazy as L
4245import qualified Data.ByteString.Lazy.Char8 as L8
@@ -47,18 +50,28 @@ import qualified Data.Text as T
4750import qualified Data.Text.Encoding as TE
4851import Data.IORef
4952import Data.Monoid (mempty, mappend)
53+import Data.Time.Clock (getCurrentTime)
5054
51−type Session = ReaderT Application (StateT ClientState IO)
55+type Session = ReaderT Application (ST.StateT ClientState IO)
56+
57+type ClientCookies = Map ByteString Cookie.SetCookie
5258
5359data ClientState = ClientState
54− { _clientCookies :: Map ByteString ByteString
60+ { clientCookies :: ClientCookies
5561 }
5662
63+getClientCookies :: Session ClientCookies
64+getClientCookies = clientCookies <$> lift ST.get
65+
66+modifyClientCookies :: (ClientCookies -> ClientCookies) -> Session ()
67+modifyClientCookies f =
68+ lift (ST.modify (\cs -> cs { clientCookies = f $ clientCookies cs }))
69+
5770initState :: ClientState
5871initState = ClientState Map.empty
5972
6073runSession :: Session a -> Application -> IO a
61−runSession session app = evalStateT (runReaderT session app) initState
74+runSession session app = ST.evalStateT (runReaderT session app) initState
6275
6376data SRequest = SRequest
6477 { simpleRequest :: Request
@@ -93,6 +106,40 @@ setRawPathInfo r rawPinfo =
93106 dropFrontSlash ("":path) = path
94107 dropFrontSlash path = path
95108
109+addCookiesToRequest :: Request -> Session Request
110+addCookiesToRequest req = do
111+ oldClientCookies <- getClientCookies
112+ let requestPath = "/" `T.append` T.intercalate "/" (pathInfo req)
113+ currentUTCTime <- liftIO getCurrentTime
114+ let cookiesForRequest =
115+ Map.filter
116+ (\c -> checkCookieTime currentUTCTime c
117+ && checkCookiePath requestPath c)
118+ oldClientCookies
119+ let cookiePairs = [ (Cookie.setCookieName c, Cookie.setCookieValue c)
120+ | c <- map snd $ Map.toList cookiesForRequest
121+ ]
122+ let cookieValue = toByteString $ Cookie.renderCookies cookiePairs
123+ return $ req { requestHeaders = ("Cookie", cookieValue):requestHeaders req }
124+ where checkCookieTime t c =
125+ case Cookie.setCookieExpires c of
126+ Nothing -> True
127+ Just t' -> t < t'
128+ checkCookiePath p c =
129+ case Cookie.setCookiePath c of
130+ Nothing -> True
131+ Just p' -> p' `S8.isPrefixOf` TE.encodeUtf8 p
132+
133+extractSetCookieFromSResponse :: SResponse -> Session SResponse
134+extractSetCookieFromSResponse response = do
135+ let setCookieHeaders =
136+ filter (("Set-Cookie"==) . fst) $ simpleHeaders response
137+ let newClientCookies = map (Cookie.parseSetCookie . snd) setCookieHeaders
138+ modifyClientCookies
139+ (Map.union
140+ (Map.fromList [(Cookie.setCookieName c, c) | c <- newClientCookies ]))
141+ return response
142+
96143srequest :: SRequest -> Session SResponse
97144srequest (SRequest req bod) = do
98145 app <- ask
@@ -103,12 +150,12 @@ srequest (SRequest req bod) = do
103150 [] -> ([], S.empty)
104151 x:y -> (y, x)
105152 }
106− liftIO $ do
153+ req'' <- addCookiesToRequest req'
154+ response <- liftIO $ do
107155 ref <- newIORef $ error "runResponse gave no result"
108− ResponseReceived <- app req' (runResponse ref)
156+ ResponseReceived <- app req'' (runResponse ref)
109157 readIORef ref
110− -- FIXME cookie processing
111− --return sres
158+ extractSetCookieFromSResponse response
112159
113160runResponse :: IORef SResponse -> Response -> IO ResponseReceived
114161runResponse ref res = do