feat: implement JWT caching (#2928)

This commit is contained in:
Taimoor Zaeem
2023-09-25 14:46:55 -03:00
committed by GitHub
parent 90e3a5e29f
commit a6e3eda5b2
19 changed files with 206 additions and 19 deletions
+16
View File
@@ -4,6 +4,7 @@
module PostgREST.AppState
( AppState
, AuthResult(..)
, destroy
, getConfig
, getSchemaCache
@@ -12,6 +13,7 @@ module PostgREST.AppState
, getPgVersion
, getRetryNextIn
, getTime
, getJwtCache
, init
, initWithPool
, logWithZTime
@@ -24,8 +26,11 @@ module PostgREST.AppState
, runListener
) where
import qualified Data.Aeson as JSON
import qualified Data.Aeson.KeyMap as KM
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Cache as C
import Data.Either.Combinators (whenLeft)
import qualified Data.Text.Encoding as T
import Hasql.Connection (acquire)
@@ -62,6 +67,11 @@ import PostgREST.SchemaCache.Identifiers (dumpQi)
import Protolude
data AuthResult = AuthResult
{ authClaims :: KM.KeyMap JSON.Value
, authRole :: BS.ByteString
}
data AppState = AppState
-- | Database connection pool
{ statePool :: SQL.Pool
@@ -87,6 +97,8 @@ data AppState = AppState
, stateRetryNextIn :: IORef Int
-- | Logs a pool error with a debounce
, debounceLogAcquisitionTimeout :: IO ()
-- | JWT Cache
, jwtCache :: C.Cache ByteString AuthResult
}
init :: AppConfig -> IO AppState
@@ -108,6 +120,7 @@ initWithPool pool conf = do
<*> myThreadId
<*> newIORef 0
<*> pure (pure ())
<*> C.newCache Nothing
debLogTimeout <-
@@ -188,6 +201,9 @@ putConfig = atomicWriteIORef . stateConf
getTime :: AppState -> IO UTCTime
getTime = stateGetTime
getJwtCache :: AppState -> C.Cache ByteString AuthResult
getJwtCache = jwtCache
-- | Log to stderr with local time
logWithZTime :: AppState -> Text -> IO ()
logWithZTime appState txt = do
+48 -16
View File
@@ -26,6 +26,8 @@ import qualified Data.Aeson.KeyMap as KM
import qualified Data.Aeson.Types as JSON
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy.Char8 as LBS
import qualified Data.Cache as C
import qualified Data.Scientific as Sci
import qualified Data.Vault.Lazy as Vault
import qualified Data.Vector as V
import qualified Network.HTTP.Types.Header as HTTP
@@ -36,22 +38,20 @@ import Control.Lens (set)
import Control.Monad.Except (liftEither)
import Data.Either.Combinators (mapLeft)
import Data.List (lookup)
import Data.Time.Clock (UTCTime)
import Data.Time.Clock (UTCTime, nominalDiffTimeToSeconds)
import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
import System.Clock (TimeSpec (..))
import System.IO.Unsafe (unsafePerformIO)
import System.TimeIt (timeItT)
import PostgREST.AppState (AppState, getConfig, getTime)
import PostgREST.AppState (AppState, AuthResult (..), getConfig,
getJwtCache, getTime)
import PostgREST.Config (AppConfig (..), JSPath, JSPathExp (..))
import PostgREST.Error (Error (..))
import Protolude
data AuthResult = AuthResult
{ authClaims :: KM.KeyMap JSON.Value
, authRole :: BS.ByteString
}
-- | Receives the JWT secret and audience (from config) and a JWT and returns a
-- JSON object of JWT claims.
parseToken :: Monad m =>
@@ -107,16 +107,48 @@ middleware appState app req respond = do
let token = fromMaybe "" $ Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
parseJwt = runExceptT $ parseToken conf (LBS.fromStrict token) time >>= parseClaims conf
if configDbPlanEnabled conf
then do
(dur,authResult) <- timeItT parseJwt
let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
app req' respond
else do
authResult <- parseJwt
let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
app req' respond
-- If DbPlanEnabled -> calculate JWT validation time
-- If JwtCacheMaxLifetime -> cache JWT validation result
req' <- case (configDbPlanEnabled conf, configJwtCacheMaxLifetime conf) of
(True, 0) -> do
(dur, authResult) <- timeItT parseJwt
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
(True, maxLifetime) -> do
(dur, authResult) <- timeItT $ getJWTFromCache appState token maxLifetime parseJwt time
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
(False, 0) -> do
authResult <- parseJwt
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
(False, maxLifetime) -> do
authResult <- getJWTFromCache appState token maxLifetime parseJwt time
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
app req' respond
-- | Used to retrieve and insert JWT to JWT Cache
getJWTFromCache :: AppState -> ByteString -> Int -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult)
getJWTFromCache appState token maxLifetime parseJwt utc = do
checkCache <- C.lookup (getJwtCache appState) token
authResult <- maybe parseJwt (pure . Right) checkCache
case (authResult,checkCache) of
(Right res, Nothing) -> C.insert' (getJwtCache appState) (getTimeSpec res maxLifetime utc) token res
_ -> pure ()
return authResult
-- Used to extract JWT exp claim and add to JWT Cache
getTimeSpec :: AuthResult -> Int -> UTCTime -> Maybe TimeSpec
getTimeSpec res maxLifetime utc = do
let expireJSON = KM.lookup "exp" (authClaims res)
utcToSecs = floor . nominalDiffTimeToSeconds . utcTimeToPOSIXSeconds
sciToInt = fromMaybe 0 . Sci.toBoundedInteger
case expireJSON of
Just (JSON.Number seconds) -> Just $ TimeSpec (sciToInt seconds - utcToSecs utc) 0
_ -> Just $ TimeSpec (fromIntegral maxLifetime :: Int64) 0
authResultKey :: Vault.Key (Either Error AuthResult)
authResultKey = unsafePerformIO Vault.newKey
+4 -1
View File
@@ -162,7 +162,7 @@ exampleConfigFile =
|## Time in seconds after which to recycle unused pool connections
|# db-pool-max-idletime = 30
|
|## Allow autmatic database connection retrying
|## Allow automatic database connection retrying
|# db-pool-automatic-recovery = true
|
|## Stored proc to exec immediately after auth
@@ -205,6 +205,9 @@ exampleConfigFile =
|# jwt-secret = "secret_with_at_least_32_characters"
|jwt-secret-is-base64 = false
|
|## Enables and set JWT Cache max lifetime, disables caching with 0
|# jwt-cache-max-lifetime = 0
|
|## Logging level, the admitted values are: crit, error, warn and info.
|log-level = "error"
|
+3
View File
@@ -97,6 +97,7 @@ data AppConfig = AppConfig
, configJwtRoleClaimKey :: JSPath
, configJwtSecret :: Maybe BS.ByteString
, configJwtSecretIsBase64 :: Bool
, configJwtCacheMaxLifetime :: Int
, configLogLevel :: LogLevel
, configOpenApiMode :: OpenAPIMode
, configOpenApiSecurityActive :: Bool
@@ -162,6 +163,7 @@ toText conf =
,("jwt-role-claim-key", q . T.intercalate mempty . fmap dumpJSPath . configJwtRoleClaimKey)
,("jwt-secret", q . T.decodeUtf8 . showJwtSecret)
,("jwt-secret-is-base64", T.toLower . show . configJwtSecretIsBase64)
,("jwt-cache-max-lifetime", show . configJwtCacheMaxLifetime)
,("log-level", q . dumpLogLevel . configLogLevel)
,("openapi-mode", q . dumpOpenApiMode . configOpenApiMode)
,("openapi-security-active", T.toLower . show . configOpenApiSecurityActive)
@@ -265,6 +267,7 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
<*> (fromMaybe False <$> optWithAlias
(optBool "jwt-secret-is-base64")
(optBool "secret-is-base64"))
<*> (fromMaybe 0 <$> optInt "jwt-cache-max-lifetime")
<*> parseLogLevel "log-level"
<*> parseOpenAPIMode "openapi-mode"
<*> (fromMaybe False <$> optBool "openapi-security-active")