feat: implement JWT Caching

CHANGELOG.md

@@ -10,6 +10,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
  - #1614, Add `db-pool-automatic-recovery` configuration to disable connection retrying - @taimoorzaeem
  - #2492, Allow full response control when raising exceptions - @taimoorzaeem, @laurenceisla
  - #2771, Add `Server-Timing` header with JWT duration - @taimoorzaeem
+ - #2698, Add config `jwt-cache-max-lifetime` and implement JWT caching - @taimoorzaeem
 
 ### Fixed
 

postgrest.cabal

@@ -80,8 +80,10 @@ library
                     , auto-update               >= 0.1.4 && < 0.2
                     , base64-bytestring         >= 1 && < 1.3
                     , bytestring                >= 0.10.8 && < 0.12
+                    , 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
                     , contravariant-extras      >= 0.3.3 && < 0.4

src/PostgREST/AppState.hs

@@ -4,6 +4,7 @@
 
 module PostgREST.AppState
   ( AppState
+  , AuthResult(..)
   , destroy
   , getConfig
   , getSchemaCache
@@ -12,6 +13,7 @@ module PostgREST.AppState
   , getPgVersion
   , getRetryNextIn
   , getTime
+  , getJwtCache
   , init
   , initWithPool
   , logWithZTime
@@ -24,8 +26,11 @@ module PostgREST.AppState
   , runListener
   ) where
 
+import qualified Data.Aeson                 as JSON
+import qualified Data.Aeson.KeyMap          as KM
 import qualified Data.ByteString.Char8      as BS
 import qualified Data.ByteString.Lazy       as LBS
+import qualified Data.Cache                 as C
 import           Data.Either.Combinators    (whenLeft)
 import qualified Data.Text.Encoding         as T
 import           Hasql.Connection           (acquire)
@@ -62,6 +67,11 @@ import PostgREST.SchemaCache.Identifiers (dumpQi)
 import Protolude
 
 
+data AuthResult = AuthResult
+  { authClaims :: KM.KeyMap JSON.Value
+  , authRole   :: BS.ByteString
+  }
+
 data AppState = AppState
   -- | Database connection pool
   { statePool                     :: SQL.Pool
@@ -87,6 +97,8 @@ data AppState = AppState
   , stateRetryNextIn              :: IORef Int
   -- | Logs a pool error with a debounce
   , debounceLogAcquisitionTimeout :: IO ()
+  -- | JWT Cache
+  , jwtCache                      :: C.Cache ByteString AuthResult
   }
 
 init :: AppConfig -> IO AppState
@@ -108,6 +120,7 @@ initWithPool pool conf = do
     <*> myThreadId
     <*> newIORef 0
     <*> pure (pure ())
+    <*> C.newCache Nothing
 
 
   debLogTimeout <-
@@ -188,6 +201,9 @@ putConfig = atomicWriteIORef . stateConf
 getTime :: AppState -> IO UTCTime
 getTime = stateGetTime
 
+getJwtCache :: AppState -> C.Cache ByteString AuthResult
+getJwtCache = jwtCache
+
 -- | Log to stderr with local time
 logWithZTime :: AppState -> Text -> IO ()
 logWithZTime appState txt = do

src/PostgREST/Auth.hs

@@ -26,6 +26,8 @@ import qualified Data.Aeson.KeyMap               as KM
 import qualified Data.Aeson.Types                as JSON
 import qualified Data.ByteString                 as BS
 import qualified Data.ByteString.Lazy.Char8      as LBS
+import qualified Data.Cache                      as C
+import qualified Data.Scientific                 as Sci
 import qualified Data.Vault.Lazy                 as Vault
 import qualified Data.Vector                     as V
 import qualified Network.HTTP.Types.Header       as HTTP
@@ -36,22 +38,20 @@ import Control.Lens            (set)
 import Control.Monad.Except    (liftEither)
 import Data.Either.Combinators (mapLeft)
 import Data.List               (lookup)
-import Data.Time.Clock         (UTCTime)
+import Data.Time.Clock         (UTCTime, nominalDiffTimeToSeconds)
+import Data.Time.Clock.POSIX   (utcTimeToPOSIXSeconds)
+import System.Clock            (TimeSpec (..))
 import System.IO.Unsafe        (unsafePerformIO)
 import System.TimeIt           (timeItT)
 
-import PostgREST.AppState (AppState, getConfig, getTime)
+import PostgREST.AppState (AppState, AuthResult (..), getConfig,
+                           getJwtCache, getTime)
 import PostgREST.Config   (AppConfig (..), JSPath, JSPathExp (..))
 import PostgREST.Error    (Error (..))
 
 import Protolude
 
 
-data AuthResult = AuthResult
-  { authClaims :: KM.KeyMap JSON.Value
-  , authRole   :: BS.ByteString
-  }
-
 -- | Receives the JWT secret and audience (from config) and a JWT and returns a
 -- JSON object of JWT claims.
 parseToken :: Monad m =>
@@ -107,16 +107,48 @@ middleware appState app req respond = do
   let token  = fromMaybe "" $ Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
       parseJwt = runExceptT $ parseToken conf (LBS.fromStrict token) time >>= parseClaims conf
 
-  if configDbPlanEnabled conf
-    then do
-      (dur,authResult) <- timeItT parseJwt
-      let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
-      app req' respond
-    else do
-      authResult <- parseJwt
-      let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
-      app req' respond
-
+-- If DbPlanEnabled       -> calculate JWT validation time
+-- If JwtCacheMaxLifetime -> cache JWT validation result
+  req' <- case (configDbPlanEnabled 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 $ getJWTFromCache appState token maxLifetime parseJwt time
+          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 <- getJWTFromCache appState token maxLifetime parseJwt time
+          return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
+
+  app req' respond
+
+-- | Used to retrieve and insert JWT to JWT Cache
+getJWTFromCache :: AppState -> ByteString -> Int -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult)
+getJWTFromCache appState token maxLifetime parseJwt utc = do
+  checkCache <- C.lookup (getJwtCache appState) token
+  authResult <- maybe parseJwt (pure . Right) checkCache
+
+  case (authResult,checkCache) of
+    (Right res, Nothing) -> C.insert' (getJwtCache appState) (getTimeSpec res maxLifetime utc) token res
+    _                    -> pure ()
+
+  return authResult
+
+-- Used to extract JWT exp claim and add to JWT Cache
+getTimeSpec :: AuthResult -> Int -> UTCTime -> Maybe 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) -> Just $ TimeSpec (sciToInt seconds - utcToSecs utc) 0
+    _                          -> Just $ TimeSpec (fromIntegral maxLifetime :: Int64) 0
 
 authResultKey :: Vault.Key (Either Error AuthResult)
 authResultKey = unsafePerformIO Vault.newKey

src/PostgREST/Config.hs

@@ -163,7 +163,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 . configJwtCaching)
+      ,("jwt-cache-max-lifetime",                   show . configJwtCacheMaxLifetime)
       ,("log-level",                 q . dumpLogLevel . configLogLevel)
       ,("openapi-mode",              q . dumpOpenApiMode . configOpenApiMode)
       ,("openapi-security-active",       T.toLower . show . configOpenApiSecurityActive)

test/io/test_io.py

@@ -1,6 +1,6 @@
 "Unit tests for Input/Ouput of PostgREST seen as a black box."
 
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
 from operator import attrgetter
 import os
 import re
@@ -1095,3 +1095,122 @@ def test_fail_with_automatic_recovery_disabled_and_terminated_using_query(defaul
 
         exitCode = wait_until_exit(postgrest)
         assert exitCode == 1
+
+
+def test_server_timing_jwt_should_decrease_on_subsequent_requests(defaultenv):
+    "assert that server-timing duration for JWT should decrease on subsequent requests"
+
+    env = {
+        **defaultenv,
+        "PGRST_DB_PLAN_ENABLED": "true",
+        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_SECRET": "@/dev/stdin",
+        "PGRST_DB_CONFIG": "false",
+    }
+
+    headers = jwtauthheader(
+        {
+            "role": "postgrest_test_author",
+            "exp": int(
+                (datetime.now(timezone.utc) + timedelta(minutes=30)).timestamp()
+            ),
+        },
+        SECRET,
+    )
+
+    with run(stdin=SECRET.encode(), env=env) as postgrest:
+        first_dur_text = postgrest.session.get(
+            "/authors_only", headers=headers
+        ).headers["Server-Timing"]
+        second_dur_text = postgrest.session.get(
+            "/authors_only", headers=headers
+        ).headers["Server-Timing"]
+
+        first_dur = float(first_dur_text[8:])  # skip "jwt;dur="
+        second_dur = float(second_dur_text[8:])
+
+        # their difference should be atleast 300, implying
+        # that JWT Caching is working as expected
+        assert (first_dur - second_dur) > 300.0
+
+
+# just added to complete code coverage
+def test_jwt_caching_works_with_db_plan_disabled(defaultenv):
+    "assert that JWT caching words even when Server-Timing header is not returned"
+
+    env = {
+        **defaultenv,
+        "PGRST_DB_PLAN_ENABLED": "false",
+        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_SECRET": "@/dev/stdin",
+        "PGRST_DB_CONFIG": "false",
+    }
+
+    headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET)
+
+    with run(stdin=SECRET.encode(), env=env) as postgrest:
+        first_request = postgrest.session.get("/authors_only", headers=headers)
+        second_request = postgrest.session.get("/authors_only", headers=headers)
+
+        # in this case we don't get server-timing in response headers
+        # so we can't compare durations, we just check if request succeeds
+        assert first_request.status_code == 200 and second_request.status_code == 200
+
+
+def test_server_timing_jwt_should_not_decrease_when_caching_disabled(defaultenv):
+    "assert than jwt duration should not decrease when disabled"
+
+    env = {
+        **defaultenv,
+        "PGRST_DB_PLAN_ENABLED": "true",
+        "PGRST_JWT_CACHE_MAX_LIFETIME": "0",  # cache disabled
+        "PGRST_JWT_SECRET": "@/dev/stdin",
+        "PGRST_DB_CONFIG": "false",
+    }
+
+    headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET)
+
+    with run(stdin=SECRET.encode(), env=env) as postgrest:
+        warmup_req = postgrest.session.get("/authors_only", headers=headers)
+        first_dur_text = postgrest.session.get(
+            "/authors_only", headers=headers
+        ).headers["Server-Timing"]
+        second_dur_text = postgrest.session.get(
+            "/authors_only", headers=headers
+        ).headers["Server-Timing"]
+
+        first_dur = float(first_dur_text[8:])  # skip "jwt;dur="
+        second_dur = float(second_dur_text[8:])
+
+        # their difference should be less than 100
+        # implying that token is not cached
+        assert (first_dur - second_dur) < 100.0
+
+
+def test_jwt_cache_with_no_exp_claim(defaultenv):
+    "assert than jwt duration should decrease"
+
+    env = {
+        **defaultenv,
+        "PGRST_DB_PLAN_ENABLED": "true",
+        "PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
+        "PGRST_JWT_SECRET": "@/dev/stdin",
+        "PGRST_DB_CONFIG": "false",
+    }
+
+    headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET)  # no exp
+
+    with run(stdin=SECRET.encode(), env=env) as postgrest:
+        first_dur_text = postgrest.session.get(
+            "/authors_only", headers=headers
+        ).headers["Server-Timing"]
+        second_dur_text = postgrest.session.get(
+            "/authors_only", headers=headers
+        ).headers["Server-Timing"]
+
+        first_dur = float(first_dur_text[8:])  # skip "jwt;dur="
+        second_dur = float(second_dur_text[8:])
+
+        # their difference should be less than 100
+        # implying that token is not cached
+        assert (first_dur - second_dur) > 300.0

test/memory/memory-tests.sh

@@ -102,7 +102,7 @@ postJsonArrayTest(){
 
 echo "Running memory usage tests.."
 
-jsonKeyTest "1M" "POST" "/rpc/leak?columns=blob" "16M"
+jsonKeyTest "1M" "POST" "/rpc/leak?columns=blob" "23M"
 jsonKeyTest "1M" "POST" "/leak?columns=blob" "16M"
 jsonKeyTest "1M" "PATCH" "/leak?id=eq.1&columns=blob" "16M"