Auth Servant

App.hs

{-# LANGUAGE TypeOperators, DataKinds, OverloadedStrings, DeriveGeneric, FlexibleContexts #-}
module Lib
    ( app
    , User(..)
    ) where

import Servant
import Servant.API
import Data.Proxy
import Network.Wai
import Servant.Auth.Server
import Data.Aeson (ToJSON, FromJSON)
import GHC.Generics (Generic)
import Control.Monad.IO.Class (liftIO, MonadIO)
import API(test)

type ReaderAPI = "ep1" :> Get '[JSON] Int 
            :<|> "ep2" :> Get '[JSON] String

type Unprotected =
    "login"
        :> ReqBody '[JSON] Login
        :> PostNoContent '[JSON] (Headers '[ Header "Set-Cookie" SetCookie
                                           , Header "Set-Cookie" SetCookie]
                                  NoContent)
        :<|> Raw

type API auths = (Auth auths User :> ReaderAPI) :<|> Unprotected

api = Proxy :: Proxy (API '[JWT, Cookie])

unprotected :: CookieSettings -> JWTSettings -> Server Unprotected
unprotected cs jwts = checkCreds cs jwts :<|> serveDirectoryFileServer "example/static"

checkCreds :: CookieSettings
           -> JWTSettings
           -> Login
           -> Handler (Headers '[ Header "Set-Cookie" SetCookie
                                , Header "Set-Cookie" SetCookie]
                               NoContent)
checkCreds cookieSettings jwtSettings login = do
  user <- liftIO $ authenticate login
  case user of
    Nothing -> throwError err401
    Just usr -> do
      mApplyCookies <- liftIO $ acceptLogin cookieSettings jwtSettings usr
      case mApplyCookies of
        Nothing           -> throwError err401
        Just applyCookies -> return $ applyCookies NoContent

authenticate :: MonadIO m => Login -> m (Maybe User)
authenticate (Login "truong" "dung") = pure $ Just (User "alice" "alice@gmail.com")
authenticate _ = pure Nothing

data Login = Login { username :: String, password :: String }
  deriving (Eq, Show, Read, Generic)

data User = User { name :: String, email :: String }
  deriving (Eq, Show, Read, Generic)

instance ToJSON User
instance ToJWT User
instance FromJSON User
instance FromJWT User

instance ToJSON Login
instance FromJSON Login

checkJWT :: MonadIO m => User -> m (Maybe Bool)
checkJWT (User "alice" "alice@gmail.com") = return $ Just True
checkJWT _ = return Nothing

-- readerServer :: AuthResult User -> Server ReaderAPI
readerServer xs (Authenticated user) =
  withLog user (return 1797) :<|> withLog user (return (name user <> xs)) 
  where
    withLog user action = do
      usr <- liftIO $ checkJWT user
      -- liftIO $ print user
      case usr of
        Just True -> action
        _ -> throwError err401
    
readerServer xs a@(_) = throwAll err401          

-- server :: CookieSettings -> JWTSettings -> ReaderT String (Server (API auths))
server dbConf cs jwts = (readerServer dbConf) :<|> (unprotected cs jwts)

-- nt :: Config -> AppM a -> Handler a
--nt conf x = runReaderT x conf

-- app :: Config -> Application
app dbConf jwtCfg cfg = serveWithContext api cfg (server dbConf defaultCookieSettings jwtCfg)--hoistServer api (nt conf) server

Spec.hs

{-# LANGUAGE OverloadedStrings, FlexibleContexts, QuasiQuotes #-}
import Lib (app, User(..))
import Test.Hspec
import Test.Hspec.Wai
import Test.Hspec.Wai.JSON
import Network.HTTP.Types
import Servant.Server
import Servant.Auth.Server
import Control.Monad.IO.Class (liftIO)
import Crypto.JOSE.JWK
import Data.Monoid ((<>))
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.String as DS
import qualified Data.ByteString.Char8 as BC
import qualified Data.ByteString.Lazy as BL
import Test.Hspec.Wai.Matcher

main :: IO ()
main = do
  myKey <- generateKey
  hspec (spec "hi" myKey)

bearerType jwt = BC.pack "Bearer " <> BL.toStrict jwt

headers :: BL.ByteString -> [Header]
headers jwt = [
    (hAuthorization,  bearerType jwt)
  ]

jsonHeader = [
    (hContentType, "application/json")
  ]  

validUser = User "alice" "alice@gmail.com"
invalidUser = User "test" "fdsfd@fdsfd.com"

mkHeaders user jwtCfg = do
  jwtRes <-  makeJWT user jwtCfg Nothing
  case jwtRes of
    Left _ -> return $ headers ""
    Right jwt -> return $ headers jwt

spec :: String -> JWK  -> Spec
spec dbConf myKey = do
  let jwtCfg = defaultJWTSettings myKey
      cfg = defaultCookieSettings :. jwtCfg :. EmptyContext
  with (return (app dbConf jwtCfg cfg)) $
    describe "Protected API" $ do
      it "responds with 401" $ do
        header <- liftIO $ mkHeaders invalidUser jwtCfg
        (request methodGet  "/ep1" header "") `shouldRespondWith` 401
      it "responds with 200" $ do
        header <- liftIO $ mkHeaders validUser jwtCfg
        (request methodGet  "/ep1" header "") `shouldRespondWith` 200
      it "responds with 1797" $ do
        header <- liftIO $ mkHeaders validUser jwtCfg
        (request methodGet  "/ep1" header "") `shouldRespondWith` "1797"
      it   "responds with 200" $ do
        header <- liftIO $ mkHeaders validUser jwtCfg
        (request methodGet  "/ep2" header "") `shouldRespondWith` 200
      it "responds with usernamehi" $ do
        header <- liftIO $ mkHeaders validUser jwtCfg
        let expected = bodyEquals (DS.fromString . show $ (name validUser <> "hi"))
        request methodGet  "/ep2" header "" `shouldRespondWith` (ResponseMatcher 200 [] expected)
      it "succeed login" $ do
        request methodPost "/login" jsonHeader [json|{username: "truong", password: "dung"}|] `shouldRespondWith` 204
      it "fails login" $ do
        request methodPost "/login" jsonHeader [json|{username: "truong1", password: "dung"}|] `shouldRespondWith` 401
      it "fails login with malformed request" $ do
        request methodPost "/login" jsonHeader "" `shouldRespondWith` 400
        request methodPost "/login" jsonHeader [json|{username: "truong1"}|] `shouldRespondWith` 400