feat: implement JWT Caching

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
@@ -184,6 +186,7 @@ test-suite spec
                       Feature.Auth.AudienceJwtSecretSpec
                       Feature.Auth.AuthSpec
                       Feature.Auth.BinaryJwtSecretSpec
+                      Feature.Auth.JwtCachingSpec
                       Feature.Auth.NoAnonSpec
                       Feature.Auth.NoJwtSpec
                       Feature.ConcurrentSpec

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
@@ -37,21 +39,18 @@ import Control.Monad.Except    (liftEither)
 import Data.Either.Combinators (mapLeft)
 import Data.List               (lookup)
 import Data.Time.Clock         (UTCTime)
+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 +106,49 @@ 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 JwtCaching    -> cache JWT validation result
+  case (configDbPlanEnabled conf, configJwtCaching conf) of
+    (True, True)   -> do
+          (dur, authResult) <- timeItT $ getJWTFromCache appState token parseJwt
+          let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
+          app req' respond
+
+    (True, False)  -> do
+          (dur, authResult) <- timeItT parseJwt
+          let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
+          app req' respond
+
+    (False, True)  -> do
+          authResult <- getJWTFromCache appState token parseJwt
+          let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
+          app req' respond
+
+    (False, False) -> do
+          authResult <- parseJwt
+          let req' = req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
+          app req' respond
+
+-- Used to extract JWT exp claim and add to JWT Cache
+getTimeSpec :: AuthResult -> Maybe TimeSpec
+getTimeSpec res = do
+  let sciToInt = fromMaybe 0 . Sci.toBoundedInteger
+  expireJSON <- KM.lookup "exp" (authClaims res)
+  case expireJSON of
+    JSON.Number seconds -> Just $ TimeSpec (sciToInt seconds) 0
+    _                   -> Just $ TimeSpec 0 0 -- set timeSpec to 0 so it expires immediately, hence not cached
+
+-- | Used to retrieve and insert JWT to JWT Cache
+getJWTFromCache :: AppState -> ByteString -> IO (Either Error AuthResult) -> IO (Either Error AuthResult)
+getJWTFromCache appState token parseJwt = 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) token res
+    _                    -> pure ()
+
+  return authResult
 
 authResultKey :: Vault.Key (Either Error AuthResult)
 authResultKey = unsafePerformIO Vault.newKey

test/spec/Feature/Auth/JwtCachingSpec.hs

@@ -0,0 +1,39 @@
+module Feature.Auth.JwtCachingSpec where
+
+import Network.Wai (Application)
+
+import Network.HTTP.Types
+import Test.Hspec
+import Test.Hspec.Wai
+import Test.Hspec.Wai.JSON
+
+import Protolude  hiding (get)
+import SpecHelper
+
+spec :: SpecWith ((), Application)
+spec = describe "JWT Caching" $ do
+  let auth = authHeaderJWT "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJyb2xlIjoicG9zdGdyZXN0X3Rlc3RfYXV0aG9yIn0.Xod-F15qsGL0WhdOCr2j3DdKuTw9QJERVgoFD3vGaWA"
+
+  it "jwt Cache warmup" $
+    request methodPost "/rpc/privileged_hello" [auth] [json|{"name": "jdoe"}|]
+      `shouldRespondWith` [json|"Privileged hello to jdoe"|]
+      { matchStatus = 200
+      , matchHeaders = ["Server-Timing" <:> "jwt;dur="
+                       , matchContentTypeJson]
+      }
+
+  it "JWT Cache, token should be cached at this request" $
+    request methodPost "/rpc/privileged_hello" [auth] [json|{"name": "jdoe"}|]
+      `shouldRespondWith` [json|"Privileged hello to jdoe"|]
+      { matchStatus = 200
+      , matchHeaders = ["Server-Timing" <:> "jwt;dur="
+                       , matchContentTypeJson]
+      }
+
+  it "JWT Claims should be retrieved from the cache" $
+    request methodPost "/rpc/privileged_hello" [auth] [json|{"name": "jdoe"}|]
+      `shouldRespondWith` [json|"Privileged hello to jdoe"|]
+      { matchStatus = 200
+      , matchHeaders = ["Server-Timing" <:> "jwt;dur="
+                       , matchContentTypeJson]
+      }

test/spec/Main.hs

@@ -20,6 +20,7 @@ import qualified Feature.Auth.AsymmetricJwtSpec
 import qualified Feature.Auth.AudienceJwtSecretSpec
 import qualified Feature.Auth.AuthSpec
 import qualified Feature.Auth.BinaryJwtSecretSpec
+import qualified Feature.Auth.JwtCachingSpec
 import qualified Feature.Auth.NoAnonSpec
 import qualified Feature.Auth.NoJwtSpec
 import qualified Feature.ConcurrentSpec
@@ -111,6 +112,7 @@ main = do
       pgSafeUpdateApp      = app testPgSafeUpdateEnabledCfg
       obsApp               = app testObservabilityCfg
       serverTiming         = app testCfgServerTiming
+      jwtCaching           = app testCfgJwtCaching
 
       extraSearchPathApp   = appDbs testCfgExtraSearchPath
       unicodeApp           = appDbs testUnicodeCfg
@@ -210,6 +212,10 @@ main = do
     parallel $ before asymJwkSetApp $
       describe "Feature.Auth.AsymmetricJwtSpec" Feature.Auth.AsymmetricJwtSpec.spec
 
+    -- this test runs with jwtCaching
+    parallel $ before jwtCaching $
+      describe "Feature.Auth.JwtCachingSpec" Feature.Auth.JwtCachingSpec.spec
+
     -- this test runs with a nonexistent db-schema
     parallel $ before nonexistentSchemaApp $
       describe "Feature.Query.ErrorSpec" Feature.Query.ErrorSpec.spec