perf: JWT Cache based on Sieve algorithm

Changes:
1. Refactoring and some cleanup of JWT handling code:
* Instead of caching AuthResult cache decoded claims (which signature was verified). Validating claims and determining role is done after cache lookup
* Cleaned up API so that usage of it is simplified: lookupJwtCache cache key >>= parseClaims configJwtAud time
* Handling of JwtCacheState initialization and updates of configuration is encapsulated in Auth.JwtCache module

2. Generic high performance (hopefully) scalable, dynamically resizeable cache implementation based on stm, stm-hamt and sieve algorithm. It also integrates with PostgREST measurements infrastructure providing usage stats (ie. hit ratio, evictions count)

Author: mkleczek

CHANGELOG.md

@@ -11,6 +11,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
   + It now shows the invalid schema in the `message` field.
   + The exposed schemas are now listed in the `hint` instead of the `message` field.
 - Improve error details of `PGRST301` error by @taimoorzaeem in #4051
+- JWT cache based on sieve algorithm
 
 ## [13.0.4] - 2025-06-17
 

docs/references/auth.rst

@@ -98,9 +98,11 @@ The ``Bearer`` header value can be used with or without capitalization(``bearer`
 JWT Caching
 -----------
 
-PostgREST validates ``JWTs`` on every request. We can cache ``JWTs`` to avoid this performance overhead.
+PostgREST validates ``JWTs`` on every request. Asymmetric signature validation (such as RSA) is slow and we can cache ``JWT`` validation results to avoid this performance overhead.
 
-To enable JWT caching, the config :code:`jwt-cache-max-lifetime` is to be set. It is the maximum number of seconds for which the cache stores the JWT validation results. The cache uses the :code:`exp` claim to set the cache entry lifetime. If the JWT does not have an :code:`exp` claim, it uses the config value. See :ref:`jwt-cache-max-lifetime` for more details.
+JWT caching is automatically enabled if ref:`jwt-secret` is set to an asymmetric key. Otherwise it is disabled and can be enabled by setting the config :code:`jwt-cache-max-entries` to a value greater than 0. Setting it to 0 disables caching regardless of ref:`jwt-secret`.
+
+See :ref:`jwt-cache-max-entries` for more details.
 
 .. note::
 

docs/references/configuration.rst

@@ -658,20 +658,20 @@ jwt-secret-is-base64
 
   When this is set to :code:`true`, the value derived from :code:`jwt-secret` will be treated as a base64 encoded secret.
 
-.. _jwt-cache-max-lifetime:
+.. _jwt-cache-max-entries:
 
-jwt-cache-max-lifetime
+jwt-cache-max-entries
 ----------------------
 
   =============== =================================
   **Type**        Int
-  **Default**     0
+  **Default**     1000
   **Reloadable**  Y
-  **Environment** PGRST_JWT_CACHE_MAX_LIFETIME
-  **In-Database** pgrst.jwt_cache_max_lifetime
+  **Environment** PGRST_JWT_CACHE_MAX_ENTRIES
+  **In-Database** pgrst.jwt_cache_max_entries
   =============== =================================
 
-  Maximum number of seconds of lifetime for cached entries. The default :code:`0` disables caching. See :ref:`jwt_caching`.
+  Maximum number of entries in JWT cache. The value :code:`0` disables JWT caching. See :ref:`jwt_caching`.
 
 .. _log-level:
 

nix/tools/loadtest.nix

@@ -61,7 +61,6 @@ let
         export PGRST_DB_TX_END="rollback-allow-override"
         export PGRST_LOG_LEVEL="crit"
         export PGRST_JWT_SECRET="reallyreallyreallyreallyverysafe"
-        export PGRST_JWT_CACHE_MAX_LIFETIME="86400"
 
         mkdir -p "$(dirname "$_arg_output")"
         abs_output="$(realpath "$_arg_output")"
@@ -72,7 +71,7 @@ let
             ${genTargetsHS} "$_arg_testdir"/gen_targets.http
 
             if [ "$_arg_jwtcache" = "off" ]; then
-              export PGRST_JWT_CACHE_MAX_LIFETIME="0"
+              export PGRST_JWT_CACHE_MAX_ENTRIES="0"
             fi
 
             # shellcheck disable=SC2145
@@ -89,7 +88,7 @@ let
             export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json"
 
             if [ "$_arg_jwtcache" = "off" ]; then
-              export PGRST_JWT_CACHE_MAX_LIFETIME="0"
+              export PGRST_JWT_CACHE_MAX_ENTRIES="0"
             fi
 
             # shellcheck disable=SC2145

postgrest.cabal

@@ -52,6 +52,7 @@ library
                       PostgREST.Auth.Jwt
                       PostgREST.Auth.JwtCache
                       PostgREST.Auth.Types
+                      PostgREST.Cache.Sieve
                       PostgREST.CLI
                       PostgREST.Config
                       PostgREST.Config.Database
@@ -154,6 +155,10 @@ library
                     -- https://github.com/kazu-yamamoto/logger/commit/3a71ca70afdbb93d4ecf0083eeba1fbbbcab3fc3
                     , wai-logger                >= 2.4.0
                     , warp                      >= 3.3.19 && < 3.4
+                    , stm                       >= 2.5 && < 3
+                    , stm-hamt                  >= 1.2 && < 2
+                    , focus                     >= 1.0 && < 2
+                    , some                      >= 1.0.4.1 && < 2
                       -- -fno-spec-constr may help keep compile time memory use in check,
                       --   see https://gitlab.haskell.org/ghc/ghc/issues/16017#note_219304
                       -- -optP-Wno-nonportable-include-path
@@ -208,6 +213,7 @@ test-suite spec
                       Feature.Auth.AudienceJwtSecretSpec
                       Feature.Auth.AuthSpec
                       Feature.Auth.BinaryJwtSecretSpec
+                      Feature.Auth.JwtCacheSpec
                       Feature.Auth.NoAnonSpec
                       Feature.Auth.NoJwtSecretSpec
                       Feature.ConcurrentSpec
@@ -265,6 +271,7 @@ test-suite spec
                     , hasql-transaction >= 1.0.1 && < 1.1
                     , heredoc           >= 0.2 && < 0.3
                     , hspec             >= 2.3 && < 2.12
+                    , hspec-expectations >= 0.8.4 && < 0.9
                     , hspec-wai         >= 0.10 && < 0.12
                     , hspec-wai-json    >= 0.10 && < 0.12
                     , http-types        >= 0.12.3 && < 0.13
@@ -274,6 +281,7 @@ test-suite spec
                     , monad-control     >= 1.0.1 && < 1.1
                     , postgrest
                     , process           >= 1.4.2 && < 1.7
+                    , prometheus-client >= 1.1.1 && < 1.2.0
                     , protolude         >= 0.3.1 && < 0.4
                     , regex-tdfa        >= 1.2.2 && < 1.4
                     , scientific        >= 0.3.4 && < 0.4

src/PostgREST/AppState.hs

@@ -57,7 +57,7 @@ import Data.IORef         (IORef, atomicWriteIORef, newIORef,
                            readIORef)
 import Data.Time.Clock    (UTCTime, getCurrentTime)
 
-import PostgREST.Auth.JwtCache           (JwtCacheState)
+import PostgREST.Auth.JwtCache           (JwtCacheState, update)
 import PostgREST.Config                  (AppConfig (..),
                                           addFallbackAppName,
                                           readAppConfig)
@@ -127,14 +127,13 @@ init conf@AppConfig{configLogLevel, configDbPoolSize} = do
 
   observer $ AppStartObs prettyVersion
 
-  jwtCacheState <- JwtCache.init
   pool <- initPool conf observer
   (sock, adminSock) <- initSockets conf
-  state' <- initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState observer
+  state' <- initWithPool (sock, adminSock) pool conf loggerState metricsState observer
   pure state' { stateSocketREST = sock, stateSocketAdmin = adminSock}
 
-initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> JwtCache.JwtCacheState -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
-initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState observer = do
+initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
+initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
 
   appState <- AppState pool
     <$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
@@ -150,7 +149,7 @@ initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState
     <*> pure sock
     <*> pure adminSock
     <*> pure observer
-    <*> pure jwtCacheState
+    <*> JwtCache.init conf observer
     <*> pure loggerState
     <*> pure metricsState
 
@@ -471,10 +470,7 @@ readInDbConfig startingUp appState@AppState{stateObserver=observer} = do
       -- After the config has reloaded, jwt-secret might have changed, so
       -- if it has changed, it is important to invalidate the jwt cache
       -- entries, because they were cached using the old secret
-      if configJwtSecret conf == configJwtSecret newConf then
-        pass
-      else
-        JwtCache.emptyCache (getJwtCacheState appState) -- atomic O(1) operation
+      update (getJwtCacheState appState) newConf
 
       if startingUp then
         pass

src/PostgREST/Auth.hs

@@ -30,50 +30,32 @@ import System.TimeIt    (timeItT)
 
 import PostgREST.AppState      (AppState, getConfig, getJwtCacheState,
                                 getTime)
+import PostgREST.Auth.Jwt      (parseClaims)
 import PostgREST.Auth.JwtCache (lookupJwtCache)
 import PostgREST.Auth.Types    (AuthResult (..))
 import PostgREST.Config        (AppConfig (..))
-import PostgREST.Error         (Error (..), JwtError (..))
+import PostgREST.Error         (Error (..))
 
-import qualified Data.Aeson.KeyMap  as KM
-import           PostgREST.Auth.Jwt (parseAndDecodeClaims,
-                                     parseClaims)
-import           Protolude
+import Protolude
 
--- | Validate authorization header.
+-- | Validate authorization header
 --   Parse and store JWT claims for future use in the request.
 middleware :: AppState -> Wai.Middleware
 middleware appState app req respond = do
-  cfg@AppConfig{..} <- getConfig appState
+  conf@AppConfig{..} <- getConfig appState
   time <- getTime appState
 
   let token  = Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
-      parseAuthToken = maybe (const $ throwError (JwtErr JwtSecretMissing)) parseAndDecodeClaims configJWKS
-      parseJwt = runExceptT $ maybe (pure KM.empty) parseAuthToken token >>= parseClaims cfg time
+      parseJwt = runExceptT $ lookupJwtCache jwtCacheState token >>= parseClaims conf time
       jwtCacheState = getJwtCacheState appState
 
--- If ServerTimingEnabled -> calculate JWT validation time
--- If JwtCacheMaxLifetime -> cache JWT validation result
-  req' <- case (configServerTimingEnabled, configJwtCacheMaxLifetime) 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 $ case token of
-            Just tkn -> lookupJwtCache jwtCacheState tkn maxLifetime parseJwt time
-            Nothing  -> parseJwt
-          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 <- case token of
-            Just tkn -> lookupJwtCache jwtCacheState tkn maxLifetime parseJwt time
-            Nothing -> parseJwt
-          return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
+  -- If ServerTimingEnabled -> calculate JWT validation time
+  req' <- if configServerTimingEnabled then do
+      (dur, authResult) <- timeItT parseJwt
+      pure $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
+    else do
+      authResult <- parseJwt
+      pure $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
 
   app req' respond