feat: change JWT cache to limited LRU based cache

BREAKING CHANGE

Our JWT cache implementation had no upper bound for number of cache
entries. This caused OOM errors. Additionally, the purge mechanism for
expired entries was quite slow. This changes our implementation to a
LRU based cache which limits the amount of cached entries.

Author: taimoorzaeem

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
 
@@ -54,6 +55,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:
 

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,99 @@ Description : PostgREST Jwt Authentication Result Cache.
 
 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
+
+      -- lookup key = token
+      -- MAKE SURE WE UPDATE THE CACHE ON ALL PATHS AFTER LOOKUP (changes cache)
+      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 afte 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

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
+      |## Maximum number of auth token that can be cached
+      |# jwt-cache-max-entries = 1000
       |
       |## Logging level, the admitted values are: crit, error, warn, info and debug.
       |log-level = "error"