feat: JWT cache implementation based on sieve algorithm

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:
-
-JWT Claims Validation
----------------------
+``kid`` verification
+^^^^^^^^^^^^^^^^^^^^
 
-JWT ``exp``, ``iat`` , ``nbf`` 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:
 
-The time-based JWT claims specified in `RFC 7519 <https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.4>`_ are validated:
+- If the JWT contains a ``kid`` parameter, then PostgREST will look for the JWK in the :ref:`jwt-secret`.
 
-- ``exp`` Expiration Time
-- ``iat`` Issued At
-- ``nbf`` Not Before
+  + 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.
 
-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.
+- If the JWT does not have a ``kid`` parameter, then PostgREST will validate the token against each JWK in the :ref:`jwt-secret`.
 
-.. _jwt_aud_validation:
+.. _jwt_aud_verification:
 
-JWT ``aud`` Validation
-~~~~~~~~~~~~~~~~~~~~~~
+``aud`` verification
+~~~~~~~~~~~~~~~~~~~~
 
-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
-~~~~~~~~~~~~~~~~~~~~~~
-
-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`.
+JWT Claims Validation
+---------------------
 
-  + 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.
+The time-based JWT claims specified in `RFC 7519 <https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.4>`_ are validated:
 
-- 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.
+
+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.
 
-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.
+: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

@@ -316,7 +316,7 @@ db-extra-search-path
   Multiple schemas can be added in a comma-separated string, e.g. ``public, extensions``.
 
 .. important::
-  
+
   We default this config to ``public`` because it is the most common schema used to install PostgreSQL extensions such as :ref:`PostGIS <ww_postgis>`. You can disable this by setting this config to ``""``.
 
 .. _db-hoisted-tx-settings:
@@ -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