refactor: simplify control flow in App.postgrest

Currently, authentication and response execution each unwrap ExceptT with separate runExceptT calls, which split the main request flow across nested pattern matching and Either handling. Control flow is complex and difficult to understand.

The goal of this change is to make request execution as sequential
monadic code with clear error handling.

To implement that, request handling is now run in ExceptT over WriterT (Last ByteString) IO monad stack. Auth role is written after authentication succeeds and further returned along the response. Thanks to it response observation generation is centralized at the end of request handling.

It was necessary to abstract monad stack in getAuthResult, lookupJwtCache, postgrestResponse, and withTiming to enable introduction of WriterT.

Author: mkleczek

src/PostgREST/App.hs

@@ -9,6 +9,7 @@ Some of its functionality includes:
 - Producing HTTP Headers according to RFCs.
 - Content Negotiation
 -}
+{-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE NamedFieldPuns      #-}
 {-# LANGUAGE RecordWildCards     #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -61,6 +62,7 @@ import PostgREST.SchemaCache          (SchemaCache (..))
 import PostgREST.TimeIt               (timeItT)
 import PostgREST.Version              (docsVersion, prettyVersion)
 
+import           Control.Monad.Writer
 import qualified Data.ByteString.Char8     as BS
 import qualified Data.List                 as L
 import           Data.Streaming.Network    (bindPortTCP)
@@ -127,24 +129,21 @@ postgrest appState connWorker =
       appConf@AppConfig{..} <- AppState.getConfig appState -- the config must be read again because it can reload
       maybeSchemaCache <- AppState.getSchemaCache appState
 
-      let observer = AppState.getObserver appState
-          bearerAuth = ApiRequest.userBearerAuth req
+      (response, user) <-
+        -- has to be before runExceptT to make sure role is not lost on error
+        runWriterT $
+        -- handle errors
+        fmap (either (Error.errorResponseFor configClientErrorVerbosity) identity) $
+        runExceptT $
+        do
+          (jwtTime, authResult@AuthResult{..}) <- withTiming appConf $
+            Auth.getAuthResult appState $ ApiRequest.userBearerAuth req
 
-      response <- do
-        authResultE <- runExceptT $ withTiming appConf $
-            liftIO (Auth.getAuthResult appState bearerAuth) >>= liftEither
+          tell $ pure authRole
 
-        case authResultE of
-          Left err -> do
-            let resp = Error.errorResponseFor configClientErrorVerbosity err
-            observer $ genResponseObs Nothing req resp
-            pure resp
+          postgrestResponse appState appConf maybeSchemaCache jwtTime authResult req
 
-          Right (jwtTime, authResult@AuthResult{..}) -> do
-            resp <- either (Error.errorResponseFor configClientErrorVerbosity) identity <$>
-              runExceptT (postgrestResponse appState appConf maybeSchemaCache jwtTime authResult req)
-            observer $ genResponseObs (Just authRole) req resp
-            pure resp
+      AppState.getObserver appState $ genResponseObs (getLast user) req response
 
       -- Launch the connWorker when the connection is down. The postgrest
       -- function can respond successfully (with a stale schema cache) before
@@ -164,13 +163,14 @@ postgrest appState connWorker =
       ResponseObs user req (Wai.responseStatus resp) (WaiHeader.contentLength $ Wai.responseHeaders resp)
 
 postgrestResponse
-  :: AppState.AppState
+  :: (MonadError Error m, MonadIO m)
+  => AppState.AppState
   -> AppConfig
   -> Maybe SchemaCache
   -> Maybe Double
   -> AuthResult
   -> Wai.Request
-  -> ExceptT Error IO Wai.Response
+  -> m Wai.Response
 postgrestResponse appState conf@AppConfig{..} maybeSchemaCache jwtTime authResult@AuthResult{..} req = do
   let observer = AppState.getObserver appState
 
@@ -179,12 +179,12 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache jwtTime authResul
       Just sCache ->
         return sCache
       Nothing -> do
-        lift $ observer SchemaCacheEmptyObs
+        liftIO $ observer SchemaCacheEmptyObs
         throwError Error.NoSchemaCacheError
 
   let prefs = ApiRequest.userPreferences conf req (dbTimezones sCache)
 
-  body <- lift $ Wai.strictRequestBody req
+  body <- liftIO $ Wai.strictRequestBody req
 
   (parseTime, apiReq@ApiRequest{..}) <- withTiming conf $ liftEither . mapLeft Error.ApiRequestErr $ ApiRequest.userApiRequest conf prefs req body
   (planTime, plan)                   <- withTiming conf $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
@@ -197,21 +197,21 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache jwtTime authResul
     case tx of
       MainTx.NoDbTx r -> pure r
       MainTx.DbTx{..} -> do
-        dbRes <- lift $ AppState.usePool appState (dqTransaction dqIsoLevel dqTxMode $ runExceptT dqDbHandler)
+        dbRes <- liftIO $ AppState.usePool appState (dqTransaction dqIsoLevel dqTxMode $ runExceptT dqDbHandler)
         let eitherResp = join $ mapLeft (Error.PgErr . Error.PgError (Just authRole /= configDbAnonRole)) dbRes
 
         -- TODO: we use obsQuery twice, one here and one below because in case of an error with the usePool above, the request will finish here and return an error message.
         -- This is because of a combination of ExceptT + our Error module which has Wai.responseLBS.
         -- This needs refactoring so only the below obsQuery is used.
-        lift $ whenLeft eitherResp $ obsQuery . Error.status
+        liftIO $ whenLeft eitherResp $ obsQuery . Error.status
         liftEither eitherResp
 
   (respTime, resp) <- withTiming conf $ do
     let response = Response.actionResponse txResult apiReq (T.decodeUtf8 prettyVersion, docsVersion) conf sCache
         status' = either Error.status Response.pgrstStatus response
 
     -- TODO: see above obsQuery, only this obsQuery should remain after refactoring (because the QueryObs depends on the status)
-    lift $ obsQuery status'
+    liftIO $ obsQuery status'
     liftEither response
 
   return $ toWaiResponse (ServerTiming jwtTime parseTime planTime txTime respTime) resp
@@ -230,7 +230,7 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache jwtTime authResul
     varyHeaderPresent :: [HTTP.Header] -> Bool
     varyHeaderPresent = any (\(h, _v) -> h == HTTP.hVary)
 
-withTiming :: AppConfig -> ExceptT e IO a -> ExceptT e IO (Maybe Double, a)
+withTiming :: (MonadError e m, MonadIO m) => AppConfig -> m a -> m (Maybe Double, a)
 withTiming AppConfig{configServerTimingEnabled} f = if configServerTimingEnabled
     then do
       (t, r) <- timeItT f

src/PostgREST/Auth.hs

@@ -10,6 +10,7 @@ Authentication should always be implemented in an external service.
 In the test suite there is an example of simple login function that can be used for a
 very simple authentication system inside the PostgreSQL database.
 -}
+{-# LANGUAGE FlexibleContexts #-}
 module PostgREST.Auth
   ( getAuthResult )
   where
@@ -25,14 +26,9 @@ import Protolude
 
 -- | Perform authentication and authorization
 --   Parse JWT and return AuthResult
-getAuthResult :: AppState -> Maybe ByteString -> IO (Either Error AuthResult)
+getAuthResult :: (MonadError Error m, MonadIO m) => AppState -> Maybe ByteString -> m AuthResult
 getAuthResult appState token = do
-  conf <- getConfig appState
-  time <- getTime appState
+  conf <- liftIO $ getConfig appState
+  time <- liftIO $ getTime appState
 
-  let jwtCacheState = getJwtCacheState appState
-      parseJwt = runExceptT $ do
-        claims <- lookupJwtCache jwtCacheState token
-        parseClaims conf time claims
-
-  parseJwt
+  parseClaims conf time =<< lookupJwtCache (getJwtCacheState appState) token

src/PostgREST/Auth/JwtCache.hs

@@ -5,6 +5,7 @@ Description : PostgREST JWT validation results Cache.
 This module provides functions to deal with the JWT cache.
 -}
 {-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE FlexibleContexts          #-}
 {-# LANGUAGE FlexibleInstances         #-}
 {-# LANGUAGE LambdaCase                #-}
 {-# LANGUAGE MultiParamTypeClasses     #-}
@@ -42,7 +43,7 @@ import           Protolude
 data JwtCacheState = JwtCacheState ObservationHandler (IORef JwtCache)
 
 class CacheVariant m v where
-  cached :: SC.Cache m ByteString v -> ByteString -> ExceptT Error IO JSON.Object
+  cached :: (MonadError Error n, MonadIO n) => SC.Cache m ByteString v -> ByteString -> n JSON.Object
 
 {-|
 Jwt caching can have three different configurations:
@@ -60,12 +61,12 @@ data JwtCache =
   forall m v. CacheVariant m v => JwtCache JwkSet (TVar Int) (SC.Cache m ByteString v)
 
 instance CacheVariant IO (Either Error JSON.Object) where
-  cached c = lift . SC.cached c >=> liftEither
+  cached c = liftIO . SC.cached c >=> liftEither
 
 instance CacheVariant (ExceptT Error IO) JSON.Object where
-  cached = SC.cached
+  cached c = liftIO . runExceptT . SC.cached c >=> liftEither
 
-decode :: JwtCache -> ByteString -> ExceptT Error IO JSON.Object
+decode :: (MonadError Error m, MonadIO m) => JwtCache -> ByteString -> m JSON.Object
 decode JwtNoJwks        = const $ throwError (JwtErr JwtSecretMissing)
 decode (JwtNoCache key) = parseAndDecodeClaims key
 decode (JwtCache _ _ c) = cached c
@@ -110,5 +111,5 @@ newJwtCache AppConfig{configJWKS, configJwtCacheMaxEntries} observationHandler =
             (const . const $ lift $ observationHandler JwtCacheEviction) -- evictions metrics
             alwaysValid) -- no invalidation for now
 
-lookupJwtCache :: JwtCacheState -> Maybe ByteString -> ExceptT Error IO JSON.Object
+lookupJwtCache :: (MonadError Error m, MonadIO m) => JwtCacheState -> Maybe ByteString -> m JSON.Object
 lookupJwtCache (JwtCacheState _ cacheState) k = liftIO (readIORef cacheState) >>= flip (maybe (pure KM.empty)) k . decode