Add functions to yesod-test to support the new CSRF middleware

yesod-test/ChangeLog.md

@@ -1,3 +1,8 @@
+## 1.4.3.2
+
+* Add `addTokenFromCookie` and `addTokenFromCookieNamedToHeaderNamed`, which support the new CSRF token middleware [#1058](https://github.com/yesodweb/yesod/pull/1058)
+* Add `getRequestCookies`, which returns the cookies from the most recent request [#1058](https://github.com/yesodweb/yesod/pull/1058)
+
 ## 1.4.3.1
 
 * Improved README

yesod-test/Yesod/Test.hs

@@ -77,6 +77,8 @@ module Yesod.Test
     , addToken_
     , addNonce
     , addNonce_
+    , addTokenFromCookie
+    , addTokenFromCookieNamedToHeaderNamed
 
     -- * Assertions
     , assertEqual
@@ -93,6 +95,7 @@ module Yesod.Test
     -- * Grab information
     , getTestYesod
     , getResponse
+    , getRequestCookies
 
     -- * Debug output
     , printBody
@@ -133,6 +136,7 @@ import qualified Data.Map as M
 import qualified Web.Cookie as Cookie
 import qualified Blaze.ByteString.Builder as Builder
 import Data.Time.Clock (getCurrentTime)
+import Control.Applicative ((<$>))
 
 -- | The state used in a single test case defined using 'yit'
 --
@@ -589,6 +593,65 @@ addToken_ scope = do
 addToken :: RequestBuilder site ()
 addToken = addToken_ ""
 
+-- | Calls 'addTokenFromCookieNamedToHeaderNamed' with the 'defaultCsrfCookieName' and 'defaultCsrfHeaderName'.
+--
+-- Use this function if you're using the CSRF middleware from "Yesod.Core" and haven't customized the cookie or header name.
+--
+-- ==== __Examples__
+--
+-- > request $ do
+-- >   addTokenFromCookie
+--
+-- Since 1.4.3.2
+addTokenFromCookie :: RequestBuilder site ()
+addTokenFromCookie = addTokenFromCookieNamedToHeaderNamed defaultCsrfCookieName defaultCsrfHeaderName
+
+-- | Looks up the CSRF token stored in the cookie with the given name and adds it to the request headers. An error is thrown if the cookie can't be found.
+--
+-- Use this function if you're using the CSRF middleware from "Yesod.Core" and have customized the cookie or header name.
+--
+-- See "Yesod.Core.Handler" for details on this approach to CSRF protection.
+--
+-- ==== __Examples__
+--
+-- > import Data.CaseInsensitive (CI)
+-- > request $ do
+-- >   addTokenFromCookieNamedToHeaderNamed "cookieName" (CI "headerName")
+--
+-- Since 1.4.3.2
+addTokenFromCookieNamedToHeaderNamed :: ByteString -- ^ The name of the cookie
+                                     -> CI ByteString -- ^ The name of the header
+                                     -> RequestBuilder site ()
+addTokenFromCookieNamedToHeaderNamed cookieName headerName = do
+  cookies <- getRequestCookies
+  case M.lookup cookieName cookies of
+        Just csrfCookie -> addRequestHeader (headerName, Cookie.setCookieValue csrfCookie)
+        Nothing -> failure $ T.concat
+          [ "addTokenFromCookieNamedToHeaderNamed failed to lookup CSRF cookie with name: "
+          , T.pack $ show cookieName
+          , ". Cookies were: "
+          , T.pack $ show cookies
+          ]
+
+-- | Returns the 'Cookies' from the most recent request. If a request hasn't been made, an error is raised.
+--
+-- ==== __Examples__
+--
+-- > request $ do
+-- >   cookies <- getRequestCookies
+-- >   liftIO $ putStrLn $ "Cookies are: " ++ show cookies
+--
+-- Since 1.4.3.2
+getRequestCookies :: RequestBuilder site Cookies
+getRequestCookies = do
+  requestBuilderData <- ST.get
+  headers <- case simpleHeaders <$> rbdResponse requestBuilderData of
+                  Just h -> return h
+                  Nothing -> failure "getRequestCookies: No request has been made yet; the cookies can't be looked up."
+
+  return $ M.fromList $ map (\c -> (Cookie.setCookieName c, c)) (parseSetCookies headers)
+
+
 -- | Perform a POST request to @url@.
 --
 -- ==== __Examples__
@@ -759,7 +822,7 @@ request reqBuilder = do
             { httpVersion = H.http11
             }
         }) app
-    let newCookies = map (Cookie.parseSetCookie . snd) $ DL.filter (("Set-Cookie"==) . fst) $ simpleHeaders response
+    let newCookies = parseSetCookies $ simpleHeaders response
         cookies' = M.fromList [(Cookie.setCookieName c, c) | c <- newCookies] `M.union` cookies
     ST.put $ YesodExampleData app site cookies' (Just response)
   where
@@ -846,6 +909,10 @@ request reqBuilder = do
       , queryString = urlQuery
       }
 
+
+parseSetCookies :: [H.Header] -> [Cookie.SetCookie]
+parseSetCookies headers = map (Cookie.parseSetCookie . snd) $ DL.filter (("Set-Cookie"==) . fst) $ headers
+
 -- Yes, just a shortcut
 failure :: (MonadIO a) => T.Text -> a b
 failure reason = (liftIO $ HUnit.assertFailure $ T.unpack reason) >> error ""

yesod-test/test/main.hs

@@ -2,6 +2,10 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE QuasiQuotes #-}
+{-# LANGUAGE TypeFamilies #-}
+
 import Test.HUnit hiding (Test)
 import Test.Hspec
 
@@ -166,6 +170,26 @@ main = hspec $ do
             statusIs 200
             printBody
             bodyContains "Foo"
+    describe "CSRF with cookies/headers" $ yesodSpec CsrfApp $ do
+        yit "Should receive a CSRF cookie and add its value to the headers" $ do
+            get ("/" :: Text)
+            statusIs 200
+
+            request $ do
+                setMethod "POST"
+                setUrl ("/" :: Text)
+                addTokenFromCookie
+            statusIs 200
+        yit "Should 403 requests if we don't add the CSRF token" $ do
+            get ("/" :: Text)
+            statusIs 200
+
+            request $ do
+                setMethod "POST"
+                setUrl ("/" :: Text)
+            statusIs 403
+
+
 
 instance RenderMessage LiteApp FormMessage where
     renderMessage _ _ = defaultFormMessage
@@ -210,3 +234,26 @@ cookieApp = liteApp $ do
             setMessage "Foo"
             redirect ("/cookie/home" :: Text)
             return ()
+
+data CsrfApp = CsrfApp
+
+mkYesod "CsrfApp" [parseRoutes|
+/ HomeR GET POST
+|]
+
+instance Yesod CsrfApp where
+    yesodMiddleware = defaultYesodMiddleware . defaultCsrfMiddleware
+
+getHomeR :: Handler Html
+getHomeR = defaultLayout
+    [whamlet|
+        <p>
+            Welcome to my test application.
+    |]
+
+postHomeR :: Handler Html
+postHomeR = defaultLayout
+    [whamlet|
+        <p>
+            Welcome to my test application.
+    |]

yesod-test/yesod-test.cabal

@@ -37,7 +37,7 @@ library
                    , time
                    , blaze-builder
                    , cookie
-                   , yesod-core                >= 1.4
+                   , yesod-core                >= 1.4.14
 
     exposed-modules: Yesod.Test
                      Yesod.Test.CssQuery