Beef up testing

Co-authored-by: Georgi Lyubenov <godzbanebane@gmail.com>

changelog.d/1571

@@ -0,0 +1,8 @@
+synopsis: Support UVerb in servant-auth-server
+prs: #1571
+issues: #1570
+description: {
+UVerb endpoints are now supported by servant-auth-server and can be used under the
+Auth combinator when writing servers. It is still unsupported by
+servant-auth-client.
+}

servant-auth/servant-auth-server/servant-auth-server.cabal

@@ -129,6 +129,7 @@ test-suite spec
     , lens-aeson  >= 1.0.2     && < 1.2
     , warp        >= 3.2.25    && < 3.4
     , wreq        >= 0.5.2.1   && < 0.6
+    , text        >= 1.2.3.0   && < 2.1
   other-modules:
       Servant.Auth.ServerSpec
   default-language: Haskell2010

servant-auth/servant-auth-server/src/Servant/Auth/Server/Internal.hs

@@ -54,15 +54,12 @@ instance ( n ~ 'S ('S 'Z)
 
       makeCookies :: AuthResult v -> IO (SetCookieList ('S ('S 'Z)))
       makeCookies authResult = do
-        xsrf <- makeXsrfCookie cookieSettings
-        fmap (Just xsrf `SetCookieCons`) $
         case authResult of
           (Authenticated v) -> do
             ejwt <- makeSessionCookie cookieSettings jwtSettings v
-              case ejwt of
-                Nothing  -> return $ Nothing `SetCookieCons` SetCookieNil
-                Just jwt -> return $ Just jwt `SetCookieCons` SetCookieNil
-            _ -> return $ Nothing `SetCookieCons` SetCookieNil
+            xsrf <- makeXsrfCookie cookieSettings
+            return $ Just xsrf `SetCookieCons` (ejwt `SetCookieCons` SetCookieNil)
+          _ -> return $ Nothing `SetCookieCons` (Nothing `SetCookieCons` SetCookieNil)
 
       go :: (AuthResult v -> ServerT api Handler)
          -> (AuthResult v, SetCookieList n)

servant-auth/servant-auth-server/src/Servant/Auth/Server/Internal/AddSetCookie.hs

@@ -11,6 +11,7 @@ import           Data.Tagged              (Tagged (..))
 import qualified Network.HTTP.Types       as HTTP
 import           Network.Wai              (mapResponseHeaders)
 import           Servant
+import           Servant.API.UVerb.Union
 import           Servant.API.Generic
 import           Servant.Server.Generic
 import           Web.Cookie
@@ -33,12 +34,24 @@ type family AddSetCookieApiVerb a where
   AddSetCookieApiVerb (Headers ls a) = Headers (Header "Set-Cookie" SetCookie ': ls) a
   AddSetCookieApiVerb a = Headers '[Header "Set-Cookie" SetCookie] a
 
+#if MIN_VERSION_servant_server(0,18,1)
+type family MapAddSetCookieApiVerb (as :: [*]) where
+  MapAddSetCookieApiVerb '[] = '[]
+  MapAddSetCookieApiVerb (a ': as) = (AddSetCookieApiVerb a ': MapAddSetCookieApiVerb as)
+#endif
+
 type family AddSetCookieApi a :: *
 type instance AddSetCookieApi (a :> b) = a :> AddSetCookieApi b
 type instance AddSetCookieApi (a :<|> b) = AddSetCookieApi a :<|> AddSetCookieApi b
+#if MIN_VERSION_servant_server(0,19,0)
 type instance AddSetCookieApi (NamedRoutes api) = AddSetCookieApi (ToServantApi api)
+#endif
 type instance AddSetCookieApi (Verb method stat ctyps a)
   = Verb method stat ctyps (AddSetCookieApiVerb a)
+#if MIN_VERSION_servant_server(0,18,1)
+type instance AddSetCookieApi (UVerb method ctyps as)
+  = UVerb method ctyps (MapAddSetCookieApiVerb as)
+#endif
 type instance AddSetCookieApi Raw = Raw
 #if MIN_VERSION_servant_server(0,15,0)
 type instance AddSetCookieApi (Stream method stat framing ctyps a)
@@ -57,7 +70,7 @@ instance {-# OVERLAPS #-} AddSetCookies ('S n) oldb newb
   => AddSetCookies ('S n) (a -> oldb) (a -> newb) where
   addSetCookies cookies oldfn = addSetCookies cookies . oldfn
 
-instance AddSetCookies 'Z orig orig where
+instance (orig1 ~ orig2) => AddSetCookies 'Z orig1 orig2 where
   addSetCookies _ = id
 
 instance {-# OVERLAPPABLE #-}

servant-auth/servant-auth-server/test/Servant/Auth/ServerSpec.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE TypeApplications #-}
 module Servant.Auth.ServerSpec (spec) where
 
 #if !MIN_VERSION_servant_server(0,16,0)
@@ -25,6 +26,7 @@ import           Data.Aeson                          (FromJSON, ToJSON, Value,
 import           Data.Aeson.Lens                     (_JSON)
 import qualified Data.ByteString                     as BS
 import qualified Data.ByteString.Lazy                as BSL
+import           Data.Text                           (Text, pack)
 import           Data.CaseInsensitive                (mk)
 import           Data.Foldable                       (find)
 import           Data.Monoid
@@ -40,6 +42,7 @@ import           Network.HTTP.Types                  (Status, status200,
 import           Network.Wai                         (responseLBS)
 import           Network.Wai.Handler.Warp            (testWithApplication)
 import           Network.Wreq                        (Options, auth, basicAuth,
+                                                      checkResponse,
                                                       cookieExpiryTime, cookies,
                                                       defaults, get, getWith, postWith,
                                                       header, oauth2Bearer,
@@ -183,8 +186,21 @@ cookieAuthSpec
         it "fails with no XSRF header or cookie" $ \port -> property
                                                          $ \(user :: User) -> do
           jwt <- createJWT theKey (newJWSHeader ((), HS256)) (claims $ toJSON user)
-          opts <- addJwtToCookie cookieCfg jwt
-          getWith opts (url port) `shouldHTTPErrorWith` status401
+          opts' <- addJwtToCookie cookieCfg jwt
+          let opts = opts' & checkResponse .~ Just mempty
+          resp <- getWith opts (url port)
+          resp ^. responseStatus `shouldBe` status401
+          (resp ^. responseCookieJar) `shouldNotHaveCookies` ["XSRF-TOKEN"]
+
+          -- Validating that the XSRF cookie isn't added for UVerb routes either.
+          -- These routes can return a 401 response directly without using throwError / throwAll,
+          -- which revealed a bug:
+          --
+          -- https://github.com/haskell-servant/servant/issues/1570#issuecomment-1076374449
+          resp <- getWith opts (url port ++ "/uverb")
+          resp ^. responseStatus `shouldBe` status401
+          (resp ^. responseCookieJar) `shouldNotHaveCookies` ["XSRF-TOKEN"]
+
 
         it "succeeds if XSRF header and cookie match, and JWT is valid" $ \port -> property
                                                                         $ \(user :: User) -> do
@@ -406,13 +422,14 @@ type API auths
     = Auth auths User :>
         ( Get '[JSON] Int
        :<|> ReqBody '[JSON] Int :> Post '[JSON] Int
-       :<|> NamedRoutes DummyRoutes
+       :<|> "named" :> NamedRoutes DummyRoutes
        :<|> "header" :> Get '[JSON] (Headers '[Header "Blah" Int] Int)
 #if MIN_VERSION_servant_server(0,15,0)
        :<|> "stream" :> StreamGet NoFraming OctetStream (SourceIO BS.ByteString)
 #endif
        :<|> "raw" :> Raw
         )
+      :<|> "uverb" :> Auth auths User :> UVerb 'GET '[JSON] '[WithStatus 200 Int, WithStatus 401 Text, WithStatus 403 Text]
       :<|> "login" :> ReqBody '[JSON] User :> Post '[JSON] (Headers '[ Header "Set-Cookie" SetCookie
                                                                      , Header "Set-Cookie" SetCookie ] NoContent)
       :<|> "logout" :> Get '[JSON] (Headers '[ Header "Set-Cookie" SetCookie
@@ -490,6 +507,11 @@ server ccfg =
                         :<|> raw
         Indefinite -> throwAll err401
         _ -> throwAll err403
+    ) :<|>
+    (\authResult -> case authResult of
+      Authenticated usr -> respond (WithStatus @200 (42 :: Int))
+      Indefinite -> respond (WithStatus @401 $ pack "Authentication required")
+      _ -> respond (WithStatus @403 $ pack "Forbidden")
     )
     :<|> getLogin
     :<|> getLogout
@@ -570,6 +592,12 @@ shouldMatchCookieNames cj patterns
     = fmap cookie_name (destroyCookieJar cj)
     `shouldMatchList` patterns
 
+shouldNotHaveCookies :: HCli.CookieJar -> [BS.ByteString] -> Expectation
+shouldNotHaveCookies cj patterns =
+    sequence_ $ (\cookieName -> cookieNames `shouldNotContain` [cookieName]) <$> patterns
+  where cookieNames :: [BS.ByteString]
+        cookieNames = cookie_name <$> destroyCookieJar cj
+
 shouldMatchCookieNameValues :: HCli.CookieJar -> [(BS.ByteString, BS.ByteString)] -> Expectation
 shouldMatchCookieNameValues cj patterns
     = fmap ((,) <$> cookie_name <*> cookie_value) (destroyCookieJar cj)

servant/src/Servant/API/ResponseHeaders.hs

@@ -8,6 +8,7 @@
 {-# LANGUAGE PolyKinds              #-}
 {-# LANGUAGE ScopedTypeVariables    #-}
 {-# LANGUAGE TypeFamilies           #-}
+{-# LANGUAGE TypeFamilyDependencies #-}
 {-# LANGUAGE TypeOperators          #-}
 {-# LANGUAGE UndecidableInstances   #-}
 {-# OPTIONS_HADDOCK not-home        #-}
@@ -51,6 +52,9 @@ import           Prelude ()
 import           Prelude.Compat
 import           Servant.API.Header
                  (Header)
+import           Servant.API.UVerb.Union
+import qualified Data.SOP.BasicFunctors as SOP
+import qualified Data.SOP.NS as SOP
 
 -- | Response Header objects. You should never need to construct one directly.
 -- Instead, use 'addOptionalHeader'.
@@ -170,6 +174,25 @@ instance {-# OVERLAPPABLE #-} ( KnownSymbol h, ToHttpApiData v , new ~ Headers '
          => AddHeader h v a new where
     addOptionalHeader hdr resp = Headers resp (HCons hdr HNil)
 
+-- Instances to decorate all responses in a 'Union' with headers. The functional
+-- dependencies force us to consider singleton lists as the base case in the
+-- recursion (it is impossible to determine h and v otherwise from old / new
+-- responses if the list is empty).
+instance (AddHeader h v old new) => AddHeader h v (Union '[old]) (Union '[new]) where
+  addOptionalHeader hdr resp =
+    SOP.Z $ SOP.I $ addOptionalHeader hdr $ SOP.unI $ SOP.unZ $ resp
+
+instance
+  ( AddHeader h v old new, AddHeader h v (Union oldrest) (Union newrest)
+  -- This ensures that the remainder of the response list is _not_ empty
+  -- It is necessary to prevent the two instances for union types from
+  -- overlapping.
+  , oldrest ~ (a ': as), newrest ~ (b ': bs))
+  => AddHeader h v (Union (old ': (a ': as))) (Union (new ': (b ': bs))) where
+  addOptionalHeader hdr resp = case resp of
+    SOP.Z (SOP.I rHead) -> SOP.Z $ SOP.I $ addOptionalHeader hdr rHead
+    SOP.S rOthers -> SOP.S $ addOptionalHeader hdr rOthers
+
 -- | @addHeader@ adds a header to a response. Note that it changes the type of
 -- the value in the following ways:
 --

servant/src/Servant/API/UVerb.hs

@@ -38,6 +38,7 @@ import GHC.TypeLits (Nat)
 import Network.HTTP.Types (Status, StdMethod)
 import Servant.API.ContentTypes (JSON, PlainText, FormUrlEncoded, OctetStream, NoContent, MimeRender(mimeRender), MimeUnrender(mimeUnrender))
 import Servant.API.Status (KnownStatus, statusVal)
+import Servant.API.ResponseHeaders (Headers)
 import Servant.API.UVerb.Union
 
 class KnownStatus (StatusOf a) => HasStatus (a :: *) where
@@ -86,6 +87,8 @@ newtype WithStatus (k :: Nat) a = WithStatus a
 instance KnownStatus n => HasStatus (WithStatus n a) where
   type StatusOf (WithStatus n a) = n
 
+instance HasStatus a => HasStatus (Headers ls a) where
+  type StatusOf (Headers ls a) = StatusOf a
 
 -- | A variant of 'Verb' that can have any of a number of response values and status codes.
 --