feat: JWT cache implementation based on sieve algorithm (#4084)

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)

CHANGELOG.md

@@ -12,6 +12,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
   + The exposed schemas are now listed in the `hint` instead of the `message` field.
 - Improve error details of `PGRST301` error by @taimoorzaeem in #4051
 
+### Changed
+- #4084, Implemented fixed size JWT cache based on sieve algorithm
 ### Fixed
 
 - Fix OpenAPI broken docs link by @taimoorzaeem in #4080

docs/references/auth.rst

@@ -94,10 +94,12 @@ JWT Generation
 
 You can create a valid JWT either from inside your database (see :ref:`sql_user_management`) or via an external service (see :ref:`external_auth`).
 
-JWT Keys
---------
+.. _jwt_signature:
 
-PostgREST supports both symmetric and asymmetric keys for signing and verifying the token.
+JWT Signature Verification
+--------------------------
+
+PostgREST supports both symmetric and asymmetric keys for verifying the signature of the token.
 
 Symmetric Keys
 ~~~~~~~~~~~~~~
@@ -153,28 +155,25 @@ You can specify the literal value as we saw earlier, or reference a filename to
 
   jwt-secret = "@rsa.jwk.pub"
 
-.. _jwt_claims_validation:
+``kid`` verification
+^^^^^^^^^^^^^^^^^^^^
 
-JWT Claims Validation
----------------------
+PostgREST has built-in verification of the `key ID parameter <https://www.rfc-editor.org/rfc/rfc7517#section-4.5>`_, useful when working with a JWK Set.
+It goes as follows:
 
-JWT ``exp``, ``iat`` , ``nbf`` Validation
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+- If the JWT contains a ``kid`` parameter, then PostgREST will look for the JWK in the :ref:`jwt-secret`.
 
-The time-based JWT claims specified in `RFC 7519 <https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.4>`_ are validated:
+  + If no JWK matches the same ``kid`` value (or if they do not have a ``kid``), then the token will be rejected with a :ref:`401 Unauthorized <pgrst301>` error.
+  + If a JWK matches the ``kid`` value then it will validate the token against that JWK accordingly.
 
-- ``exp`` Expiration Time
-- ``iat`` Issued At
-- ``nbf`` Not Before
+- If the JWT does not have a ``kid`` parameter, then PostgREST will validate the token against each JWK in the :ref:`jwt-secret`.
 
-We allow a 30-second clock skew when validating the above claims. In other words, we give an extra 30 seconds before the JWT is rejected if there is a slight discrepancy in the timestamps.
+.. _jwt_aud_verification:
 
-.. _jwt_aud_validation:
+``aud`` verification
+~~~~~~~~~~~~~~~~~~~~
 
-JWT ``aud`` Validation
-~~~~~~~~~~~~~~~~~~~~~~
-
-PostgREST has built-in validation of the `JWT audience claim <https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.3>`_.
+PostgREST has built-in verification of the `JWT audience claim <https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.3>`_.
 It works this way:
 
 - If :ref:`jwt-aud` is not set (the default), PostgREST identifies with all audiences and allows the JWT for any ``aud`` claim.
@@ -185,33 +184,37 @@ It works this way:
   + If the match fails or if the ``aud`` value is not a string or array of strings, then the token will be rejected with a :ref:`401 Unauthorized <pgrst303>` error.
   + If the ``aud`` key **is not present** or if its value is ``null`` or ``[]``, PostgREST will interpret this token as allowed for all audiences and will complete the request.
 
-JWK ``kid`` validation
-~~~~~~~~~~~~~~~~~~~~~~
+.. _jwt_claims_validation:
 
-PostgREST has built-in validation of the `key ID parameter <https://www.rfc-editor.org/rfc/rfc7517#section-4.5>`_, useful when working with a JWK Set.
-It goes as follows:
+JWT Claims Validation
+---------------------
 
-- If the JWT contains a ``kid`` parameter, then PostgREST will look for the JWK in the :ref:`jwt-secret`.
+The time-based JWT claims specified in `RFC 7519 <https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.4>`_ are validated:
 
-  + If no JWK matches the same ``kid`` value (or if they do not have a ``kid``), then the token will be rejected with a :ref:`401 Unauthorized <pgrst301>` error.
-  + If a JWK matches the ``kid`` value then it will validate the token against that JWK accordingly.
-
-- If the JWT does not have a ``kid`` parameter, then PostgREST will validate the token against each JWK in the :ref:`jwt-secret`.
+- ``exp`` Expiration Time
+- ``iat`` Issued At
+- ``nbf`` Not Before
 
+We allow a 30-second clock skew when validating the above claims. In other words, we give an extra 30 seconds before the JWT is rejected if there is a slight discrepancy in the timestamps.
 
 .. _jwt_caching:
 
 JWT Cache
 ---------
 
-PostgREST validates ``JWTs`` on every request. We can cache ``JWTs`` to avoid this performance overhead.
+JWT signature validation (specially :ref:`asym_keys` such as RSA) is slow, 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.
+The JWT cache is bounded and uses the `SIEVE algorithm <https://cachemon.github.io/SIEVE-website>`_ for efficient eviction. The cache is enabled by default and can be configured with :ref:`jwt-cache-max-entries`.
+
+It's recommended to leave the JWT cache enabled as our load tests indicate ~20% more throughput for simple GET requests when using it. This while reducing CPU utilization in exchange for a bit more memory.
+
+:ref:`jwt_cache_metrics` are available.
 
 .. note::
 
-  You can use the :ref:`server-timing_header` to see the effect of JWT caching.
+  - If the ``jwt-secret`` is changed and the config is reloaded, the JWT cache will reset.
+  - JWTs that pass :ref:`jwt_signature` are cached, regardless if they pass :ref:`jwt_claims_validation`. We do this to ensure responses stays fast under common failure cases (such as expired JWTs).
+  - You can use the :ref:`server-timing_header` to see the peformance benefit of JWT caching.
 
 .. _jwt_role_extract:
 

docs/references/configuration.rst

@@ -603,7 +603,7 @@ jwt-aud
   **In-Database** pgrst.jwt_aud
   =============== =================================
 
-    Specifies an audience for the JWT ``aud`` claim. See :ref:`jwt_aud_validation`.
+    Specifies an audience for the JWT ``aud`` claim. See :ref:`jwt_aud_verification`.
 
 .. _jwt-role-claim-key:
 
@@ -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:
 

docs/references/observability.rst

@@ -201,6 +201,40 @@ pgrst_db_pool_max
 
 Max pool connections.
 
+.. _jwt_cache_metrics:
+
+JWT Cache Metrics
+-----------------
+
+Metrics related to the :ref:`jwt_caching`.
+
+pgrst_jwt_cache_requests_total
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+======== =======
+**Type** Counter
+======== =======
+
+The total number of JWT cache lookups.
+
+pgrst_jwt_cache_hits_total
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+======== =======
+**Type** Counter
+======== =======
+
+The total number of JWT cache hits.
+
+pgrst_jwt_cache_evictions_total
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+======== =======
+**Type** Counter
+======== =======
+
+The total number of JWT cache evictions.
+
 Traces
 ======
 

nix/tools/loadtest.nix

@@ -59,6 +59,9 @@ let
         export PGRST_DB_TX_END="rollback-allow-override"
         export PGRST_LOG_LEVEL="crit"
         export PGRST_JWT_SECRET="reallyreallyreallyreallyverysafe"
+        # set previous PGRST_JWT_CACHE_MAX_LIFETIME configuration so that
+        # load test works across branches
+        # TODO clean once PGRST_JWT_CACHE_MAX_ENTRIES merged and released
         export PGRST_JWT_CACHE_MAX_LIFETIME="86400"
 
         mkdir -p "$(dirname "$_arg_output")"
@@ -67,6 +70,7 @@ let
         case "$_arg_kind" in
           jwt-hs)
             ${genTargetsHS} "$_arg_testdir"/gen_targets.http
+            export PGRST_JWT_CACHE_MAX_ENTRIES="0"
             export PGRST_JWT_CACHE_MAX_LIFETIME="0"
             ;;
 
@@ -80,6 +84,7 @@ let
 
           jwt-rsa)
             ${genTargetsHS} --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http
+            export PGRST_JWT_CACHE_MAX_ENTRIES="0"
             export PGRST_JWT_CACHE_MAX_LIFETIME="0"
             export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json"
             ;;

postgrest.cabal

@@ -50,6 +50,7 @@ library
                       PostgREST.Auth.Jwt
                       PostgREST.Auth.JwtCache
                       PostgREST.Auth.Types
+                      PostgREST.Cache.Sieve
                       PostgREST.CLI
                       PostgREST.Config
                       PostgREST.Config.Database
@@ -153,6 +154,10 @@ library
                     -- https://github.com/kazu-yamamoto/logger/commit/3a71ca70afdbb93d4ecf0083eeba1fbbbcab3fc3
                     , wai-logger                >= 2.4.0
                     , warp                      >= 3.3.19 && < 3.5
+                    , 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
@@ -207,6 +212,7 @@ test-suite spec
                       Feature.Auth.AudienceJwtSecretSpec
                       Feature.Auth.AuthSpec
                       Feature.Auth.BinaryJwtSecretSpec
+                      Feature.Auth.JwtCacheSpec
                       Feature.Auth.NoAnonSpec
                       Feature.Auth.NoJwtSecretSpec
                       Feature.ConcurrentSpec
@@ -264,6 +270,7 @@ test-suite spec
                     , hasql-transaction >= 1.0.1 && < 1.2
                     , 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
@@ -273,6 +280,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
 

src/PostgREST/Auth/JwtCache.hs

@@ -1,99 +1,114 @@
 {-|
 Module      : PostgREST.Auth.JwtCache
-Description : PostgREST Jwt Authentication Result Cache.
+Description : PostgREST JWT validation results Cache.
 
-This module provides functions to deal with the JWT cache
+This module provides functions to deal with the JWT cache.
 -}
-{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE FlexibleInstances         #-}
+{-# LANGUAGE LambdaCase                #-}
+{-# LANGUAGE MultiParamTypeClasses     #-}
+{-# LANGUAGE NamedFieldPuns            #-}
+{-# LANGUAGE StrictData                #-}
+
 module PostgREST.Auth.JwtCache
   ( init
+  , update
   , JwtCacheState
   , lookupJwtCache
-  , emptyCache
   ) where
 
 import qualified Data.Aeson        as JSON
 import qualified Data.Aeson.KeyMap as KM
-import qualified Data.Cache        as C
-import qualified Data.Scientific   as Sci
 
-import Control.Debounce
+import PostgREST.Error (Error (..), JwtError (JwtSecretMissing))
 
-import Data.Time.Clock       (UTCTime, nominalDiffTimeToSeconds)
-import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
-import System.Clock          (TimeSpec (..))
+import           Control.Concurrent.STM      (newTVarIO, readTVar,
+                                              writeTVar)
+import           Control.Concurrent.STM.TVar (TVar)
+import           Control.Monad.Error.Class   (liftEither)
+import           Data.ByteString             hiding (all, init)
+import           Data.IORef                  (IORef, newIORef,
+                                              readIORef, writeIORef)
+import           Jose.Jwk                    (JwkSet)
+import           PostgREST.Auth.Jwt          (parseAndDecodeClaims)
+import           PostgREST.Cache.Sieve       (alwaysValid)
+import qualified PostgREST.Cache.Sieve       as SC
+import           PostgREST.Config            (AppConfig (..))
+import           PostgREST.Observation       (Observation (JwtCacheEviction, JwtCacheLookup),
+                                              ObservationHandler)
+import           Protolude
 
-import PostgREST.Auth.Types (AuthResult (..))
-import PostgREST.Error      (Error (..))
+data JwtCacheState = JwtCacheState ObservationHandler (IORef JwtCache)
 
-import Protolude
+class CacheVariant m v where
+  cached :: SC.Cache m ByteString v -> ByteString -> ExceptT Error IO JSON.Object
 
--- | JWT Cache and IO action that triggers purging old entries from the cache
-data JwtCacheState = JwtCacheState
-  { jwtCache   :: C.Cache ByteString AuthResult
-  , purgeCache :: IO ()
-  }
+{-|
+Jwt caching can have three different configurations:
+* missing JWT Key (no caching and throw error when JWT token present in the request)
+* JWT cache turned off
+* JWT cache turned on
+
+All three options are represented by JwtCache data type.
+
+Handling of reconfiguration is centralized in this module.
+-}
+data JwtCache =
+  JwtNoJwks |
+  JwtNoCache JwkSet |
+  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
+
+instance CacheVariant (ExceptT Error IO) JSON.Object where
+  cached = SC.cached
+
+decode :: JwtCache -> ByteString -> ExceptT Error IO JSON.Object
+decode JwtNoJwks        = const $ throwError (JwtErr JwtSecretMissing)
+decode (JwtNoCache key) = parseAndDecodeClaims key
+decode (JwtCache _ _ c) = cached c
+
+-- | Reconfigure JWT caching and update JwtCacheState accordingly
+update :: JwtCacheState -> AppConfig -> IO ()
+update (JwtCacheState observationHandler jwtCacheState) config@AppConfig{configJWKS, configJwtCacheMaxEntries} =
+  let reinitialize =
+        newJwtCache config observationHandler
+          >>= writeIORef jwtCacheState
+  in
+  readIORef jwtCacheState >>= \case
+    (JwtCache decodingKey maxSize _) ->
+      if configJWKS /= Just decodingKey || configJwtCacheMaxEntries <= 0 then
+        -- reinitialize if key changed or cache disabled
+        reinitialize
+      else
+        -- max size changed - set it and let the cache shrink itself if necessary
+        atomically $ writeTVar maxSize configJwtCacheMaxEntries
+
+    _ -> reinitialize
+
+init :: AppConfig -> ObservationHandler -> IO JwtCacheState
+init config = fmap (<$>) JwtCacheState <*> (newJwtCache config >=> newIORef)
 
 -- | Initialize JwtCacheState
-init :: IO JwtCacheState
-init = do
-  cache <- C.newCache Nothing -- no default expiration
-  -- purgeExpired has O(n^2) complexity
-  -- so we wrap it in debounce to make sure it:
-  -- 1) is executed asynchronously
-  -- 2) only a single purge operation is running at a time
-  debounce <- mkDebounce defaultDebounceSettings
-    -- debounceFreq is set to default 1 second
-    { debounceAction = C.purgeExpired cache
-    , debounceEdge = leadingEdge
-    }
-  pure $ JwtCacheState cache debounce
+newJwtCache :: AppConfig -> ObservationHandler -> IO JwtCache
+newJwtCache AppConfig{configJWKS, configJwtCacheMaxEntries} observationHandler = do
+  maybe (pure JwtNoJwks) initCache configJWKS
+  where
+    initCache key = if configJwtCacheMaxEntries <= 0 then pure (JwtNoCache key) else createCache key configJwtCacheMaxEntries
 
--- | Used to retrieve and insert JWT to JWT Cache
-lookupJwtCache :: JwtCacheState -> ByteString -> Int -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult)
-lookupJwtCache JwtCacheState{jwtCache, purgeCache} token maxLifetime parseJwt utc = do
-  checkCache <- C.lookup jwtCache token
-  authResult <- maybe parseJwt (pure . Right) checkCache
+    createCache key maxSize = do
+          maxSizeTVar <- newTVarIO maxSize
+          JwtCache key maxSizeTVar <$>
+            notCachingErrors (readTVar maxSizeTVar) key
 
-  case (authResult,checkCache) of
-    -- From comment:
-    -- https://github.com/PostgREST/postgrest/pull/3801#discussion_r1857987914
-    --
-    -- We purge expired cache entries on a cache miss
-    -- The reasoning is that:
-    --
-    -- 1. We expect it to be rare (otherwise there is no point of the cache)
-    -- 2. It makes sure the cache is not growing (as inserting new entries
-    --    does garbage collection)
-    -- 3. Since this is time expiration based cache there is no real risk of
-    --    starvation - sooner or later we are going to have a cache miss.
+    notCachingErrors :: STM Int -> JwkSet -> IO (SC.Cache (ExceptT Error IO) ByteString JSON.Object)
+    notCachingErrors maxSize key = SC.cacheIO (SC.CacheConfig maxSize
+            (parseAndDecodeClaims key)
+            (lift . observationHandler . JwtCacheLookup) -- lookup metrics
+            (const . const $ lift $ observationHandler JwtCacheEviction) -- evictions metrics
+            alwaysValid) -- no invalidation for now
 
-    (Right res, Nothing) -> do -- cache miss
-
-      let timeSpec = getTimeSpec res maxLifetime utc
-
-      -- insert new cache entry
-      C.insert' jwtCache (Just timeSpec) token res
-
-      -- Execute IO action to purge the cache
-      -- It is assumed this action returns immidiately
-      -- so that request processing is not blocked.
-      purgeCache
-
-    _                    -> pure ()
-
-  return authResult
-
--- Used to extract JWT exp claim and add to JWT Cache
-getTimeSpec :: AuthResult -> Int -> UTCTime -> 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) -> TimeSpec (sciToInt seconds - utcToSecs utc) 0
-    _                          -> TimeSpec (fromIntegral maxLifetime :: Int64) 0
-
--- | Empty the cache (done when the config is reloaded)
-emptyCache :: JwtCacheState -> IO ()
-emptyCache JwtCacheState{jwtCache} = C.purge jwtCache
+lookupJwtCache :: JwtCacheState -> Maybe ByteString -> ExceptT Error IO JSON.Object
+lookupJwtCache (JwtCacheState _ cacheState) k = liftIO (readIORef cacheState) >>= flip (maybe (pure KM.empty)) k . decode

src/PostgREST/CLI.hs

@@ -203,8 +203,8 @@ 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
+      |## Enables JWT Cache and sets its max size, disables caching with 0
+      |# jwt-cache-max-entries = 0
       |
       |## Logging level, the admitted values are: crit, error, warn, info and debug.
       |log-level = "error"

src/PostgREST/Cache/Sieve.hs

@@ -0,0 +1,218 @@
+{-|
+Module      : PostgREST.Cache.Sieve
+Description : PostgREST cache implementation based on Sieve algorithm.
+
+This module provides implementation of a mutable cache on Sieve algorithm.
+-}
+{-# LANGUAGE DataKinds       #-}
+{-# LANGUAGE GADTs           #-}
+{-# LANGUAGE LambdaCase      #-}
+{-# LANGUAGE NamedFieldPuns  #-}
+{-# LANGUAGE PolyKinds       #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE RecursiveDo     #-}
+{-# LANGUAGE StrictData      #-}
+{-# LANGUAGE TupleSections   #-}
+
+module PostgREST.Cache.Sieve (
+      Cache
+    , CacheConfig (..)
+    , Discard (..)
+    , alwaysValid
+    , cache
+    , cacheIO
+    , cached
+)
+where
+
+import           Control.Concurrent.STM
+import           Control.Monad.Extra    (whileM)
+import           Data.Some
+import qualified Focus                  as F
+import           Protolude              hiding (elem, head)
+import qualified StmHamt.SizedHamt      as SH
+
+data ListNode k v (b :: Bool) = ListNode {
+        nextPtr        :: NodePtr k v,
+        prevNextPtrPtr :: NodePtrPtr k v,
+        elem           :: NodeElem k v b
+    }
+
+data NodeElem :: Type -> Type -> Bool -> Type where
+    Head :: {
+            entries          :: SH.SizedHamt (HamtEntry k v),
+            finger           :: NodePtrPtr k v
+        } -> NodeElem k v False
+    Entry :: Hashable k => {
+            visited :: TVar Bool,
+            ekey :: k,
+            entryValue :: v
+        } -> NodeElem k v True
+
+type HamtEntry k v = ListNode k v True
+type AnyNode k v = Some (ListNode k v)
+type NodePtr k v = TVar (AnyNode k v)
+type NodePtrPtr k v = TVar (NodePtr k v)
+
+data Discard m v = Refresh (m ()) | Invalid (m v)
+
+data Cache m k v = (MonadIO m, Hashable k) => Cache (ListNode k v False) (CacheConfig m k v)
+
+data CacheConfig m k v = CacheConfig {
+        maxSize          :: STM Int,
+        load             :: k -> m v,
+        requestListener  :: Bool -> m (),
+        evictionListener :: k -> v -> m (),
+        validator        :: m (k -> v -> Maybe (Discard m v))
+}
+
+alwaysValid :: Applicative m => m (k -> v -> Maybe (Discard m v))
+alwaysValid = pure (const . const Nothing)
+
+cacheIO :: (MonadIO m, Hashable k) => CacheConfig m k v -> IO (Cache m k v)
+cacheIO = atomically . cache
+
+cache :: (MonadIO m, Hashable k) => CacheConfig m k v -> STM (Cache m k v)
+cache cacheConfig = mdo
+    tail <- newTVar (Some head)
+    entries <- SH.new
+    finger <- newTVar tail
+    head <- ListNode tail <$> newTVar tail <*> pure Head {..}
+    pure $ Cache head cacheConfig
+
+cached :: Cache m k v -> k -> m v
+cached (Cache head@ListNode{prevNextPtrPtr=neck, elem=Head{..}} CacheConfig{..}) k = do
+    checkValid <- validator
+    tryMaybe
+        -- Fast path: lookup value, update stats and return the value if found and valid
+        ((liftIO . atomically) (lookup checkValid) >>= notify (requestListener . isJust) >>= validate)
+        -- Slow path: load/calculate value and insert it (if still not found)
+        (do
+            value <- load k
+            whileM (not <$> tryInsert value)
+            pure value)
+    where
+        tryMaybe f notFound = f >>= maybe notFound pure
+
+        notify = ((<$) <*>)
+
+        validate = fmap join . traverse (\case
+            -- valid value
+            (Right v) -> pure $ Just v
+            -- refresh value
+            (Left (Refresh act)) -> act $> Nothing
+            -- discard value and return alt result
+            (Left (Invalid res)) -> Just <$> res)
+
+        lookup checkValid = SH.focus focus (ekey . elem) k entries
+            where
+                focus = F.Focus
+                    -- not found
+                    (pure (Nothing, F.Leave))
+                    -- found
+                    -- check entry validity
+                    (\e@ListNode{elem=Entry{visited, entryValue}} ->
+                        maybe
+                            -- entry valid
+                            (mark visited True $> (Just $ Right entryValue, F.Leave))
+                            -- entry invalid
+                            -- remove it
+                            ((removeEntry e $>) . (, F.Remove) . Just . Left)
+                            (checkValid k entryValue)
+                    )
+
+        mark t b = whenM ((/= b) <$> readTVar t) (writeTVar t b)
+
+        -- perform a single entry eviction and possibly insertion atomically
+        -- returning False if could not insert
+        -- (either because entry currently pointed by the finger was visited
+        --  or because after this entry eviction the cache is still full)
+        -- so that other threads don't have to wait when visiting entries.
+        -- First check if entry is still not in the cache - this time inside transaction.
+        --
+        -- Execute evictionListener if an entry was evicted
+        tryInsert value = do
+            (result, evicted) <- liftIO . atomically $ do
+                -- Use SH.focus to performa a single lookup instead of 2
+                -- we cannot modify Hamt from inside focus
+                -- so if there is any entry to remove
+                -- we need to delete it after
+                (res, evictedKey) <- SH.focus focus (ekey . elem) k entries
+                case evictedKey of
+                    (Just Entry{ekey=entryKey, entryValue}) -> do
+                        SH.focus F.delete (ekey . elem) entryKey entries
+                        pure (res, evictionListener entryKey entryValue)
+                    Nothing -> pure (res, pure ())
+
+            evicted $> result
+            where
+                focus = F.Focus (do
+                    (hasSpace, evictedKey) <- evictionStep
+                    if hasSpace then do
+                        entry <- newLinkedEntry value
+                        -- done, maybe evicted, insert entry
+                        pure ((True, evictedKey), F.Set entry)
+                    else
+                        -- not done, maybe evicted, don't modify entries
+                        pure ((False, evictedKey), F.Leave))
+                    -- Entry found case
+                    (\ListNode{elem=Entry{visited}} -> do
+                        -- mark as visited
+                        mark visited True
+                        -- done, no evictions, don't modify entries
+                        pure ((True, Nothing), F.Leave))
+
+        -- if the cache is full precoesses a single node
+        -- removing it if it is marked as unvisited
+        -- or clearing visited mark
+        -- returns True if there is space in the cache
+        -- puts evictionListener in state if an entry was evicted
+        evictionStep = do
+            currDiff <- liftA2 (-) (SH.size entries) (max 1 <$> maxSize)
+            if currDiff >= 0 then do
+                -- no space in the cache
+                -- need to evict an entry
+                (nextFinger, evictedKey) <- readTVar finger >>= evict
+                writeTVar finger nextFinger
+                -- return if enough space and evicted key if any
+                pure (isJust evictedKey && currDiff == 0, evictedKey)
+            else
+                -- there is space in the cache
+                pure (True, Nothing)
+
+        evict :: TVar (Some (ListNode k v)) -> STM (NodePtr k v, Maybe (NodeElem k v True))
+        evict = readTVar >=> \case
+            (Some e@ListNode{nextPtr, prevNextPtrPtr, elem=elem@Entry{visited}}) -> do
+                ifM (readTVar visited)
+
+                    (writeTVar visited False $> (nextPtr, Nothing))
+
+                    (unlinkEntry e *> fmap (, Just elem) (readTVar prevNextPtrPtr))
+            -- skip head
+            (Some ListNode{nextPtr, elem=Head{}}) -> evict nextPtr
+
+        unlinkEntry :: HamtEntry k v -> STM ()
+        unlinkEntry (ListNode{nextPtr, prevNextPtrPtr=currPrev}) = do
+            nextEntry <- readTVar nextPtr
+            withSome nextEntry $ \e -> do
+                prevNextPtr <- readTVar currPrev
+                writeTVar (prevNextPtrPtr e) prevNextPtr
+                writeTVar prevNextPtr nextEntry
+
+        newLinkedEntry v = do
+            oldNeckNextPtr <- readTVar neck
+            newNeckNextPtr <- newTVar (Some head)
+            newNeck <- ListNode newNeckNextPtr <$>
+                newTVar oldNeckNextPtr <*>
+                (Entry <$> newTVar False <*> pure k <*> pure v)
+            -- update pointers
+            writeTVar oldNeckNextPtr (Some newNeck)
+            writeTVar neck newNeckNextPtr
+            -- return HAMT entry
+            pure newNeck
+
+        removeEntry = fmap (*>) unlinkEntry <*> adjustFinger
+
+        adjustFinger ListNode{nextPtr, prevNextPtrPtr} =
+            whenM ((nextPtr ==) <$> readTVar finger) $
+                readTVar prevNextPtrPtr >>= writeTVar finger

src/PostgREST/Config.hs

@@ -97,7 +97,7 @@ data AppConfig = AppConfig
   , configJwtRoleClaimKey          :: JSPath
   , configJwtSecret                :: Maybe BS.ByteString
   , configJwtSecretIsBase64        :: Bool
-  , configJwtCacheMaxLifetime      :: Int
+  , configJwtCacheMaxEntries       :: Int
   , configLogLevel                 :: LogLevel
   , configLogQuery                 :: LogQuery
   , configOpenApiMode              :: OpenAPIMode
@@ -177,7 +177,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)
+      ,("jwt-cache-max-entries",         show . configJwtCacheMaxEntries)
       ,("log-level",                 q . dumpLogLevel . configLogLevel)
       ,("log-query",                 q . dumpLogQuery . configLogQuery)
       ,("openapi-mode",              q . dumpOpenApiMode . configOpenApiMode)
@@ -287,7 +287,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")
+    <*> (fromMaybe 1000 <$> optInt "jwt-cache-max-entries")
     <*> parseLogLevel "log-level"
     <*> parseLogQuery "log-query"
     <*> parseOpenAPIMode "openapi-mode"

src/PostgREST/Logger.hs

@@ -100,6 +100,12 @@ observationLogger loggerState logLevel obs = case obs of
   o@PoolRequestFullfilled ->
     when (logLevel >= LogDebug) $ do
       logWithZTime loggerState $ observationMessage o
+  o@JwtCacheEviction ->
+    when (logLevel >= LogDebug) $ do
+      logWithZTime loggerState $ observationMessage o
+  o@(JwtCacheLookup _) ->
+    when (logLevel >= LogDebug) $ do
+      logWithZTime loggerState $ observationMessage o
   o ->
     logWithZTime loggerState $ observationMessage o
 

src/PostgREST/Metrics.hs

@@ -26,7 +26,10 @@ data MetricsState =
     poolWaiting          :: Gauge,
     poolMaxSize          :: Gauge,
     schemaCacheLoads     :: Vector Label1 Counter,
-    schemaCacheQueryTime :: Gauge
+    schemaCacheQueryTime :: Gauge,
+    jwtCacheRequests     :: Counter,
+    jwtCacheHits         :: Counter,
+    jwtCacheEvictions    :: Counter
   }
 
 init :: Int -> IO MetricsState
@@ -37,7 +40,10 @@ init configDbPoolSize = do
     register (gauge (Info "pgrst_db_pool_waiting" "Requests waiting to acquire a pool connection")) <*>
     register (gauge (Info "pgrst_db_pool_max" "Max pool connections")) <*>
     register (vector "status" $ counter (Info "pgrst_schema_cache_loads_total" "The total number of times the schema cache was loaded")) <*>
-    register (gauge (Info "pgrst_schema_cache_query_time_seconds" "The query time in seconds of the last schema cache load"))
+    register (gauge (Info "pgrst_schema_cache_query_time_seconds" "The query time in seconds of the last schema cache load")) <*>
+    register (counter (Info "pgrst_jwt_cache_requests_total" "The total number of JWT cache lookups")) <*>
+    register (counter (Info "pgrst_jwt_cache_hits_total" "The total number of JWT cache hits")) <*>
+    register (counter (Info "pgrst_jwt_cache_evictions_total" "The total number of JWT cache evictions"))
   setGauge (poolMaxSize metricState) (fromIntegral configDbPoolSize)
   pure metricState
 
@@ -63,6 +69,9 @@ observationMetrics MetricsState{..} obs = case obs of
     setGauge schemaCacheQueryTime resTime
   SchemaCacheErrorObs{} -> do
     withLabel schemaCacheLoads "FAIL" incCounter
+  JwtCacheLookup True -> incCounter jwtCacheRequests *> incCounter jwtCacheHits
+  JwtCacheLookup False -> incCounter jwtCacheRequests
+  JwtCacheEviction -> incCounter jwtCacheEvictions
   _ ->
     pure ()
 

src/PostgREST/Observation.hs

@@ -60,6 +60,8 @@ data Observation
   | HasqlPoolObs SQL.Observation
   | PoolRequest
   | PoolRequestFullfilled
+  | JwtCacheLookup Bool
+  | JwtCacheEviction
 
 data ObsFatalError = ServerAuthError | ServerPgrstBug | ServerError42P05 | ServerError08P01
 
@@ -151,6 +153,10 @@ observationMessage = \case
     "Trying to borrow a connection from pool"
   PoolRequestFullfilled ->
     "Borrowed a connection from the pool"
+  JwtCacheLookup _ ->
+    "Looked up a JWT in JWT cache"
+  JwtCacheEviction ->
+    "Evicted entry from JWT cache"
   where
     showMillis :: Double -> Text
     showMillis x = toS $ showFFloat (Just 1) (x * 1000) ""

test/io/configs/expected/aliases.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"aliased\""
 jwt-secret = ""
 jwt-secret-is-base64 = true
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/boolean-numeric.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"role\""
 jwt-secret = ""
 jwt-secret-is-base64 = true
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/boolean-string.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"role\""
 jwt-secret = ""
 jwt-secret-is-base64 = true
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/defaults.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"role\""
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/jwt-role-claim-key1.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"roles\"[?(@ == \"role1\")]"
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/jwt-role-claim-key2.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"roles\"[?(@ != \"role1\")]"
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/jwt-role-claim-key3.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"roles\"[?(@ ^== \"role1\")]"
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/jwt-role-claim-key4.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"roles\"[?(@ ==^ \"role1\")]"
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/jwt-role-claim-key5.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"roles\"[?(@ *== \"role1\")]"
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/expected/no-defaults-with-db-other-authenticator.config

@@ -23,7 +23,7 @@ jwt-aud = "https://otherexample.org"
 jwt-role-claim-key = ".\"other\".\"pre_config_role\""
 jwt-secret = "ODERREALLYREALLYREALLYREALLYVERYSAFE"
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 7200
+jwt-cache-max-entries = 86400
 log-level = "info"
 log-query = "main-query"
 openapi-mode = "disabled"

test/io/configs/expected/no-defaults-with-db.config

@@ -23,7 +23,7 @@ jwt-aud = "https://example.org"
 jwt-role-claim-key = ".\"a\".\"role\""
 jwt-secret = "OVERRIDE=REALLY=REALLY=REALLY=REALLY=VERY=SAFE"
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 3600
+jwt-cache-max-entries = 86400
 log-level = "info"
 log-query = "main-query"
 openapi-mode = "ignore-privileges"

test/io/configs/expected/no-defaults.config

@@ -23,7 +23,7 @@ jwt-aud = "https://postgrest.org"
 jwt-role-claim-key = ".\"user\"[0].\"real-role\""
 jwt-secret = "c2VjdXJpdHl0aHJvdWdob2JzY3VyaXR5aW5iYXNlNjQ="
 jwt-secret-is-base64 = true
-jwt-cache-max-lifetime = 86400
+jwt-cache-max-entries = 86400
 log-level = "info"
 log-query = "main-query"
 openapi-mode = "ignore-privileges"

test/io/configs/expected/types.config

@@ -23,7 +23,7 @@ jwt-aud = ""
 jwt-role-claim-key = ".\"role\""
 jwt-secret = ""
 jwt-secret-is-base64 = false
-jwt-cache-max-lifetime = 0
+jwt-cache-max-entries = 1000
 log-level = "error"
 log-query = "disabled"
 openapi-mode = "follow-privileges"

test/io/configs/no-defaults-env.yaml

@@ -26,7 +26,7 @@ PGRST_JWT_AUD: 'https://postgrest.org'
 PGRST_JWT_ROLE_CLAIM_KEY: '.user[0]."real-role"'
 PGRST_JWT_SECRET: c2VjdXJpdHl0aHJvdWdob2JzY3VyaXR5aW5iYXNlNjQ=
 PGRST_JWT_SECRET_IS_BASE64: true
-PGRST_JWT_CACHE_MAX_LIFETIME: 86400
+PGRST_JWT_CACHE_MAX_ENTRIES: 86400
 PGRST_LOG_LEVEL: info
 PGRST_LOG_QUERY: 'main-query'
 PGRST_OPENAPI_MODE: 'ignore-privileges'

test/io/configs/no-defaults.config

@@ -23,7 +23,7 @@ jwt-aud = "https://postgrest.org"
 jwt-role-claim-key = ".user[0].\"real-role\""
 jwt-secret = "c2VjdXJpdHl0aHJvdWdob2JzY3VyaXR5aW5iYXNlNjQ="
 jwt-secret-is-base64 = true
-jwt-cache-max-lifetime = 86400
+jwt-cache-max-entries = 86400
 log-level = "info"
 log-query = "main-query"
 openapi-mode = "ignore-privileges"

test/io/db_config.sql

@@ -14,7 +14,7 @@ ALTER ROLE db_config_authenticator SET pgrst.db_root_spec = 'root';
 ALTER ROLE db_config_authenticator SET pgrst.db_schemas = 'test, tenant1, tenant2';
 ALTER ROLE db_config_authenticator SET pgrst.db_tx_end = 'commit-allow-override';
 ALTER ROLE db_config_authenticator SET pgrst.jwt_aud = 'https://example.org';
-ALTER ROLE db_config_authenticator SET pgrst.jwt_cache_max_lifetime = '3600';
+ALTER ROLE db_config_authenticator SET pgrst.jwt_cache_max_entries = '86400';
 ALTER ROLE db_config_authenticator SET pgrst.jwt_role_claim_key = '."a"."role"';
 ALTER ROLE db_config_authenticator SET pgrst.jwt_secret = 'REALLY=REALLY=REALLY=REALLY=VERY=SAFE';
 ALTER ROLE db_config_authenticator SET pgrst.jwt_secret_is_base64 = 'false';
@@ -68,7 +68,7 @@ ALTER ROLE other_authenticator SET pgrst.db_schemas = 'test, other_tenant1, othe
 ALTER ROLE other_authenticator SET pgrst.jwt_aud = 'https://otherexample.org';
 ALTER ROLE other_authenticator SET pgrst.jwt_secret = 'ODERREALLYREALLYREALLYREALLYVERYSAFE';
 ALTER ROLE other_authenticator SET pgrst.jwt_secret_is_base64 = 'false';
-ALTER ROLE other_authenticator SET pgrst.jwt_cache_max_lifetime = '7200';
+ALTER ROLE other_authenticator SET pgrst.jwt_cache_max_entries = '86400';
 ALTER ROLE other_authenticator SET pgrst.openapi_mode = 'disabled';
 ALTER ROLE other_authenticator SET pgrst.openapi_security_active = 'false';
 ALTER ROLE other_authenticator SET pgrst.openapi_server_proxy_uri = 'https://otherexample.org/api';

test/io/test_io.py

@@ -152,7 +152,7 @@ def test_jwt_errors(defaultenv):
     env = {
         **defaultenv,
         "PGRST_SERVER_TIMING_ENABLED": "true",
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",
         "PGRST_JWT_SECRET": SECRET,
     }
 
@@ -165,7 +165,7 @@ def test_jwt_errors(defaultenv):
     env = {
         **defaultenv,
         "PGRST_SERVER_TIMING_ENABLED": "false",
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",
         "PGRST_JWT_SECRET": SECRET,
     }
 
@@ -1446,7 +1446,7 @@ def test_jwt_cache_server_timing(defaultenv):
     env = {
         **defaultenv,
         "PGRST_SERVER_TIMING_ENABLED": "true",
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",
         "PGRST_JWT_SECRET": SECRET,
         "PGRST_DB_CONFIG": "false",
     }
@@ -1482,7 +1482,7 @@ def test_jwt_cache_without_server_timing(defaultenv):
     env = {
         **defaultenv,
         "PGRST_SERVER_TIMING_ENABLED": "false",
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",
         "PGRST_JWT_SECRET": SECRET,
         "PGRST_DB_CONFIG": "false",
     }
@@ -1503,7 +1503,7 @@ def test_jwt_cache_without_exp_claim(defaultenv):
     env = {
         **defaultenv,
         "PGRST_SERVER_TIMING_ENABLED": "true",
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",
         "PGRST_JWT_SECRET": SECRET,
         "PGRST_DB_CONFIG": "false",
     }
@@ -1772,7 +1772,7 @@ def test_jwt_cache_purges_expired_entries(defaultenv):
 
     env = {
         **defaultenv,
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",
         "PGRST_JWT_SECRET": SECRET,
         "PGRST_DB_CONFIG": "false",
     }
@@ -1838,10 +1838,10 @@ def test_log_pool_req_observation(level, defaultenv):
         postgrest.session.get("/authors_only", headers=headers)
 
         if level == "debug":
-            output = postgrest.read_stdout(nlines=4)
-            assert pool_req in output[0]
-            assert pool_req_fullfill in output[3]
-            assert len(output) == 4
+            output = postgrest.read_stdout(nlines=5)
+            assert pool_req in output[1]
+            assert pool_req_fullfill in output[4]
+            assert len(output) == 5
         elif level == "info":
             output = postgrest.read_stdout(nlines=4)
             assert len(output) == 1
@@ -1882,7 +1882,7 @@ def test_invalidate_jwt_cache_when_secret_changes(tmp_path, defaultenv):
         **defaultenv,
         "PGRST_JWT_SECRET": f"@{external_secret_file}",
         "PGRST_DB_CHANNEL_ENABLED": "true",
-        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",  # enable cache
+        "PGRST_JWT_CACHE_MAX_ENTRIES": "86400",  # enable cache
         "PGRST_DB_ANON_ROLE": "postgrest_test_anonymous",  # required for NOTIFY
     }
 

test/spec/Feature/Auth/JwtCacheSpec.hs

@@ -0,0 +1,167 @@
+{-# LANGUAGE AllowAmbiguousTypes       #-}
+{-# LANGUAGE DataKinds                 #-}
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE FlexibleContexts          #-}
+{-# LANGUAGE ImpredicativeTypes        #-}
+{-# LANGUAGE OverloadedStrings         #-}
+{-# LANGUAGE ScopedTypeVariables       #-}
+{-# LANGUAGE TypeApplications          #-}
+module Feature.Auth.JwtCacheSpec
+
+where
+
+import Network.Wai (Application)
+
+import Network.HTTP.Types
+import Test.Hspec         (Expectation, SpecWith, describe, it,
+                           shouldBe)
+import Test.Hspec.Wai
+
+import Data.String                     (String)
+import PostgREST.Metrics               (MetricsState (..))
+import Prometheus                      (getCounter)
+import Protolude
+import SpecHelper
+import Test.Hspec.Expectations.Contrib (annotate)
+import Test.Hspec.Wai.JSON             (json)
+
+spec :: SpecWith (MetricsState, Application)
+spec = describe "Server started with JWT and metrics enabled" $ do
+  it "Should not have JWT in cache" $ do
+    let auth = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe1"}|]
+
+    expectCounters
+      [
+        requests (+ 1)
+      , hits     (+ 0)
+      ] $
+
+         request methodGet "/authors_only" [auth] ""
+
+  it "Should have JWT in cache" $ do
+    let auth = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe2"}|]
+
+    expectCounters
+      [
+        requests (+ 2)
+      , hits     (+ 1)
+      ] $
+
+         request methodGet "/authors_only" [auth] "" `shouldRespondWith` 200
+      *> request methodGet "/authors_only" [auth] "" `shouldRespondWith` 200
+
+  it "Should not cache invalid JWTs" $ do
+    let auth = authHeaderJWT "some random bytes"
+
+    expectCounters
+      [
+        requests (+ 2)
+      , hits     (+ 0)
+      ] $
+
+         request methodGet "/authors_only" [auth] "" `shouldRespondWith` 401
+      *> request methodGet "/authors_only" [auth] "" `shouldRespondWith` 401
+
+  it "Should cache expired JWTs" $ do
+    let auth = genToken [json|{"exp": 1, "role": "postgrest_test_author", "id": "jdoe2"}|]
+
+    expectCounters
+      [
+        requests (+ 2)
+      , hits     (+ 1)
+      ] $
+
+         request methodGet "/authors_only" [auth] "" `shouldRespondWith` 401
+      *> request methodGet "/authors_only" [auth] "" `shouldRespondWith` 401
+
+  it "Should evict entries from the JWT cache (jwt cache max is 2)" $ do
+    let jwt1 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe3"}|]
+        jwt2 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe4"}|]
+        jwt3 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe5"}|]
+
+    expectCounters
+      [
+        requests  (+ 6)
+      , hits      (+ 0)
+      , evictions (+ 4)
+      ] $
+
+         request methodGet "/authors_only" [jwt1] ""
+      *> request methodGet "/authors_only" [jwt2] ""
+      *> request methodGet "/authors_only" [jwt3] ""
+      *> request methodGet "/authors_only" [jwt1] ""
+      *> request methodGet "/authors_only" [jwt2] ""
+      *> request methodGet "/authors_only" [jwt3] ""
+
+  it "Should not evict entries from the JWT cache in FIFO order" $ do
+    let jwt1 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe6"}|]
+        jwt2 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe7"}|]
+        jwt3 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe8"}|]
+
+    expectCounters
+      [
+        requests  (+ 6)
+      , hits      (+ 3)
+      , evictions (+ 1)
+      ] $
+
+         request methodGet "/authors_only" [jwt1] ""
+      *> request methodGet "/authors_only" [jwt2] ""
+      -- this one should hit the cache
+      *> request methodGet "/authors_only" [jwt1] ""
+      -- this one should trigger eviction of jwt2 (not FIFO)
+      *> request methodGet "/authors_only" [jwt3] ""
+      -- these two should hit the cache
+      *> request methodGet "/authors_only" [jwt1] ""
+      *> request methodGet "/authors_only" [jwt3] ""
+
+  -- This one makes sure we test the scenario when finger
+  -- has to move through the whole list first and pass the head
+  -- The test case was added based on coverage report
+  -- showing this scenario was not covered by previous tests
+  it "Should evict entries even though all were hit" $ do
+    let jwt1 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe9"}|]
+        jwt2 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe10"}|]
+        jwt3 = genToken [json|{"exp": 9999999999, "role": "postgrest_test_author", "id": "jdoe11"}|]
+
+    expectCounters
+      [
+        requests  (+ 7)
+      , hits      (+ 4)
+      , evictions (+ 1)
+      ] $
+
+         request methodGet "/authors_only" [jwt1] ""
+      *> request methodGet "/authors_only" [jwt2] ""
+      -- these two should hit the cache
+      *> request methodGet "/authors_only" [jwt1] ""
+      *> request methodGet "/authors_only" [jwt2] ""
+      -- this one should trigger eviction of jwt1
+      *> request methodGet "/authors_only" [jwt3] ""
+      -- these two should hit the cache
+      *> request methodGet "/authors_only" [jwt2] ""
+      *> request methodGet "/authors_only" [jwt3] ""
+
+  where
+      counterToInt = second (fmap (round @Double @Int) . getCounter)
+      expectCounters = stateCheck . fmap (\(g, h) -> StateCheck (counterToInt . g) (flip shouldBe . h))
+      genToken = authHeaderJWT . generateJWT
+      requests = (,) (getF @"jwtCacheRequests")
+      hits = (,) (getF @"jwtCacheHits")
+      evictions = (,) (getF @"jwtCacheEvictions")
+
+
+-- should be moved to helpers???
+getF :: forall s r a. (KnownSymbol s, HasField s r a) => r -> (String, a)
+getF r = (symbolVal (Proxy @s), getField @s r)
+
+data StateCheck st = forall a. (Show a, Eq a) => StateCheck (st -> (String, WaiSession st a)) (a -> a -> Expectation)
+
+stateCheck :: (Traversable t) => t (StateCheck st) -> WaiSession st a -> WaiSession st ()
+stateCheck checks act = do
+  metrics <- getState
+  expectations <- traverse (\(StateCheck g expect) -> let (msg, m) = g metrics in m >>= createExpectation msg m . expect) checks
+  void act
+  sequenceA_ expectations
+  where
+    createExpectation msg metrics expect = pure $ metrics >>= liftIO . annotate msg . expect

test/spec/Main.hs

@@ -15,15 +15,15 @@ import PostgREST.SchemaCache     (querySchemaCache)
 import Protolude                 hiding (toList, toS)
 import SpecHelper
 
-import qualified PostgREST.AppState      as AppState
-import qualified PostgREST.Auth.JwtCache as JwtCache
-import qualified PostgREST.Logger        as Logger
-import qualified PostgREST.Metrics       as Metrics
+import qualified PostgREST.AppState as AppState
+import qualified PostgREST.Logger   as Logger
+import qualified PostgREST.Metrics  as Metrics
 
 import qualified Feature.Auth.AsymmetricJwtSpec
 import qualified Feature.Auth.AudienceJwtSecretSpec
 import qualified Feature.Auth.AuthSpec
 import qualified Feature.Auth.BinaryJwtSecretSpec
+import qualified Feature.Auth.JwtCacheSpec
 import qualified Feature.Auth.NoAnonSpec
 import qualified Feature.Auth.NoJwtSecretSpec
 import qualified Feature.ConcurrentSpec
@@ -85,24 +85,23 @@ main = do
   -- cached schema cache so most tests run fast
   baseSchemaCache <- loadSCache pool testCfg
   sockets <- AppState.initSockets testCfg
-  jwtCacheState <- JwtCache.init
   loggerState <- Logger.init
   metricsState <- Metrics.init (configDbPoolSize testCfg)
 
   let
-    initApp sCache config = do
-      appState <- AppState.initWithPool sockets pool config jwtCacheState loggerState metricsState (const $ pure ())
+    initApp sCache st config = do
+      appState <- AppState.initWithPool sockets pool config loggerState metricsState (Metrics.observationMetrics metricsState)
       AppState.putPgVersion appState actualPgVersion
       AppState.putSchemaCache appState (Just sCache)
-      return ((), postgrest (configLogLevel config) appState (pure ()))
+      return (st, postgrest (configLogLevel config) appState (pure ()))
 
     -- For tests that run with the same schema cache
-    app = initApp baseSchemaCache
+    app = initApp baseSchemaCache ()
 
     -- For tests that run with a different SchemaCache (depends on configSchemas)
     appDbs config = do
       customSchemaCache <- loadSCache pool config
-      initApp customSchemaCache config
+      initApp customSchemaCache () config
 
   let withApp              = app testCfg
       maxRowsApp           = app testMaxRowsCfg
@@ -276,6 +275,9 @@ main = do
     before pgSafeUpdateApp $
       describe "Feature.Query.PgSafeUpdateSpec.spec" Feature.Query.PgSafeUpdateSpec.spec
 
+    before (initApp baseSchemaCache metricsState testCfgJwtCache) $
+      describe "Feature.Auth.JwtCacheSpec" Feature.Auth.JwtCacheSpec.spec
+
   where
     loadSCache pool conf =
       either (panic.show) id <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache conf)

test/spec/SpecHelper.hs

@@ -140,7 +140,7 @@ baseCfg = let secret = encodeUtf8 "reallyreallyreallyreallyverysafe" in
   , configJwtRoleClaimKey           = [JSPKey "role"]
   , configJwtSecret                 = Just secret
   , configJwtSecretIsBase64         = False
-  , configJwtCacheMaxLifetime       = 0
+  , configJwtCacheMaxEntries        = 10
   , configLogLevel                  = LogCrit
   , configLogQuery                  = LogQueryDisabled
   , configOpenApiMode               = OAFollowPriv
@@ -205,6 +205,14 @@ testCfgBinaryJWT =
   , configJWKS = rightToMaybe $ parseSecret generateSecret
   }
 
+testCfgJwtCache :: AppConfig
+testCfgJwtCache =
+  baseCfg {
+    configJwtSecret = Just generateSecret
+  , configJWKS = rightToMaybe $ parseSecret generateSecret
+  , configJwtCacheMaxEntries = 2
+  }
+
 testCfgAudienceJWT :: AppConfig
 testCfgAudienceJWT =
   baseCfg {