Merge pull request #1417 from psibi/header-yesod

Add replaceOrAddHeader function to Yesod.Core.Handler module

yesod-core/ChangeLog.md

@@ -1,3 +1,7 @@
+## 1.4.36
+
+* Add `replaceOrAddHeader` function in Yesod.Core.Handler module. [1416](https://github.com/yesodweb/yesod/issues/1416)
+
 ## 1.4.35.1
 
 * TH fix for GHC 8.2

yesod-core/Yesod/Core/Handler.hs

@@ -10,6 +10,7 @@
 {-# LANGUAGE TypeFamilies               #-}
 {-# LANGUAGE RankNTypes                 #-}
 {-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE ScopedTypeVariables        #-}
 ---------------------------------------------------------
 --
 -- Module        : Yesod.Handler
@@ -114,6 +115,7 @@ module Yesod.Core.Handler
     , deleteCookie
     , addHeader
     , setHeader
+    , replaceOrAddHeader
     , setLanguage
       -- ** Content caching and expiration
     , cacheSeconds
@@ -206,7 +208,7 @@ import Control.Monad.Trans.Class (lift)
 
 import           Data.Aeson                    (ToJSON(..))
 import qualified Data.Text                     as T
-import           Data.Text.Encoding            (decodeUtf8With, encodeUtf8)
+import           Data.Text.Encoding            (decodeUtf8With, encodeUtf8, decodeUtf8)
 import           Data.Text.Encoding.Error      (lenientDecode)
 import qualified Data.Text.Lazy                as TL
 import           Text.Blaze.Html.Renderer.Utf8 (renderHtml)
@@ -787,6 +789,40 @@ setHeader :: MonadHandler m => Text -> Text -> m ()
 setHeader = addHeader
 {-# DEPRECATED setHeader "Please use addHeader instead" #-}
 
+-- | Replace an existing header with a new value or add a new header
+-- if not present.
+--
+-- Note that, while the data type used here is 'Text', you must provide only
+-- ASCII value to be HTTP compliant.
+--
+-- @since 1.4.36
+replaceOrAddHeader :: MonadHandler m => Text -> Text -> m ()
+replaceOrAddHeader a b =
+  modify $ \g -> g {ghsHeaders = replaceHeader (ghsHeaders g)}
+  where
+    repHeader = Header (encodeUtf8 a) (encodeUtf8 b)
+
+    sameHeaderName :: Header -> Header -> Bool
+    sameHeaderName (Header n1 _) (Header n2 _) = T.toLower (decodeUtf8 n1) == T.toLower (decodeUtf8 n2)
+    sameHeaderName _ _ = False
+
+    replaceIndividualHeader :: [Header] -> [Header]
+    replaceIndividualHeader [] = [repHeader]
+    replaceIndividualHeader xs = aux xs []
+      where
+        aux [] acc = acc ++ [repHeader]
+        aux (x:xs') acc =
+          if sameHeaderName repHeader x
+            then acc ++
+                 [repHeader] ++
+                 (filter (\header -> not (sameHeaderName header repHeader)) xs')
+            else aux xs' (acc ++ [x])
+
+    replaceHeader :: Endo [Header] -> Endo [Header]
+    replaceHeader endo =
+      let allHeaders :: [Header] = appEndo endo []
+      in Endo (\rest -> replaceIndividualHeader allHeaders ++ rest)
+
 -- | Set the Cache-Control header to indicate this response should be cached
 -- for the given number of seconds.
 cacheSeconds :: MonadHandler m => Int -> m ()

yesod-core/test/YesodCoreTest.hs

@@ -6,6 +6,7 @@ import YesodCoreTest.Exceptions
 import YesodCoreTest.Widget
 import YesodCoreTest.Media
 import YesodCoreTest.Links
+import YesodCoreTest.Header
 import YesodCoreTest.NoOverloadedStrings
 import YesodCoreTest.InternalRequest
 import YesodCoreTest.ErrorHandling
@@ -27,6 +28,7 @@ import Test.Hspec
 
 specs :: Spec
 specs = do
+      headerTest
       cleanPathTest
       exceptionsTest
       widgetTest

yesod-core/test/YesodCoreTest/Header.hs

@@ -0,0 +1,77 @@
+{-# LANGUAGE OverloadedStrings, TemplateHaskell, QuasiQuotes,
+  TypeFamilies, MultiParamTypeClasses, ViewPatterns #-}
+
+module YesodCoreTest.Header
+  ( headerTest
+  , Widget
+  , resourcesApp
+  ) where
+
+import Data.Text (Text)
+import Network.HTTP.Types (decodePathSegments)
+import Network.Wai
+import Network.Wai.Test
+import Test.Hspec
+import Yesod.Core
+
+data App =
+  App
+
+mkYesod
+  "App"
+  [parseRoutes|
+/header1 Header1R GET
+/header2 Header2R GET
+/header3 Header3R GET
+|]
+
+instance Yesod App
+
+getHeader1R :: Handler RepPlain
+getHeader1R = do
+  addHeader "hello" "world"
+  return $ RepPlain $ toContent ("header test" :: Text)
+
+getHeader2R :: Handler RepPlain
+getHeader2R = do
+  addHeader "hello" "world"
+  replaceOrAddHeader "hello" "sibi"
+  return $ RepPlain $ toContent ("header test" :: Text)
+
+getHeader3R :: Handler RepPlain
+getHeader3R = do
+  addHeader "hello" "world"
+  addHeader "michael" "snoyman"
+  addHeader "yesod" "framework"
+  replaceOrAddHeader "yesod" "book"
+  return $ RepPlain $ toContent ("header test" :: Text)
+
+runner :: Session () -> IO ()
+runner f = toWaiApp App >>= runSession f
+
+addHeaderTest :: IO ()
+addHeaderTest =
+  runner $ do
+    res <- request defaultRequest {pathInfo = decodePathSegments "/header1"}
+    assertHeader "hello" "world" res
+
+multipleHeaderTest :: IO ()
+multipleHeaderTest =
+  runner $ do
+    res <- request defaultRequest {pathInfo = decodePathSegments "/header2"}
+    assertHeader "hello" "sibi" res
+
+header3Test :: IO ()
+header3Test = do
+  runner $ do
+    res <- request defaultRequest {pathInfo = decodePathSegments "/header3"}
+    assertHeader "hello" "world" res
+    assertHeader "michael" "snoyman" res
+    assertHeader "yesod" "book" res
+
+headerTest :: Spec
+headerTest =
+  describe "Test.Header" $ do
+    it "addHeader" addHeaderTest
+    it "multiple header" multipleHeaderTest
+    it "persist headers" header3Test

yesod-core/yesod-core.cabal

@@ -1,5 +1,5 @@
 name:            yesod-core
-version:         1.4.35.1
+version:         1.4.36
 license:         MIT
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>
@@ -150,6 +150,7 @@ test-suite tests
                    YesodCoreTest.Auth
                    YesodCoreTest.Cache
                    YesodCoreTest.CleanPath
+                   YesodCoreTest.Header
                    YesodCoreTest.Csrf
                    YesodCoreTest.ErrorHandling
                    YesodCoreTest.Exceptions