From d5b66d35d563d1a4234c9833921de2512996ad02 Mon Sep 17 00:00:00 2001
From: Michael Snoyman <michael@snoyman.com>
Date: Sun, 4 Aug 2013 08:55:42 +0300
Subject: [PATCH] Allow overriding status code in error handlers

---
 yesod-core/Yesod/Core/Internal/Response.hs    | 19 ++++++++++++++++++-
 yesod-core/Yesod/Core/Internal/Run.hs         |  9 ++++++---
 .../test/YesodCoreTest/ErrorHandling.hs       | 17 ++++++++++++++++-
 yesod-core/yesod-core.cabal                   |  2 +-
 4 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/yesod-core/Yesod/Core/Internal/Response.hs b/yesod-core/Yesod/Core/Internal/Response.hs
index 8e0870f5..b71ea5c2 100644
--- a/yesod-core/Yesod/Core/Internal/Response.hs
+++ b/yesod-core/Yesod/Core/Internal/Response.hs
@@ -32,7 +32,7 @@ yarToResponse :: Monad m
               -> YesodRequest
               -> m Response
 yarToResponse (YRWai a) _ _ = return a
-yarToResponse (YRPlain s hs ct c newSess) saveSession yreq = do
+yarToResponse (YRPlain s' hs ct c newSess) saveSession yreq = do
     extraHeaders <- do
         let nsToken = maybe
                 newSess
@@ -50,6 +50,23 @@ yarToResponse (YRPlain s hs ct c newSess) saveSession yreq = do
         go (ContentSource body) = ResponseSource s finalHeaders body
         go (ContentDontEvaluate c') = go c'
     return $ go c
+  where
+    s
+        | s' == defaultStatus = H.status200
+        | otherwise = s'
+
+-- | Indicates that the user provided no specific status code to be used, and
+-- therefore the default status code should be used. For normal responses, this
+-- would be a 200 response, whereas for error responses this would be an
+-- appropriate status code.
+--
+-- For more information on motivation for this, see:
+--
+-- https://groups.google.com/d/msg/yesodweb/vHDBzyu28TM/bezCvviWp4sJ
+--
+-- Since 1.2.3.1
+defaultStatus :: H.Status
+defaultStatus = H.mkStatus (-1) "INVALID DEFAULT STATUS"
 
 -- | Convert Header to a key/value pair.
 headerToPair :: Header
diff --git a/yesod-core/Yesod/Core/Internal/Run.hs b/yesod-core/Yesod/Core/Internal/Run.hs
index 38411b58..35f1d3fd 100644
--- a/yesod-core/Yesod/Core/Internal/Run.hs
+++ b/yesod-core/Yesod/Core/Internal/Run.hs
@@ -73,15 +73,18 @@ runHandler rhe@RunHandlerEnv {..} handler yreq = withInternalState $ \resState -
     state <- liftIO $ I.readIORef istate
     let finalSession = ghsSession state
     let headers = ghsHeaders state
-    let contents = either id (HCContent H.status200 . toTypedContent) contents'
+    let contents = either id (HCContent defaultStatus . toTypedContent) contents'
     let handleError e = flip runInternalState resState $ do
             yar <- rheOnError e yreq
                 { reqSession = finalSession
                 }
             case yar of
-                YRPlain _ hs ct c sess ->
+                YRPlain status' hs ct c sess ->
                     let hs' = appEndo headers hs
-                     in return $ YRPlain (getStatus e) hs' ct c sess
+                        status
+                            | status' == defaultStatus = getStatus e
+                            | otherwise = status'
+                     in return $ YRPlain status hs' ct c sess
                 YRWai _ -> return yar
     let sendFile' ct fp p =
             return $ YRPlain H.status200 (appEndo headers []) ct (ContentFile fp p) finalSession
diff --git a/yesod-core/test/YesodCoreTest/ErrorHandling.hs b/yesod-core/test/YesodCoreTest/ErrorHandling.hs
index ee225a8a..883f2900 100644
--- a/yesod-core/test/YesodCoreTest/ErrorHandling.hs
+++ b/yesod-core/test/YesodCoreTest/ErrorHandling.hs
@@ -12,6 +12,7 @@ import Text.Hamlet (hamlet)
 import qualified Data.ByteString.Lazy as L
 import qualified Data.ByteString.Char8 as S8
 import Control.Exception (SomeException, try)
+import Network.HTTP.Types (mkStatus)
 
 data App = App
 
@@ -22,9 +23,14 @@ mkYesod "App" [parseRoutes|
 /after_runRequestBody AfterRunRequestBodyR POST
 /error-in-body ErrorInBodyR GET
 /error-in-body-noeval ErrorInBodyNoEvalR GET
+/override-status OverrideStatusR GET
 |]
 
-instance Yesod App
+overrideStatus = mkStatus 15 "OVERRIDE"
+
+instance Yesod App where
+    errorHandler (InvalidArgs ["OVERRIDE"]) = sendResponseStatus overrideStatus ("OH HAI" :: String)
+    errorHandler x = defaultErrorHandler x
 
 getHomeR :: Handler Html
 getHomeR = do
@@ -65,6 +71,9 @@ getErrorInBodyR = do
 getErrorInBodyNoEvalR :: Handler (DontFullyEvaluate Html)
 getErrorInBodyNoEvalR = fmap DontFullyEvaluate getErrorInBodyR
 
+getOverrideStatusR :: Handler ()
+getOverrideStatusR = invalidArgs ["OVERRIDE"]
+
 errorHandlingTest :: Spec
 errorHandlingTest = describe "Test.ErrorHandling" $ do
       it "says not found" caseNotFound
@@ -72,6 +81,7 @@ errorHandlingTest = describe "Test.ErrorHandling" $ do
       it "says 'There was an error' after runRequestBody" caseAfter
       it "error in body == 500" caseErrorInBody
       it "error in body, no eval == 200" caseErrorInBodyNoEval
+      it "can override status code" caseOverrideStatus
 
 runner :: Session () -> IO ()
 runner f = toWaiApp App >>= runSession f
@@ -125,3 +135,8 @@ caseErrorInBodyNoEval = do
     case eres of
         Left (_ :: SomeException) -> return ()
         Right _ -> error "Expected an exception"
+
+caseOverrideStatus :: IO ()
+caseOverrideStatus = runner $ do
+    res <- request defaultRequest { pathInfo = ["override-status"] }
+    assertStatus 15 res
diff --git a/yesod-core/yesod-core.cabal b/yesod-core/yesod-core.cabal
index 22a801b3..c44dd59b 100644
--- a/yesod-core/yesod-core.cabal
+++ b/yesod-core/yesod-core.cabal
@@ -1,5 +1,5 @@
 name:            yesod-core
-version:         1.2.3
+version:         1.2.3.1
 license:         MIT
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>