feat: change JWT cache to limited LRU based cache

BREAKING CHANGE

Author: taimoorzaeem

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
@@ -158,7 +157,7 @@ library
                       -- -optP-Wno-nonportable-include-path
                       --   prevents build failures on case-insensitive filesystems (macos),
                       --   see https://github.com/commercialhaskell/stack/issues/3918
-  ghc-options:        -Werror -Wall -fwarn-identities
+  ghc-options:        -Wall -fwarn-identities
                       -fno-spec-constr -optP-Wno-nonportable-include-path
 
   if flag(dev)

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          :: Maybe JwtCache.JwtCacheState
   , stateLogger            :: Logger.LoggerState
   , stateMetrics           :: Metrics.MetricsState
   }
@@ -120,20 +119,20 @@ 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
   pure state' { stateSocketREST = sock, stateSocketAdmin = adminSock}
 
-initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> JwtCache.JwtCacheState -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
+initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> Maybe JwtCache.JwtCacheState -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
 initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState observer = do
 
   appState <- AppState pool
@@ -311,7 +310,7 @@ putConfig = atomicWriteIORef . stateConf
 getTime :: AppState -> IO UTCTime
 getTime = stateGetTime
 
-getJwtCacheState :: AppState -> JwtCacheState
+getJwtCacheState :: AppState -> Maybe JwtCache.JwtCacheState
 getJwtCacheState = stateJwtCache
 
 getSocketREST :: AppState -> NS.Socket

src/PostgREST/Auth.hs

@@ -163,27 +163,23 @@ middleware appState app req respond = do
       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
+
+      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 }
+    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

@@ -13,67 +13,96 @@ module PostgREST.Auth.JwtCache
 
 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 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
 
 newtype JwtCacheState = JwtCacheState
-  { jwtCache :: C.Cache ByteString AuthResult
+  { jwtCacheIORef :: I.IORef (C.LRU ByteString AuthResult)
   }
 
 -- | Initialize JwtCacheState
-init :: IO JwtCacheState
-init = do
-  cache <- C.newCache Nothing -- no default expiration
-  return $ JwtCacheState cache
+init :: Int -> IO (Maybe JwtCacheState)
+init 0          = return Nothing
+init maxEntries = do
+  cache <- I.newIORef $ C.newLRU (Just $ integerFromInt maxEntries)
+  return $ Just $ JwtCacheState 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} 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
-
-      -- purge expired cache entries
-      C.purgeExpired jwtCache
-
-      -- insert new cache entry
-      C.insert' jwtCache (Just timeSpec) token res
-
-    _                    -> 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
+lookupJwtCache :: Maybe JwtCacheState -> ByteString -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult)
+lookupJwtCache Nothing _ parseJwt _ = parseJwt
+lookupJwtCache (Just JwtCacheState{jwtCacheIORef}) token parseJwt utc = do
+  -- get cache from IORef
+  jwtCache <- I.readIORef jwtCacheIORef
+
+  -- lookup key = token
+  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 res -> do
+
+            let jwtCache'' = C.insert token res jwtCache'
+
+            -- update IORef
+            I.writeIORef jwtCacheIORef jwtCache''
+
+            -- return the result
+            return $ Right res
+
+          Left e -> return $ Left e
+
+    Just res -> -- 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 res utc then do
+
+            let (jwtCache'',_) = C.delete token jwtCache'
+
+            I.writeIORef jwtCacheIORef jwtCache''
+
+            parseJwt
+
+        else do
+            
+            I.writeIORef jwtCacheIORef jwtCache'
+
+            return $ Right res
+
+
+type Expired = Bool
+
+-- | Check if exp claim is expired when looked up from cache
+isExpClaimExpired :: AuthResult -> UTCTime -> Expired
+isExpClaimExpired res 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 res)
+        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
+      |## Enables and set JWT Cache max entries, 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/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/Config/Database.hs

@@ -61,7 +61,7 @@ dbSettingsNames =
   ,"jwt_role_claim_key"
   ,"jwt_secret"
   ,"jwt_secret_is_base64"
-  ,"jwt_cache_max_lifetime"
+  ,"jwt-cache-max-entries"
   ,"openapi_mode"
   ,"openapi_security_active"
   ,"openapi_server_proxy_uri"

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"