Fix authentication example

servant-examples/auth-combinator/auth-combinator.hs

@@ -14,52 +14,80 @@ import           Network.Wai
 import           Network.Wai.Handler.Warp
 import           Servant
 import           Servant.Server.Internal
+import Data.Aeson
+import Data.ByteString (ByteString)
+import Data.ByteString.Builder.Internal (byteStringCopy)
+import Data.Monoid ((<>))
+import Data.Text (Text)
+import GHC.Generics
+import Network.HTTP.Types
+import Network.Wai
+import Network.Wai.Handler.Warp
+import Servant
+import Servant.API.Authentication
+import Servant.Server.Internal
+import Servant.Server.Internal.Authentication (strictProtect, AuthHandlers(AuthHandlers))
 
--- Pretty much stolen/adapted from
--- https://github.com/haskell-servant/HaskellSGMeetup2015/blob/master/examples/authentication-combinator/AuthenticationCombinator.hs
+-- | An example of a custom authentication framework that checks a Cookie for a
+-- value.
 
-type DBLookup = ByteString -> IO Bool
+-- | Data we will use to test for authentication
+data CookieAuth = CookieAuth { cookie :: ByteString }
 
+-- | a 'User' datatype we get once the authentication data is tested.
+type User = ByteString
+
+-- | we will look up authentication data in the database and extract a User.
+type DBLookup = CookieAuth -> IO (Maybe User)
+
+-- | method that tests for authentication and extracts a User type.
 isGoodCookie :: DBLookup
-isGoodCookie = return . (== "good password")
+isGoodCookie (CookieAuth cookie) = if cookie == "good cookie" then return (Just "one user") else return Nothing
 
-data AuthProtected
+-- | Response handlers: what do we do when authentication doesn't work.
+cookieAuthHandlers :: AuthHandlers CookieAuth
+cookieAuthHandlers = AuthHandlers missingAuth notAuthenticated
+  where
+    missingAuth :: IO Response
+    missingAuth = return $ (responseBuilder status401 [] "Missing Cookie header.")
 
-instance HasServer rest => HasServer (AuthProtected :> rest) where
-  type ServerT (AuthProtected :> rest) m = ServerT rest m
+    notAuthenticated :: CookieAuth -> IO Response
+    notAuthenticated (CookieAuth cookie) = return $
+        responseBuilder status401 [] ("Invalid cookie: " <> byteStringCopy cookie)
 
-  route Proxy a = WithRequest $ \ request ->
-    route (Proxy :: Proxy rest) $ do
-      case lookup "Cookie" (requestHeaders request) of
-        Nothing -> return $ failWith $ HttpError status401 (Just "Missing auth header.")
-        Just v  -> do
-          authGranted <- isGoodCookie v
-          if authGranted
-            then a
-            else return $ failWith $ HttpError status403 (Just "Invalid cookie.")
-
-type PrivateAPI = Get '[JSON] [PrivateData]
-
-type PublicAPI = Get '[JSON] [PublicData]
-
-type API = "private" :> AuthProtected :> PrivateAPI
-      :<|> PublicAPI
+-- | 'AuthData' is a typeclass that provides a method to extract authentication
+-- data from a 'Reqest'
+instance AuthData CookieAuth where
+    authData req = fmap CookieAuth (lookup "Cookie" (requestHeaders req))
 
+-- | some data we will return from our API that is protected
 newtype PrivateData = PrivateData { ssshhh :: Text }
   deriving (Eq, Show, Generic)
 
 instance ToJSON PrivateData
 
+-- | Some data we will return from our API that is not protected
 newtype PublicData = PublicData { somedata :: Text }
   deriving (Eq, Show, Generic)
 
 instance ToJSON PublicData
 
+-- | Private API (will require authentication) as a type
+type PrivateAPI = Get '[JSON] [PrivateData]
+
+-- | Public API (non-authenticated) as a type
+type PublicAPI = Get '[JSON] [PublicData]
+
+-- | Our full API as a type with authentication
+type API = AuthProtect CookieAuth User 'Strict :> "private" :>  PrivateAPI
+      :<|> PublicAPI
+
 api :: Proxy API
 api = Proxy
 
 server :: Server API
-server = return prvdata :<|> return pubdata
+server = strictProtect isGoodCookie (const (return prvdata)) cookieAuthHandlers
+    :<|> return pubdata
 
   where prvdata = [PrivateData "this is a secret"]
         pubdata = [PublicData "this is a public piece of data"]