Add server support for response headers

CHANGELOG.md

@@ -6,6 +6,7 @@
 * Export `toApplication` from `Servant.Server` (https://github.com/haskell-servant/servant-server/pull/29)
 * Support other Monads than just `EitherT (Int, String) IO` (https://github.com/haskell-servant/servant-server/pull/21)
 * Make methods return status code 204 if they return () (https://github.com/haskell-servant/servant-server/issues/28)
+* Add server support for response headers
 
 0.2.4
 -----

servant-server.cabal

@@ -80,6 +80,7 @@ test-suite spec
       base == 4.*
     , aeson
     , bytestring
+    , bytestring-conversion
     , directory
     , either
     , exceptions

src/Servant/Server/Internal.hs

@@ -1,45 +1,53 @@
-{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DataKinds            #-}
-{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE FlexibleContexts     #-}
-{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE FlexibleInstances    #-}
-{-# LANGUAGE TypeOperators #-}
-{-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE OverlappingInstances #-}
-{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE OverloadedStrings    #-}
+{-# LANGUAGE PolyKinds            #-}
+{-# LANGUAGE ScopedTypeVariables  #-}
+{-# LANGUAGE TypeFamilies         #-}
+{-# LANGUAGE TypeOperators        #-}
 module Servant.Server.Internal where
 
-import Control.Applicative ((<$>))
+import           Control.Applicative         ((<$>))
-import Control.Monad.Trans.Either (EitherT, runEitherT)
+import           Control.Monad.Trans.Either  (EitherT, runEitherT)
-import qualified Data.ByteString as B
+import qualified Data.ByteString             as B
-import qualified Data.ByteString.Lazy as BL
+import qualified Data.ByteString.Lazy        as BL
-import Data.IORef (newIORef, readIORef, writeIORef)
+import           Data.IORef                  (newIORef, readIORef, writeIORef)
-import Data.List (unfoldr)
+import           Data.List                   (unfoldr)
-import Data.Maybe (catMaybes, fromMaybe)
+import           Data.Maybe                  (catMaybes, fromMaybe)
-import Data.Monoid (Monoid, mempty, mappend)
+import           Data.Monoid                 (Monoid, mappend, mempty)
-import Data.String (fromString)
+import           Data.String                 (fromString)
-import Data.String.Conversions (cs, (<>))
+import           Data.String.Conversions     (cs, (<>))
-import Data.Text.Encoding (decodeUtf8, encodeUtf8)
+import           Data.Text                   (Text)
-import Data.Text (Text)
+import qualified Data.Text                   as T
-import qualified Data.Text as T
+import           Data.Text.Encoding          (decodeUtf8, encodeUtf8)
-import Data.Typeable
+import           Data.Typeable
-import GHC.TypeLits (KnownSymbol, symbolVal)
+import           GHC.TypeLits                (KnownSymbol, symbolVal)
-import Network.HTTP.Types hiding (Header)
+import           Network.HTTP.Types          hiding (Header, ResponseHeaders)
-import Network.Wai ( Response, Request, ResponseReceived, Application
+import           Network.Wai                 (Application, Request, Response,
-                   , pathInfo, requestBody, strictRequestBody
+                                              ResponseReceived, lazyRequestBody,
-                   , lazyRequestBody, requestHeaders, requestMethod,
+                                              pathInfo, rawQueryString,
-                    rawQueryString, responseLBS)
+                                              requestBody, requestHeaders,
-import Servant.API ( QueryParams, QueryParam, QueryFlag, ReqBody, Header
+                                              requestMethod, responseLBS,
-                   , MatrixParams, MatrixParam, MatrixFlag
+                                              strictRequestBody)
-                   , Capture, Get, Delete, Put, Post, Patch, Raw, (:>), (:<|>)(..))
+import           Servant.API                 ((:<|>) (..), (:>), Capture,
-import Servant.API.ContentTypes ( AllCTRender(..), AcceptHeader(..)
+                                              Delete, Get, Header, MatrixFlag,
-                                   , AllCTUnrender(..),)
+                                              MatrixParam, MatrixParams, Patch,
-import Servant.Common.Text (FromText, fromText)
+                                              Post, Put, QueryFlag, QueryParam,
+                                              QueryParams, Raw, ReqBody)
+import           Servant.API.ContentTypes    (AcceptHeader (..),
+                                              AllCTRender (..),
+                                              AllCTUnrender (..))
+import           Servant.API.ResponseHeaders (Headers, getResponse, getHeaders)
+import           Servant.Common.Text         (FromText, fromText)
 
 
 data ReqBodyState = Uncalled
                   | Called !B.ByteString
                   | Done !B.ByteString
 
+
 toApplication :: RoutingApplication -> Application
 toApplication ra request respond = do
   reqBodyRef <- newIORef Uncalled
@@ -286,6 +294,7 @@ instance ( AllCTRender ctypes a
         respond $ failWith WrongMethod
     | otherwise = respond $ failWith NotFound
 
+-- '()' ==> 204 No Content
 instance HasServer (Get ctypes ()) where
   type ServerT (Get ctypes ()) m = m ()
   route Proxy action request respond
@@ -299,6 +308,26 @@ instance HasServer (Get ctypes ()) where
         respond $ failWith WrongMethod
     | otherwise = respond $ failWith NotFound
 
+-- Add response headers
+instance ( AllCTRender ctypes v ) => HasServer (Get ctypes (Headers h v)) where
+  type ServerT (Get ctypes (Headers h v)) m = m (Headers h v)
+  route Proxy action request respond
+    | pathIsEmpty request && requestMethod request == methodGet = do
+      e <- runEitherT action
+      respond $ case e of
+        Right output -> do
+          let accH = fromMaybe "*/*" $ lookup hAccept $ requestHeaders request
+              headers = getHeaders output
+          case handleAcceptH (Proxy :: Proxy ctypes) (AcceptHeader accH) (getResponse output) of
+            Nothing -> failWith UnsupportedMediaType
+            Just (contentT, body) -> succeedWith $
+              responseLBS ok200 ( ("Content-Type" , cs contentT) : headers) body
+        Left (status, message) -> succeedWith $
+          responseLBS (mkStatus status (cs message)) [] (cs message)
+    | pathIsEmpty request && requestMethod request /= methodGet =
+        respond $ failWith WrongMethod
+    | otherwise = respond $ failWith NotFound
+
 -- | If you use 'Header' in one of the endpoints for your API,
 -- this automatically requires your server-side handler to be a function
 -- that takes an argument of the type specified by 'Header'.
@@ -377,6 +406,26 @@ instance HasServer (Post ctypes ()) where
         respond $ failWith WrongMethod
     | otherwise = respond $ failWith NotFound
 
+-- Add response headers
+instance ( AllCTRender ctypes v ) => HasServer (Post ctypes (Headers h v)) where
+  type ServerT (Post ctypes (Headers h v)) m = m (Headers h v)
+  route Proxy action request respond
+    | pathIsEmpty request && requestMethod request == methodPost = do
+      e <- runEitherT action
+      respond $ case e of
+        Right output -> do
+          let accH = fromMaybe "*/*" $ lookup hAccept $ requestHeaders request
+              headers = getHeaders output
+          case handleAcceptH (Proxy :: Proxy ctypes) (AcceptHeader accH) (getResponse output) of
+            Nothing -> failWith UnsupportedMediaType
+            Just (contentT, body) -> succeedWith $
+              responseLBS status201 ( ("Content-Type" , cs contentT) : headers) body
+        Left (status, message) -> succeedWith $
+          responseLBS (mkStatus status (cs message)) [] (cs message)
+    | pathIsEmpty request && requestMethod request /= methodPost =
+        respond $ failWith WrongMethod
+    | otherwise = respond $ failWith NotFound
+
 -- | When implementing the handler for a 'Put' endpoint,
 -- just like for 'Servant.API.Delete.Delete', 'Servant.API.Get.Get'
 -- and 'Servant.API.Post.Post', the handler code runs in the
@@ -387,7 +436,7 @@ instance HasServer (Post ctypes ()) where
 --
 -- If successfully returning a value, we use the type-level list, combined
 -- with the request's @Accept@ header, to encode the value for you
--- (returning a status code of 201). If there was no @Accept@ header or it
+-- (returning a status code of 200). If there was no @Accept@ header or it
 -- was @*/*@, we return encode using the first @Content-Type@ type on the
 -- list.
 instance ( AllCTRender ctypes a
@@ -423,6 +472,26 @@ instance HasServer (Put ctypes ()) where
         respond $ failWith WrongMethod
     | otherwise = respond $ failWith NotFound
 
+-- Add response headers
+instance ( AllCTRender ctypes v ) => HasServer (Put ctypes (Headers h v)) where
+  type ServerT (Put ctypes (Headers h v)) m = m (Headers h v)
+  route Proxy action request respond
+    | pathIsEmpty request && requestMethod request == methodPut = do
+      e <- runEitherT action
+      respond $ case e of
+        Right output -> do
+          let accH = fromMaybe "*/*" $ lookup hAccept $ requestHeaders request
+              headers = getHeaders output
+          case handleAcceptH (Proxy :: Proxy ctypes) (AcceptHeader accH) (getResponse output) of
+            Nothing -> failWith UnsupportedMediaType
+            Just (contentT, body) -> succeedWith $
+              responseLBS status200 ( ("Content-Type" , cs contentT) : headers) body
+        Left (status, message) -> succeedWith $
+          responseLBS (mkStatus status (cs message)) [] (cs message)
+    | pathIsEmpty request && requestMethod request /= methodPut =
+        respond $ failWith WrongMethod
+    | otherwise = respond $ failWith NotFound
+
 -- | When implementing the handler for a 'Patch' endpoint,
 -- just like for 'Servant.API.Delete.Delete', 'Servant.API.Get.Get'
 -- and 'Servant.API.Put.Put', the handler code runs in the
@@ -433,7 +502,7 @@ instance HasServer (Put ctypes ()) where
 --
 -- If successfully returning a value, we just require that its type has
 -- a 'ToJSON' instance and servant takes care of encoding it for you,
--- yielding status code 201 along the way.
+-- yielding status code 200 along the way.
 instance ( AllCTRender ctypes a
          ) => HasServer (Patch ctypes a) where
   type ServerT (Patch ctypes a) m = m a
@@ -467,6 +536,26 @@ instance HasServer (Patch ctypes ()) where
         respond $ failWith WrongMethod
     | otherwise = respond $ failWith NotFound
 
+-- Add response headers
+instance ( AllCTRender ctypes v ) => HasServer (Patch ctypes (Headers h v)) where
+  type ServerT (Patch ctypes (Headers h v)) m = m (Headers h v)
+  route Proxy action request respond
+    | pathIsEmpty request && requestMethod request == methodPatch = do
+      e <- runEitherT action
+      respond $ case e of
+        Right outpatch -> do
+          let accH = fromMaybe "*/*" $ lookup hAccept $ requestHeaders request
+              headers = getHeaders outpatch
+          case handleAcceptH (Proxy :: Proxy ctypes) (AcceptHeader accH) (getResponse outpatch) of
+            Nothing -> failWith UnsupportedMediaType
+            Just (contentT, body) -> succeedWith $
+              responseLBS status200 ( ("Content-Type" , cs contentT) : headers) body
+        Left (status, message) -> succeedWith $
+          responseLBS (mkStatus status (cs message)) [] (cs message)
+    | pathIsEmpty request && requestMethod request /= methodPatch =
+        respond $ failWith WrongMethod
+    | otherwise = respond $ failWith NotFound
+
 -- | If you use @'QueryParam' "author" Text@ in one of the endpoints for your API,
 -- this automatically requires your server-side handler to be a function
 -- that takes an argument of type @'Maybe' 'Text'@.

test/Servant/ServerSpec.hs

@@ -3,13 +3,16 @@
 {-# LANGUAGE OverloadedStrings   #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeOperators       #-}
+{-# LANGUAGE TypeSynonymInstances       #-}
+{-# LANGUAGE FlexibleInstances       #-}
 
 module Servant.ServerSpec where
 
 
-import           Control.Monad              (when)
+import           Control.Monad              (forM_, when)
 import           Control.Monad.Trans.Either (EitherT, left)
 import           Data.Aeson                 (FromJSON, ToJSON, decode', encode)
+import           Data.ByteString.Conversion ()
 import           Data.Char                  (toUpper)
 import           Data.Monoid                ((<>))
 import           Data.Proxy                 (Proxy (Proxy))
@@ -27,15 +30,17 @@ import           Network.Wai                (Application, Request, pathInfo,
 import           Network.Wai.Test           (defaultRequest, request,
                                              runSession, simpleBody)
 import           Test.Hspec                 (Spec, describe, it, shouldBe)
-import           Test.Hspec.Wai             (get, liftIO, matchStatus, post,
+import           Test.Hspec.Wai             (get, liftIO, matchHeaders,
-                                             request, shouldRespondWith, with)
+                                             matchStatus, post, request,
+                                             shouldRespondWith, with, (<:>))
 
-import           Servant.API                ((:<|>) (..), (:>), Capture, Delete,
+import           Servant.API                ((:<|>) (..), (:>),
-                                             Get, Header, JSON, MatrixFlag,
+                                             AddHeader (addHeader), Capture,
-                                             MatrixParam, MatrixParams, Patch,
+                                             Delete, Get, Header (..), Headers,
-                                             PlainText, Post, Put, QueryFlag,
+                                             JSON, MatrixFlag, MatrixParam,
-                                             QueryParam, QueryParams, Raw,
+                                             MatrixParams, Patch, PlainText,
-                                             ReqBody)
+                                             Post, Put, QueryFlag, QueryParam,
+                                             QueryParams, Raw, ReqBody)
 import           Servant.Server             (Server, serve)
 import           Servant.Server.Internal    (RouteMismatch (..))
 
@@ -85,6 +90,7 @@ spec = do
   rawSpec
   unionSpec
   errorsSpec
+  responseHeadersSpec
 
 
 type CaptureApi = Capture "legs" Integer :> Get '[JSON] Animal
@@ -531,6 +537,42 @@ unionSpec = do
       it "returns 404 if the path does not exist" $ do
         get "/nonexistent" `shouldRespondWith` 404
 
+type ResponseHeadersApi =
+       Get   '[JSON] (Headers '[Header "H1" Int, Header "H2" String] String)
+  :<|> Post  '[JSON] (Headers '[Header "H1" Int, Header "H2" String] String)
+  :<|> Put   '[JSON] (Headers '[Header "H1" Int, Header "H2" String] String)
+  :<|> Patch '[JSON] (Headers '[Header "H1" Int, Header "H2" String] String)
+
+
+responseHeadersServer :: Server ResponseHeadersApi
+responseHeadersServer = let h = return $ addHeader 5 $ addHeader "kilroy" "hi"
+  in h :<|> h :<|> h :<|> h
+
+
+responseHeadersSpec :: Spec
+responseHeadersSpec = describe "ResponseHeaders" $ do
+  with (return $ serve (Proxy :: Proxy ResponseHeadersApi) responseHeadersServer) $ do
+
+    let methods = [(methodGet, 200), (methodPost, 201), (methodPut, 200), (methodPatch, 200)]
+
+    it "includes the headers in the response" $
+      forM_ methods $ \(method, expected) ->
+        Test.Hspec.Wai.request method "/" [] ""
+          `shouldRespondWith` "\"hi\""{ matchHeaders = ["H1" <:> "5", "H2" <:> "kilroy"]
+                                      , matchStatus  = expected
+                                      }
+
+    it "responds with not found for non-existent endpoints" $
+      forM_ methods $ \(method,_) ->
+        Test.Hspec.Wai.request method "blahblah" [] ""
+          `shouldRespondWith` 404
+
+    it "returns 415 if the Accept header is not supported" $
+      forM_ methods $ \(method,_) ->
+        Test.Hspec.Wai.request method "" [(hAccept, "crazy/mime")] ""
+          `shouldRespondWith` 415
+
+
 -- | Test server error functionality.
 errorsSpec :: Spec
 errorsSpec = do