From 33982b2112cd4e83edf52813f753d4c7aa5cedf2 Mon Sep 17 00:00:00 2001
From: Maximilian Tagher <feedback.tagher@gmail.com>
Date: Mon, 22 Jun 2015 22:26:02 -0700
Subject: [PATCH] Add CSRF protection functions/middleware that support AJAX
 requests

---
 yesod-core/ChangeLog.md               |   4 +
 yesod-core/Yesod/Core.hs              |   6 ++
 yesod-core/Yesod/Core/Class/Yesod.hs  |  51 ++++++++++
 yesod-core/Yesod/Core/Handler.hs      | 139 +++++++++++++++++++++++++-
 yesod-core/test/YesodCoreTest.hs      |   2 +
 yesod-core/test/YesodCoreTest/Csrf.hs |  92 +++++++++++++++++
 yesod-core/yesod-core.cabal           |   1 +
 7 files changed, 293 insertions(+), 2 deletions(-)
 create mode 100644 yesod-core/test/YesodCoreTest/Csrf.hs

diff --git a/yesod-core/ChangeLog.md b/yesod-core/ChangeLog.md
index 4a230a62..4abe0533 100644
--- a/yesod-core/ChangeLog.md
+++ b/yesod-core/ChangeLog.md
@@ -1,3 +1,7 @@
+## 1.4.14
+
+* Add CSRF protection functions and middleware based on HTTP cookies and headers [#1017](https://github.com/yesodweb/yesod/pull/1017)
+
 ## 1.4.13
 
 * Add mkYesodGeneral, which allows creating sites with polymorphic type parameters [#1055](https://github.com/yesodweb/yesod/pull/1055)
diff --git a/yesod-core/Yesod/Core.hs b/yesod-core/Yesod/Core.hs
index c1a55280..2a157ca5 100644
--- a/yesod-core/Yesod/Core.hs
+++ b/yesod-core/Yesod/Core.hs
@@ -57,6 +57,12 @@ module Yesod.Core
     , clientSessionDateCacher
     , loadClientSession
     , Header(..)
+    -- * CSRF protection
+    , defaultCsrfMiddleware
+    , defaultCsrfSetCookieMiddleware
+    , csrfSetCookieMiddleware
+    , defaultCsrfCheckMiddleware
+    , csrfCheckMiddleware
     -- * JS loaders
     , ScriptLoadPosition (..)
     , BottomOfHeadAsync
diff --git a/yesod-core/Yesod/Core/Class/Yesod.hs b/yesod-core/Yesod/Core/Class/Yesod.hs
index 14224af8..0f6408bd 100644
--- a/yesod-core/Yesod/Core/Class/Yesod.hs
+++ b/yesod-core/Yesod/Core/Class/Yesod.hs
@@ -55,6 +55,7 @@ import           Yesod.Core.Types
 import           Yesod.Core.Internal.Session
 import           Yesod.Core.Widget
 import Control.Monad.Trans.Class (lift)
+import Data.CaseInsensitive (CI)
 
 -- | Define settings for a Yesod applications. All methods have intelligent
 -- defaults, and therefore no implementation is required.
@@ -411,6 +412,56 @@ authorizationCheck = do
                               void $ notAuthenticated
             Unauthorized s' -> permissionDenied s'
 
+-- | Calls 'csrfCheckMiddleware' with 'isWriteRequest', 'defaultCsrfHeaderName', and 'defaultCsrfParamName' as parameters.
+--
+-- Since 1.4.14
+defaultCsrfCheckMiddleware :: Yesod site => HandlerT site IO res -> HandlerT site IO res
+defaultCsrfCheckMiddleware handler = do
+    csrfCheckMiddleware
+        handler
+        (getCurrentRoute >>= maybe (return False) isWriteRequest)
+        defaultCsrfHeaderName
+        defaultCsrfParamName
+
+-- | Looks up the CSRF token from the request headers or POST parameters. If the value doesn't match the token stored in the session,
+-- this function throws a 'PermissionDenied' error.
+--
+-- For details, see the "AJAX CSRF protection" section of 'Yesod.Core.Handler'.
+--
+-- Since 1.4.14
+csrfCheckMiddleware :: Yesod site
+                    => HandlerT site IO res
+                    -> HandlerT site IO Bool -- ^ Whether or not to perform the CSRF check.
+                    -> CI S8.ByteString -- ^ The header name to lookup the CSRF token from.
+                    -> Text -- ^ The POST parameter name to lookup the CSRF token from.
+                    -> HandlerT site IO res
+csrfCheckMiddleware handler shouldCheckFn headerName paramName = do
+    shouldCheck <- shouldCheckFn
+    when shouldCheck (checkCsrfHeaderOrParam headerName paramName)
+    handler
+
+-- | Calls 'csrfSetCookieMiddleware' with the 'defaultCsrfCookieName'.
+--
+-- Since 1.4.14
+defaultCsrfSetCookieMiddleware :: Yesod site => HandlerT site IO res -> HandlerT site IO res
+defaultCsrfSetCookieMiddleware handler = csrfSetCookieMiddleware handler (def { setCookieName = defaultCsrfCookieName })
+
+-- | Takes a 'SetCookie' and overrides its value with a CSRF token, then sets the cookie. See 'setCsrfCookieWithCookie'.
+--
+-- For details, see the "AJAX CSRF protection" section of 'Yesod.Core.Handler'.
+--
+-- Since 1.4.14
+csrfSetCookieMiddleware :: Yesod site => HandlerT site IO res -> SetCookie -> HandlerT site IO res
+csrfSetCookieMiddleware handler cookie = setCsrfCookieWithCookie cookie >> handler
+
+-- | Calls 'defaultCsrfSetCookieMiddleware' and 'defaultCsrfCheckMiddleware'. Use this midle 
+--
+-- For details, see the "AJAX CSRF protection" section of 'Yesod.Core.Handler'.
+--
+-- Since 1.4.14
+defaultCsrfMiddleware :: Yesod site => HandlerT site IO res -> HandlerT site IO res
+defaultCsrfMiddleware = defaultCsrfSetCookieMiddleware . defaultCsrfCheckMiddleware
+
 -- | Convert a widget to a 'PageContent'.
 widgetToPageContent :: (Eq (Route site), Yesod site)
                     => WidgetT site IO ()
diff --git a/yesod-core/Yesod/Core/Handler.hs b/yesod-core/Yesod/Core/Handler.hs
index dd048e9f..e3e57d78 100644
--- a/yesod-core/Yesod/Core/Handler.hs
+++ b/yesod-core/Yesod/Core/Handler.hs
@@ -153,6 +153,24 @@ module Yesod.Core.Handler
     , cached
     , cachedBy
     , stripHandlerT
+      -- * AJAX CSRF protection
+
+      -- $ajaxCSRFOverview
+
+      -- ** Setting CSRF Cookies
+    , setCsrfCookie
+    , setCsrfCookieWithCookie
+    , defaultCsrfCookieName
+      -- ** Looking up CSRF Headers
+    , checkCsrfHeaderNamed
+    , hasValidCsrfHeaderNamed
+    , defaultCsrfHeaderName
+      -- ** Looking up CSRF POST Parameters
+    , hasValidCsrfParamNamed
+    , checkCsrfParamNamed
+    , defaultCsrfParamName
+      -- ** Checking CSRF Headers or POST Parameters
+    , checkCsrfHeaderOrParam
     ) where
 
 import           Data.Time                     (UTCTime, addUTCTime,
@@ -186,6 +204,8 @@ import qualified Data.ByteString               as S
 import qualified Data.ByteString.Lazy          as L
 import qualified Data.Map                      as Map
 
+import           Data.Byteable                 (constEqBytes)
+
 import           Control.Arrow                 ((***))
 import qualified Data.ByteString.Char8         as S8
 import           Data.Monoid                   (Endo (..), mappend, mempty)
@@ -219,6 +239,8 @@ import Data.Conduit (Source, transPipe, Flush (Flush), yield, Producer
                    )
 import qualified Yesod.Core.TypeCache as Cache
 import qualified Data.Word8 as W8
+import qualified Data.Foldable as Fold
+import Data.Default
 
 get :: MonadHandler m => m GHState
 get = liftHandlerT $ HandlerT $ I.readIORef . handlerState
@@ -479,10 +501,10 @@ setUltDestReferer = do
 redirectUltDest :: (RedirectUrl (HandlerSite m) url, MonadHandler m)
                 => url -- ^ default destination if nothing in session
                 -> m a
-redirectUltDest def = do
+redirectUltDest defaultDestination = do
     mdest <- lookupSession ultDestKey
     deleteSession ultDestKey
-    maybe (redirect def) redirect mdest
+    maybe (redirect defaultDestination) redirect mdest
 
 -- | Remove a previously set ultimate destination. See 'setUltDest'.
 clearUltDest :: MonadHandler m => m ()
@@ -1264,3 +1286,116 @@ stripHandlerT (HandlerT f) getSub toMaster newRoute = HandlerT $ \hd -> do
             }
         , handlerToParent = toMaster
         }
+
+-- $ajaxCSRFOverview
+-- When a user has authenticated with your site, all requests made from the browser to your server will include the session information that you use to verify that the user is logged in. 
+-- Unfortunately, this allows attackers to make unwanted requests on behalf of the user by e.g. submitting an HTTP request to your site when the user visits theirs. 
+-- This is known as a <https://en.wikipedia.org/wiki/Cross-site_request_forgery Cross Site Request Forgery> (CSRF) attack.
+--
+-- To combat this attack, you need a way to verify that the request is valid.
+-- This is achieved by generating a random string ("token"), storing it in your encrypted session so that the server can look it up (see 'reqToken'), and adding the token to HTTP requests made to your server.
+-- When a request comes in, the token in the request is compared to the one from the encrypted session. If they match, you can be sure the request is valid.
+--
+-- Yesod implements this behavior in two ways:
+--
+-- (1) The yesod-form package <http://www.yesodweb.com/book/forms#forms_running_forms stores the CSRF token in a hidden field> in the form, then validates it with functions like 'Yesod.Form.Functions.runFormPost'.
+--
+-- (2) Yesod can store the CSRF token in a cookie which is accessible by Javascript. Requests made by Javascript can lookup this cookie and add it as a header to requests. The server then checks the token in the header against the one in the encrypted session.
+--
+-- The form-based approach has the advantage of working for users with Javascript disabled, while adding the token to the headers with Javascript allows things like submitting JSON or binary data in AJAX requests. Yesod supports checking for a CSRF token in either the POST parameters of the form ('checkCsrfHeaderNamed'), the headers ('checkCsrfHeaderNamed'), or both options ('checkCsrfHeaderOrParam').
+--
+-- The easiest way to check both sources is to add the 'defaultCsrfMiddleware' to your Yesod Middleware.
+
+-- | The default cookie name for the CSRF token ("XSRF-TOKEN").
+-- 
+-- Since 1.4.14
+defaultCsrfCookieName :: S8.ByteString
+defaultCsrfCookieName = "XSRF-TOKEN"
+
+-- | Sets a cookie with a CSRF token, using 'defaultCsrfCookieName' for the cookie name.
+-- 
+-- Since 1.4.14
+setCsrfCookie :: MonadHandler m => m ()
+setCsrfCookie = setCsrfCookieWithCookie def { setCookieName = defaultCsrfCookieName }
+
+-- | Takes a 'SetCookie' and overrides its value with a CSRF token, then sets the cookie.
+-- 
+-- Since 1.4.14
+setCsrfCookieWithCookie :: MonadHandler m => SetCookie -> m ()
+setCsrfCookieWithCookie cookie  = do
+    mCsrfToken <- reqToken <$> getRequest
+    Fold.forM_ mCsrfToken (\token -> setCookie $ cookie { setCookieValue = encodeUtf8 token })
+
+-- | The default header name for the CSRF token ("X-XSRF-TOKEN").
+-- 
+-- Since 1.4.14
+defaultCsrfHeaderName :: CI S8.ByteString
+defaultCsrfHeaderName = "X-XSRF-TOKEN"
+
+-- | Takes a header name to lookup a CSRF token. If the value doesn't match the token stored in the session,
+-- this function throws a 'PermissionDenied' error.
+-- 
+-- Since 1.4.14
+checkCsrfHeaderNamed :: MonadHandler m => CI S8.ByteString -> m ()
+checkCsrfHeaderNamed headerName = do
+  valid <- hasValidCsrfHeaderNamed headerName
+  unless valid (permissionDenied csrfErrorMessage)
+
+-- | Takes a header name to lookup a CSRF token, and returns whether the value matches the token stored in the session.
+-- 
+-- Since 1.4.14
+hasValidCsrfHeaderNamed :: MonadHandler m => CI S8.ByteString -> m Bool
+hasValidCsrfHeaderNamed headerName = do
+  mCsrfToken  <- reqToken <$> getRequest
+  mXsrfHeader <- lookupHeader headerName
+
+  return $ validCsrf mCsrfToken mXsrfHeader
+
+-- CSRF Parameter checking
+
+-- | The default parameter name for the CSRF token ("_token")
+--
+-- Since 1.4.14
+defaultCsrfParamName :: Text
+defaultCsrfParamName = "_token"
+
+-- | Takes a POST parameter name to lookup a CSRF token. If the value doesn't match the token stored in the session,
+-- this function throws a 'PermissionDenied' error.
+--
+-- Since 1.4.14
+checkCsrfParamNamed :: MonadHandler m => Text -> m ()
+checkCsrfParamNamed paramName = do
+  valid <- hasValidCsrfParamNamed paramName
+  unless valid (permissionDenied csrfErrorMessage)
+
+-- | Takes a POST parameter name to lookup a CSRF token, and returns whether the value matches the token stored in the session.
+--
+-- Since 1.4.14
+hasValidCsrfParamNamed :: MonadHandler m => Text -> m Bool
+hasValidCsrfParamNamed paramName = do
+  mCsrfToken  <- reqToken <$> getRequest
+  mCsrfParam <- lookupPostParam paramName
+
+  return $ validCsrf mCsrfToken (encodeUtf8 <$> mCsrfParam)
+
+-- | Checks that a valid CSRF token is present in either the request headers or POST parameters.
+-- If the value doesn't match the token stored in the session, this function throws a 'PermissionDenied' error.
+--
+-- Since 1.4.14
+checkCsrfHeaderOrParam :: MonadHandler m
+                       => CI S8.ByteString -- ^ The header name to lookup the CSRF token
+                       -> Text -- ^ The POST parameter name to lookup the CSRF token
+                       -> m ()
+checkCsrfHeaderOrParam headerName paramName = do
+  validHeader <- hasValidCsrfHeaderNamed headerName
+  validParam <- hasValidCsrfParamNamed paramName
+  unless (validHeader || validParam) (permissionDenied csrfErrorMessage)
+
+validCsrf :: Maybe Text -> Maybe S.ByteString -> Bool
+-- It's important to use constant-time comparison (constEqBytes) in order to avoid timing attacks.
+validCsrf (Just token) (Just param) = encodeUtf8 token `constEqBytes` param
+validCsrf Nothing            _param = True
+validCsrf (Just _token)     Nothing = False
+
+csrfErrorMessage :: Text
+csrfErrorMessage = "A valid CSRF token wasn't present in HTTP headers or POST parameters. Check the Yesod.Core.Handler docs of the yesod-core package for details on CSRF protection."
\ No newline at end of file
diff --git a/yesod-core/test/YesodCoreTest.hs b/yesod-core/test/YesodCoreTest.hs
index 9d40de40..94a104db 100644
--- a/yesod-core/test/YesodCoreTest.hs
+++ b/yesod-core/test/YesodCoreTest.hs
@@ -21,6 +21,7 @@ import qualified YesodCoreTest.Reps as Reps
 import qualified YesodCoreTest.Auth as Auth
 import qualified YesodCoreTest.LiteApp as LiteApp
 import qualified YesodCoreTest.Ssl as Ssl
+import qualified YesodCoreTest.Csrf as Csrf
 
 import Test.Hspec
 
@@ -47,3 +48,4 @@ specs = do
       LiteApp.specs
       Ssl.unsecSpec
       Ssl.sslOnlySpec
+      Csrf.csrfSpec
diff --git a/yesod-core/test/YesodCoreTest/Csrf.hs b/yesod-core/test/YesodCoreTest/Csrf.hs
new file mode 100644
index 00000000..8121b6da
--- /dev/null
+++ b/yesod-core/test/YesodCoreTest/Csrf.hs
@@ -0,0 +1,92 @@
+{-# LANGUAGE TypeFamilies, QuasiQuotes, TemplateHaskell, MultiParamTypeClasses, OverloadedStrings #-}
+
+module YesodCoreTest.Csrf (csrfSpec, Widget, resourcesApp) where
+
+import Yesod.Core
+
+import Test.Hspec
+import Network.Wai
+import Network.Wai.Test
+import Web.Cookie
+import qualified Data.Map as Map
+import Data.ByteString.Lazy (fromStrict)
+import Data.Monoid ((<>))
+
+data App = App
+
+mkYesod "App" [parseRoutes|
+/ HomeR GET POST
+|]
+
+instance Yesod App 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.
+    |]
+
+runner :: Session () -> IO ()
+runner f = toWaiApp App >>= runSession f
+
+csrfSpec :: Spec
+csrfSpec = describe "A Yesod application with the defaultCsrfMiddleware" $ do
+    it "serves a includes a cookie in a GET request" $ runner $ do
+        res <- request defaultRequest
+        assertStatus 200 res
+        assertClientCookieExists "Should have an XSRF-TOKEN cookie" defaultCsrfCookieName
+
+    it "200s write requests with the correct CSRF header, but no param" $ runner $ do
+        getRes <- request defaultRequest
+        assertStatus 200 getRes
+        csrfValue <- fmap setCookieValue requireCsrfCookie
+        postRes <- request (defaultRequest { requestMethod = "POST", requestHeaders = [(defaultCsrfHeaderName, csrfValue)] })
+        assertStatus 200 postRes
+
+    it "200s write requests with the correct CSRF param, but no header" $ runner $ do
+        getRes <- request defaultRequest
+        assertStatus 200 getRes
+        csrfValue <- fmap setCookieValue requireCsrfCookie
+
+        let body = "_token=" <> csrfValue
+        postRes <- srequest $ SRequest (defaultRequest { requestMethod = "POST", requestHeaders = [("Content-Type","application/x-www-form-urlencoded")] }) (fromStrict body)
+        assertStatus 200 postRes
+
+
+    it "403s write requests without the CSRF header" $ runner $ do
+        res <- request (defaultRequest { requestMethod = "POST" })
+        assertStatus 403 res
+
+    it "403s write requests with the wrong CSRF header" $ runner $ do
+        getRes <- request defaultRequest
+        assertStatus 200 getRes
+        csrfValue <- fmap setCookieValue requireCsrfCookie
+
+        res <- request (defaultRequest { requestMethod = "POST", requestHeaders = [(defaultCsrfHeaderName, csrfValue <> "foo")] })
+        assertStatus 403 res
+        
+    it "403s write requests with the wrong CSRF param" $ runner $ do
+        getRes <- request defaultRequest
+        assertStatus 200 getRes
+        csrfValue <- fmap setCookieValue requireCsrfCookie
+
+        let body = "_token=" <> (csrfValue <> "foo")
+        postRes <- srequest $ SRequest (defaultRequest { requestMethod = "POST", requestHeaders = [("Content-Type","application/x-www-form-urlencoded")] }) (fromStrict body)
+        assertStatus 403 postRes
+
+
+requireCsrfCookie :: Session SetCookie
+requireCsrfCookie = do
+    cookies <- getClientCookies
+    case Map.lookup defaultCsrfCookieName cookies of
+        Just c -> return c
+        Nothing -> error "Failed to lookup CSRF cookie"
diff --git a/yesod-core/yesod-core.cabal b/yesod-core/yesod-core.cabal
index 0385f859..ef348885 100644
--- a/yesod-core/yesod-core.cabal
+++ b/yesod-core/yesod-core.cabal
@@ -69,6 +69,7 @@ library
                    , word8
                    , auto-update
                    , semigroups
+                   , byteable
 
     exposed-modules: Yesod.Core
                      Yesod.Core.Content