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)
This commit is contained in:
Michal Kleczek
2025-07-29 18:51:41 -05:00
committed by GitHub
parent ac155a9391
commit 77ff11de95
35 changed files with 664 additions and 203 deletions
+1 -1
View File
@@ -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"
@@ -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"
@@ -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"
+1 -1
View File
@@ -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"
@@ -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"
@@ -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"
@@ -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"
@@ -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"
@@ -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"
@@ -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"
@@ -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"
+1 -1
View File
@@ -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"
+1 -1
View File
@@ -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"
+1 -1
View File
@@ -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'
+1 -1
View File
@@ -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"
+2 -2
View File
@@ -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';
+11 -11
View File
@@ -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
}
+167
View File
@@ -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
+12 -10
View File
@@ -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)
+9 -1
View File
@@ -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 {