From 839b56b032c388a92b0f07b80c07b13de7a8701c Mon Sep 17 00:00:00 2001
From: Sibi Prabakaran <sibi@psibi.in>
Date: Thu, 13 Jul 2017 11:10:54 +0530
Subject: [PATCH] Implement replaceOrAddHeader function

---
 yesod-core/Yesod/Core/Handler.hs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/yesod-core/Yesod/Core/Handler.hs b/yesod-core/Yesod/Core/Handler.hs
index 0afced9a..4b580cc9 100644
--- a/yesod-core/Yesod/Core/Handler.hs
+++ b/yesod-core/Yesod/Core/Handler.hs
@@ -114,6 +114,7 @@ module Yesod.Core.Handler
     , deleteCookie
     , addHeader
     , setHeader
+    , replaceOrAddHeader
     , setLanguage
       -- ** Content caching and expiration
     , cacheSeconds
@@ -787,6 +788,23 @@ setHeader :: MonadHandler m => Text -> Text -> m ()
 setHeader = addHeader
 {-# DEPRECATED setHeader "Please use addHeader instead" #-}
 
+replaceOrAddHeader :: MonadHandler m => Text -> Text -> m ()
+replaceOrAddHeader a b =
+  let header = Header (encodeUtf8 a) (encodeUtf8 b)
+  in modify $ \g -> g {ghsHeaders = replaceHeader header (ghsHeaders g)}
+  where
+    sameHeaderName :: Header -> Header -> Bool
+    sameHeaderName (Header n1 _) (Header n2 _) = n1 == n2
+    sameHeaderName _ _ = False
+
+    replaceHeader :: Header -> Endo [Header] -> Endo [Header]
+    replaceHeader header endo =
+      let allHeaders :: [Header] = appEndo endo []
+      in Endo
+           (\y ->
+              (header : y) ++
+              filter (\x -> not (sameHeaderName x header)) allHeaders)
+
 -- | Set the Cache-Control header to indicate this response should be cached
 -- for the given number of seconds.
 cacheSeconds :: MonadHandler m => Int -> m ()