From a797c2e5d4a1f99b2e80e20deced86ff51cff0a6 Mon Sep 17 00:00:00 2001
From: Bob Long <robertjflong@gmail.com>
Date: Sun, 1 May 2016 16:31:01 +0100
Subject: [PATCH] Add laxSameSiteSessions and strictSameSiteSessions

---
 yesod-core/Yesod/Core.hs                      |  2 ++
 yesod-core/Yesod/Core/Class/Yesod.hs          | 18 ++++++++++++++-
 yesod-core/test/YesodCoreTest.hs              |  3 ++-
 yesod-core/test/YesodCoreTest/Ssl.hs          | 16 ++++++++++++-
 .../test/YesodCoreTest/StubLaxSameSite.hs     | 23 +++++++++++++++++++
 .../test/YesodCoreTest/StubStrictSameSite.hs  | 23 +++++++++++++++++++
 6 files changed, 82 insertions(+), 3 deletions(-)
 create mode 100644 yesod-core/test/YesodCoreTest/StubLaxSameSite.hs
 create mode 100644 yesod-core/test/YesodCoreTest/StubStrictSameSite.hs

diff --git a/yesod-core/Yesod/Core.hs b/yesod-core/Yesod/Core.hs
index a63eba8b..83b1fe0e 100644
--- a/yesod-core/Yesod/Core.hs
+++ b/yesod-core/Yesod/Core.hs
@@ -53,6 +53,8 @@ module Yesod.Core
     , envClientSessionBackend
     , clientSessionBackend
     , sslOnlySessions
+    , laxSameSiteSessions
+    , strictSameSiteSessions
     , sslOnlyMiddleware
     , clientSessionDateCacher
     , loadClientSession
diff --git a/yesod-core/Yesod/Core/Class/Yesod.hs b/yesod-core/Yesod/Core/Class/Yesod.hs
index 5147fed7..f1eca3f8 100644
--- a/yesod-core/Yesod/Core/Class/Yesod.hs
+++ b/yesod-core/Yesod/Core/Class/Yesod.hs
@@ -49,7 +49,7 @@ import qualified Text.Blaze.Html5                   as TBH
 import           Text.Hamlet
 import           Text.Julius
 import qualified Web.ClientSession                  as CS
-import           Web.Cookie                         (parseCookies)
+import           Web.Cookie                         (parseCookies, sameSiteLax, sameSiteStrict, SameSiteOption)
 import           Web.Cookie                         (SetCookie (..))
 import           Yesod.Core.Types
 import           Yesod.Core.Internal.Session
@@ -366,6 +366,22 @@ sslOnlySessions = (fmap . fmap) secureSessionCookies
     setSecureBit cookie = cookie { setCookieSecure = True }
     secureSessionCookies = customizeSessionCookies setSecureBit
 
+-- | Helps defend against CSRF attacks by setting the SameSite attribute on
+-- session cookies to "Lax".
+laxSameSiteSessions :: IO (Maybe SessionBackend) -> IO (Maybe SessionBackend)
+laxSameSiteSessions = sameSiteSession sameSiteLax
+
+-- | Helps defend against CSRF attacks by setting the SameSite attribute on
+-- session cookies to "Strict".
+strictSameSiteSessions :: IO (Maybe SessionBackend) -> IO (Maybe SessionBackend)
+strictSameSiteSessions = sameSiteSession sameSiteStrict
+
+sameSiteSession :: SameSiteOption -> IO (Maybe SessionBackend) -> IO (Maybe SessionBackend)
+sameSiteSession s = (fmap . fmap) secureSessionCookies
+  where
+    sameSite cookie = cookie { setCookieSameSite = (pure s) }
+    secureSessionCookies = customizeSessionCookies sameSite
+
 -- | Apply a Strict-Transport-Security header with the specified timeout to
 -- all responses so that browsers will rewrite all http links to https
 -- until the timeout expires. For security, the max-age of the STS header
diff --git a/yesod-core/test/YesodCoreTest.hs b/yesod-core/test/YesodCoreTest.hs
index 94a104db..7c0db6fa 100644
--- a/yesod-core/test/YesodCoreTest.hs
+++ b/yesod-core/test/YesodCoreTest.hs
@@ -1,5 +1,5 @@
 {-# LANGUAGE CPP                        #-}
-module YesodCoreTest (specs) where 
+module YesodCoreTest (specs) where
 
 import YesodCoreTest.CleanPath
 import YesodCoreTest.Exceptions
@@ -48,4 +48,5 @@ specs = do
       LiteApp.specs
       Ssl.unsecSpec
       Ssl.sslOnlySpec
+      Ssl.sameSiteSpec
       Csrf.csrfSpec
diff --git a/yesod-core/test/YesodCoreTest/Ssl.hs b/yesod-core/test/YesodCoreTest/Ssl.hs
index b6162c0f..c1411b93 100644
--- a/yesod-core/test/YesodCoreTest/Ssl.hs
+++ b/yesod-core/test/YesodCoreTest/Ssl.hs
@@ -1,6 +1,8 @@
 {-# LANGUAGE TypeFamilies, QuasiQuotes, TemplateHaskell, MultiParamTypeClasses, OverloadedStrings #-}
-module YesodCoreTest.Ssl ( sslOnlySpec, unsecSpec ) where
+module YesodCoreTest.Ssl ( sslOnlySpec, unsecSpec, sameSiteSpec ) where
 import qualified YesodCoreTest.StubSslOnly as Ssl
+import qualified YesodCoreTest.StubLaxSameSite as LaxSameSite
+import qualified YesodCoreTest.StubStrictSameSite as StrictSameSite
 import qualified YesodCoreTest.StubUnsecured as Unsecured
 import Yesod.Core
 import Test.Hspec
@@ -62,3 +64,15 @@ unsecSpec = describe "A Yesod application with sslOnly off" $ do
   where
     atHome = homeFixtureFor Unsecured.App
     isNotSecure c = not $ Cookie.setCookieSecure c
+
+sameSiteSpec :: Spec
+sameSiteSpec = describe "A Yesod application" $ do
+    it "can set a Lax SameSite option" $
+        laxHome $ "_SESSION" `cookieShouldSatisfy` isLax
+    it "can set a Strict SameSite option" $
+        strictHome $ "_SESSION" `cookieShouldSatisfy` isStrict
+  where
+    laxHome = homeFixtureFor LaxSameSite.App
+    strictHome = homeFixtureFor StrictSameSite.App
+    isLax = (== Just Cookie.sameSiteLax) . Cookie.setCookieSameSite
+    isStrict = (== Just Cookie.sameSiteStrict) . Cookie.setCookieSameSite
diff --git a/yesod-core/test/YesodCoreTest/StubLaxSameSite.hs b/yesod-core/test/YesodCoreTest/StubLaxSameSite.hs
new file mode 100644
index 00000000..365c9a07
--- /dev/null
+++ b/yesod-core/test/YesodCoreTest/StubLaxSameSite.hs
@@ -0,0 +1,23 @@
+{-# LANGUAGE TypeFamilies, QuasiQuotes, TemplateHaskell, MultiParamTypeClasses, OverloadedStrings #-}
+module YesodCoreTest.StubLaxSameSite ( App ( App ) ) where
+
+import Yesod.Core
+import qualified Web.ClientSession                  as CS
+
+data App = App
+
+mkYesod "App" [parseRoutes|
+/ HomeR GET
+|]
+
+instance Yesod App where
+    yesodMiddleware = defaultYesodMiddleware . (sslOnlyMiddleware 120)
+    makeSessionBackend _ = laxSameSiteSessions $
+        fmap Just $ defaultClientSessionBackend 120 CS.defaultKeyFile
+
+getHomeR :: Handler Html
+getHomeR = defaultLayout
+    [whamlet|
+        <p>
+            Welcome to my test application.
+    |]
diff --git a/yesod-core/test/YesodCoreTest/StubStrictSameSite.hs b/yesod-core/test/YesodCoreTest/StubStrictSameSite.hs
new file mode 100644
index 00000000..0324178e
--- /dev/null
+++ b/yesod-core/test/YesodCoreTest/StubStrictSameSite.hs
@@ -0,0 +1,23 @@
+{-# LANGUAGE TypeFamilies, QuasiQuotes, TemplateHaskell, MultiParamTypeClasses, OverloadedStrings #-}
+module YesodCoreTest.StubStrictSameSite ( App ( App ) ) where
+
+import Yesod.Core
+import qualified Web.ClientSession                  as CS
+
+data App = App
+
+mkYesod "App" [parseRoutes|
+/ HomeR GET
+|]
+
+instance Yesod App where
+    yesodMiddleware = defaultYesodMiddleware . (sslOnlyMiddleware 120)
+    makeSessionBackend _ = strictSameSiteSessions $
+        fmap Just $ defaultClientSessionBackend 120 CS.defaultKeyFile
+
+getHomeR :: Handler Html
+getHomeR = defaultLayout
+    [whamlet|
+        <p>
+            Welcome to my test application.
+    |]