feat: change JWT cache to limited LRU based cache

CHANGELOG.md

@@ -21,6 +21,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
    + The selected columns in the embedded resources are aggregated into arrays
    + Aggregates are not supported
  - #2967, Add `Proxy-Status` header for better error response - @taimoorzaeem
+ - #4003, Add config `jwt-cache-max-entries` for maximum number of cached entries - @taimoorzaeem
 
 ### Fixed
 
@@ -55,6 +56,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
    + Diagnostic error messages instead of exposed internals
    + Return new `PGRST303` error when jwt claims decoding fails
  - #3906, Return `PGRST125` and `PGRST126` errors instead of empty json - @taimoorzaeem
+ - #4003, Remove config `jwt-cache-max-lifetime` and add config `jwt-cache-max-entries` for JWT cache - @taimoorzaeem
 
 ## [12.2.10] - 2025-04-18
 

docs/postgrest.dict

@@ -23,6 +23,7 @@ Cloudflare
 config
 cors
 CORS
+cryptographic
 cryptographically
 CSV
 durations
@@ -74,6 +75,7 @@ JSON
 JSPath
 JWK
 JWT
+JWTs
 jwt
 Keycloak
 Kubernetes
@@ -84,6 +86,7 @@ Logins
 LIBPQ
 logins
 lon
+LRU
 lt
 lte
 macOS

docs/references/auth.rst

@@ -98,9 +98,9 @@ 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 caches ``JWTs`` on every request to avoid performance overhead of parsing and cryptographic operations.
 
-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.
+To disable JWT caching, the config :code:`jwt-cache-max-entries` is to be set to ``0``. It is the maximum number of JWTs for which the cache stores their validation results. If the cache reaches its maximum, the `least recently used <https://redis.io/glossary/lru-cache/>`_ entry will be removed. The cache honors :code:`exp` claim. If the JWT does not have an :code:`exp` claim, it is cached until it gets removed by the LRU policy. See :ref:`jwt-cache-max-entries` for configuration 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 JWTs that can be cached. Set to :code:`0` to disable caching. See :ref:`jwt_caching`.
 
 .. _log-level:
 

nix/tools/loadtest.nix

@@ -58,7 +58,7 @@ 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"
+        export PGRST_JWT_CACHE_MAX_ENTRIES="1000" # default
 
         mkdir -p "$(dirname "$_arg_output")"
         abs_output="$(realpath "$_arg_output")"

postgrest.cabal

@@ -99,10 +99,8 @@ library
                     , auto-update               >= 0.1.4 && < 0.2
                     , base64-bytestring         >= 1 && < 1.3
                     , bytestring                >= 0.10.8 && < 0.13
-                    , cache                     >= 0.1.3 && < 0.2.0
                     , case-insensitive          >= 1.2 && < 1.3
                     , cassava                   >= 0.4.5 && < 0.6
-                    , clock                     >= 0.8.3 && < 0.9.0
                     , configurator-pg           >= 0.2 && < 0.3
                     , containers                >= 0.5.7 && < 0.7
                     , cookie                    >= 0.4.2 && < 0.5
@@ -122,6 +120,7 @@ library
                     , jose-jwt                  >= 0.9.6 && < 0.11
                     , lens                      >= 4.14 && < 5.3
                     , lens-aeson                >= 1.0.1 && < 1.3
+                    , lrucache                  >= 1.2.0.1 && < 1.3
                     , mtl                       >= 2.2.2 && < 2.4
                     , neat-interpolation        >= 0.5 && < 0.6
                     , network                   >= 2.6 && < 3.2

src/PostgREST/AppState.hs

@@ -57,7 +57,6 @@ import Data.IORef         (IORef, atomicWriteIORef, newIORef,
                            readIORef)
 import Data.Time.Clock    (UTCTime, getCurrentTime)
 
-import PostgREST.Auth.JwtCache           (JwtCacheState)
 import PostgREST.Config                  (AppConfig (..),
                                           addFallbackAppName,
                                           readAppConfig)
@@ -105,8 +104,8 @@ data AppState = AppState
   , stateSocketAdmin       :: Maybe NS.Socket
   -- | Observation handler
   , stateObserver          :: ObservationHandler
-  -- | JWT Cache
-  , stateJwtCache          :: JwtCache.JwtCacheState
+  -- | JWT Cache, disabled when config jwt-cache-max-entries is set to 0
+  , stateJwtCache          :: IORef JwtCache.JwtCacheState
   , stateLogger            :: Logger.LoggerState
   , stateMetrics           :: Metrics.MetricsState
   }
@@ -120,14 +119,14 @@ data SchemaCacheStatus
 type AppSockets = (NS.Socket, Maybe NS.Socket)
 
 init :: AppConfig -> IO AppState
-init conf@AppConfig{configLogLevel, configDbPoolSize} = do
+init conf = do
   loggerState  <- Logger.init
-  metricsState <- Metrics.init configDbPoolSize
-  let observer = liftA2 (>>) (Logger.observationLogger loggerState configLogLevel) (Metrics.observationMetrics metricsState)
+  metricsState <- Metrics.init (configDbPoolSize conf)
+  let observer = liftA2 (>>) (Logger.observationLogger loggerState (configLogLevel conf)) (Metrics.observationMetrics metricsState)
 
   observer $ AppStartObs prettyVersion
 
-  jwtCacheState <- JwtCache.init
+  jwtCacheState <- JwtCache.init (configJwtCacheMaxEntries conf)
   pool <- initPool conf observer
   (sock, adminSock) <- initSockets conf
   state' <- initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState observer
@@ -150,7 +149,7 @@ initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState
     <*> pure sock
     <*> pure adminSock
     <*> pure observer
-    <*> pure jwtCacheState
+    <*> newIORef jwtCacheState
     <*> pure loggerState
     <*> pure metricsState
 
@@ -311,8 +310,8 @@ putConfig = atomicWriteIORef . stateConf
 getTime :: AppState -> IO UTCTime
 getTime = stateGetTime
 
-getJwtCacheState :: AppState -> JwtCacheState
-getJwtCacheState = stateJwtCache
+getJwtCacheState :: AppState -> IO JwtCache.JwtCacheState
+getJwtCacheState = readIORef . stateJwtCache
 
 getSocketREST :: AppState -> NS.Socket
 getSocketREST = stateSocketREST
@@ -473,8 +472,9 @@ readInDbConfig startingUp appState@AppState{stateObserver=observer} = do
       -- 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
+      else do
+        newJwtCacheState <- JwtCache.init (configJwtCacheMaxEntries newConf)
+        atomicWriteIORef (stateJwtCache appState) newJwtCacheState
 
       if startingUp then
         pass

src/PostgREST/Auth.hs

@@ -160,30 +160,27 @@ middleware appState app req respond = do
 
   let token  = Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
       parseJwt = runExceptT $ parseToken conf token time >>= parseClaims conf
-      jwtCacheState = getJwtCacheState appState
+
+  jwtCacheState <- getJwtCacheState appState
 
 -- If ServerTimingEnabled -> calculate JWT validation time
--- If JwtCacheMaxLifetime -> cache JWT validation result
-  req' <- case (configServerTimingEnabled 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 $ 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 }
+  req' <- if configServerTimingEnabled conf then do
+
+      (dur, authResult) <- timeItT $ case token of
+
+          Just tkn -> lookupJwtCache jwtCacheState tkn parseJwt time
+          Nothing  -> parseJwt
 
-    (False, 0)           -> do
-          authResult <- parseJwt
-          return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
+      return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
+
+    else do
+
+        authResult <- case token of
+
+            Just tkn -> lookupJwtCache jwtCacheState tkn parseJwt time
+            Nothing  -> parseJwt
 
-    (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 }
+        return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
 
   app req' respond
 

src/PostgREST/Auth/JwtCache.hs

@@ -4,96 +4,100 @@
 
 This module provides functions to deal with the JWT cache
 -}
-{-# LANGUAGE NamedFieldPuns #-}
 module PostgREST.Auth.JwtCache
   ( init
   , 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.Cache.LRU    as C
+import qualified Data.IORef        as I
 import qualified Data.Scientific   as Sci
 
-import Control.Debounce
-
 import Data.Time.Clock       (UTCTime, nominalDiffTimeToSeconds)
 import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
-import System.Clock          (TimeSpec (..))
+import GHC.Num               (integerFromInt)
 
 import PostgREST.Auth.Types (AuthResult (..))
 import PostgREST.Error      (Error (..))
 
 import Protolude
 
--- | JWT Cache and IO action that triggers purging old entries from the cache
-data JwtCacheState = JwtCacheState
-  { jwtCache   :: C.Cache ByteString AuthResult
-  , purgeCache :: IO ()
+-- | Jwt Cache State
+newtype JwtCacheState = JwtCacheState
+  { maybeJwtCache :: Maybe (I.IORef (C.LRU ByteString AuthResult))
   }
 
 -- | 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
+init :: Int -> IO JwtCacheState
+init 0          = return $ JwtCacheState Nothing
+init maxEntries = do
+  cache <- I.newIORef $ C.newLRU (Just $ integerFromInt maxEntries)
+  return $ JwtCacheState $ Just cache
+
 
 -- | 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
-
-  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.
-
-    (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 -> ByteString -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult)
+lookupJwtCache jwtCacheState token parseJwt utc = do
+  case maybeJwtCache jwtCacheState of
+    Nothing            -> parseJwt
+    Just jwtCacheIORef -> do
+      -- get cache from IORef
+      jwtCache <- I.readIORef jwtCacheIORef
+
+      -- MAKE SURE WE UPDATE THE CACHE ON ALL PATHS AFTER LOOKUP
+      -- This is because it is a pure LRU cache, so lookup returns the
+      -- the cache with new state, hence it should be updated
+      let (jwtCache', maybeVal) = C.lookup token jwtCache
+
+      case maybeVal of
+        Nothing -> do -- CACHE MISS
+
+            -- When we get a cache miss, we get the parse result, insert it
+            -- into the cache. After that, we write the cache IO ref with
+            -- updated cache
+            authResult <- parseJwt
+
+            case authResult of
+              Right result -> do
+                  -- insert token -> update cache -> return token
+                  let jwtCache'' = C.insert token result jwtCache'
+                  I.writeIORef jwtCacheIORef jwtCache''
+                  return $ Right result
+              Left e -> do
+                  -- update cache after lookup -> return error
+                  I.writeIORef jwtCacheIORef jwtCache'
+                  return $ Left e
+
+        Just result -> -- CACHE HIT
+
+            -- For cache hit, we get the result from cache, we check the
+            -- exp claim. If it expired, we delete it from cache and parse
+            -- the jwt. Otherwise, the hit result is valid, so we return it
+
+            if isExpClaimExpired result utc then do
+                -- delete token -> update cache -> parse token
+                let (jwtCache'',_) = C.delete token jwtCache'
+                I.writeIORef jwtCacheIORef jwtCache''
+                parseJwt
+            else do
+                -- update cache after lookup -> return result
+                I.writeIORef jwtCacheIORef jwtCache'
+                return $ Right result
+
+
+type Expired = Bool
+
+-- | Check if exp claim is expired when looked up from cache
+isExpClaimExpired :: AuthResult -> UTCTime -> Expired
+isExpClaimExpired result utc =
+    case expireJSON of
+      Nothing                      -> False -- if exp not present then it is valid
+      Just (JSON.Number expiredAt) -> (sciToInt expiredAt - now) < 0
+      Just _                       -> False -- if exp is not a number then valid
+      where
+        expireJSON = KM.lookup "exp" (authClaims result)
+        now = (floor . nominalDiffTimeToSeconds . utcTimeToPOSIXSeconds) utc :: Int
+        sciToInt = fromMaybe 0 . Sci.toBoundedInteger