Initial version of Sieve JWT Caching

Author: mkleczek

postgrest.cabal

@@ -49,8 +49,10 @@ library
                       PostgREST.App
                       PostgREST.AppState
                       PostgREST.Auth
+                      PostgREST.Auth.Jwt
                       PostgREST.Auth.JwtCache
                       PostgREST.Auth.Types
+                      PostgREST.Cache.Sieve
                       PostgREST.CLI
                       PostgREST.Config
                       PostgREST.Config.Database
@@ -153,6 +155,9 @@ library
                     -- https://github.com/kazu-yamamoto/logger/commit/3a71ca70afdbb93d4ecf0083eeba1fbbbcab3fc3
                     , wai-logger                >= 2.4.0
                     , warp                      >= 3.3.19 && < 3.4
+                    , stm                       >= 2.5 && < 3
+                    , stm-hamt                  >= 1.2 && < 2
+                    , focus                     >= 1.0 && < 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

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,7 +127,7 @@ init conf@AppConfig{configLogLevel, configDbPoolSize} = do
 
   observer $ AppStartObs prettyVersion
 
-  jwtCacheState <- JwtCache.init
+  jwtCacheState <- JwtCache.init conf
   pool <- initPool conf observer
   (sock, adminSock) <- initSockets conf
   state' <- initWithPool (sock, adminSock) pool conf jwtCacheState loggerState metricsState observer
@@ -471,10 +471,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 newConf $ getJwtCacheState appState
 
       if startingUp then
         pass

src/PostgREST/Auth.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE RecordWildCards #-}
 {-|
 Module      : PostgREST.Auth
 Description : PostgREST authentication functions.
@@ -10,186 +11,52 @@ Authentication should always be implemented in an external service.
 In the test suite there is an example of simple login function that can be used for a
 very simple authentication system inside the PostgreSQL database.
 -}
-{-# LANGUAGE LambdaCase      #-}
-{-# LANGUAGE RecordWildCards #-}
+
 module PostgREST.Auth
   ( getResult
   , getJwtDur
   , getRole
   , middleware
   ) where
 
-import qualified Data.Aeson                      as JSON
-import qualified Data.Aeson.Key                  as K
-import qualified Data.Aeson.KeyMap               as KM
-import qualified Data.Aeson.Types                as JSON
 import qualified Data.ByteString                 as BS
-import qualified Data.ByteString.Internal        as BS
-import qualified Data.ByteString.Lazy.Char8      as LBS
-import qualified Data.Scientific                 as Sci
-import qualified Data.Text                       as T
 import qualified Data.Vault.Lazy                 as Vault
-import qualified Data.Vector                     as V
-import qualified Jose.Jwk                        as JWT
-import qualified Jose.Jwt                        as JWT
 import qualified Network.HTTP.Types.Header       as HTTP
 import qualified Network.Wai                     as Wai
 import qualified Network.Wai.Middleware.HttpAuth as Wai
 
-import Control.Monad.Except    (liftEither)
-import Data.Either.Combinators (mapLeft)
 import Data.List               (lookup)
-import Data.Time.Clock         (UTCTime, nominalDiffTimeToSeconds)
-import Data.Time.Clock.POSIX   (utcTimeToPOSIXSeconds)
 import System.IO.Unsafe        (unsafePerformIO)
 import System.TimeIt           (timeItT)
 
 import PostgREST.AppState      (AppState, getConfig, getJwtCacheState,
                                 getTime)
 import PostgREST.Auth.JwtCache (lookupJwtCache)
+import PostgREST.Auth.Jwt      (parseClaims)
 import PostgREST.Auth.Types    (AuthResult (..))
-import PostgREST.Config        (AppConfig (..), FilterExp (..),
-                                JSPath, JSPathExp (..))
-import PostgREST.Error         (Error (..), JwtError (..))
+import PostgREST.Config        (AppConfig (..))
+import PostgREST.Error         (Error (..))
 
 import Protolude
 
--- | Receives the JWT secret and audience (from config) and a JWT and returns a
--- JSON object of JWT claims.
-parseToken :: AppConfig -> Maybe ByteString -> UTCTime -> ExceptT Error IO JSON.Value
-parseToken _ Nothing _ = return JSON.emptyObject
-parseToken _ (Just "") _ = throwE . JwtErr $ JwtDecodeError "Empty JWT is sent in Authorization header"
-parseToken AppConfig{..} (Just tkn) time = do
-  secret <- liftEither . maybeToRight (JwtErr JwtSecretMissing) $ configJWKS
-  tknWith3Parts <- liftEither $ hasThreeParts tkn
-  eitherContent <- liftIO $ JWT.decode (JWT.keys secret) Nothing tknWith3Parts
-  content <- liftEither . mapLeft (JwtErr . jwtDecodeError) $ eitherContent
-  liftEither $ mapLeft JwtErr $ verifyClaims content
-  where
-      hasThreeParts :: ByteString -> Either Error ByteString
-      hasThreeParts token = case length $ BS.split (BS.c2w '.') token of
-        3 -> Right token
-        n -> Left $ JwtErr $ JwtDecodeError ("Expected 3 parts in JWT; got " <> show n)
-      jwtDecodeError :: JWT.JwtError -> JwtError
-      -- The only errors we can get from JWT.decode function are:
-      --   BadAlgorithm
-      --   KeyError
-      --   BadCrypto
-      jwtDecodeError (JWT.KeyError _)     = JwtDecodeError "No suitable key or wrong key type"
-      jwtDecodeError (JWT.BadAlgorithm _) = JwtDecodeError "Wrong or unsupported encoding algorithm"
-      jwtDecodeError JWT.BadCrypto        = JwtDecodeError "JWT cryptographic operation failed"
-      -- Control never reaches here, the decode function only returns the above three
-      jwtDecodeError _                    = JwtDecodeError "JWT couldn't be decoded"
-
-      verifyClaims :: JWT.JwtContent -> Either JwtError JSON.Value
-      verifyClaims (JWT.Jws (_, claims)) = case JSON.decodeStrict claims of
-        Just jclaims@(JSON.Object mclaims) ->
-          verifyClaim mclaims "exp" isValidExpClaim "JWT expired" >>
-          verifyClaim mclaims "nbf" isValidNbfClaim "JWT not yet valid" >>
-          verifyClaim mclaims "iat" isValidIatClaim "JWT issued at future" >>
-          verifyClaim mclaims "aud" isValidAudClaim "JWT not in audience" >>
-          return jclaims
-        _ -> Left $ JwtClaimsError "Parsing claims failed"
-      -- TODO: We could enable JWE support here (encrypted tokens)
-      verifyClaims _ = Left $ JwtDecodeError "Unsupported token type"
-
-      verifyClaim mclaims claim func err = do
-        isValid <- maybe (Right True) func (KM.lookup claim mclaims)
-        unless isValid $ Left $ JwtClaimsError err
-
-      allowedSkewSeconds = 30 :: Int64
-      now = floor . nominalDiffTimeToSeconds $ utcTimeToPOSIXSeconds time
-      sciToInt = fromMaybe 0 . Sci.toBoundedInteger
-      allStrings = all (\case (JSON.String _) -> True; _ -> False)
-
-      isValidExpClaim :: JSON.Value -> Either JwtError Bool
-      isValidExpClaim (JSON.Number secs) = Right $ now <= (sciToInt secs + allowedSkewSeconds)
-      isValidExpClaim _ = Left $ JwtClaimsError "The JWT 'exp' claim must be a number"
-
-      isValidNbfClaim :: JSON.Value -> Either JwtError Bool
-      isValidNbfClaim (JSON.Number secs) = Right $ now >= (sciToInt secs - allowedSkewSeconds)
-      isValidNbfClaim _ = Left $ JwtClaimsError "The JWT 'nbf' claim must be a number"
-
-      isValidIatClaim :: JSON.Value -> Either JwtError Bool
-      isValidIatClaim (JSON.Number secs) = Right $ now >= (sciToInt secs - allowedSkewSeconds)
-      isValidIatClaim _ = Left $ JwtClaimsError "The JWT 'iat' claim must be a number"
-
-      isValidAudClaim :: JSON.Value -> Either JwtError Bool
-      isValidAudClaim JSON.Null = Right True -- {"aud": null} is valid for all audiences
-      isValidAudClaim (JSON.String str) = Right $ maybe (const True) (==) configJwtAudience str
-      isValidAudClaim (JSON.Array arr)
-        | null arr = Right True              -- {"aud": []} is valid for all audiences
-        | allStrings arr = Right $ maybe True (\a -> JSON.String a `elem` arr) configJwtAudience
-      isValidAudClaim _ = Left $ JwtClaimsError "The JWT 'aud' claim must be a string or an array of strings"
-
-parseClaims :: Monad m =>
-  AppConfig -> JSON.Value -> ExceptT Error m AuthResult
-parseClaims AppConfig{..} jclaims@(JSON.Object mclaims) = do
-  -- role defaults to anon if not specified in jwt
-  role <- liftEither . maybeToRight (JwtErr JwtTokenRequired) $
-    unquoted <$> walkJSPath (Just jclaims) configJwtRoleClaimKey <|> configDbAnonRole
-  return AuthResult
-           { authClaims = mclaims & KM.insert "role" (JSON.toJSON $ decodeUtf8 role)
-           , authRole = role
-           }
-  where
-    walkJSPath :: Maybe JSON.Value -> JSPath -> Maybe JSON.Value
-    walkJSPath x                      []                = x
-    walkJSPath (Just (JSON.Object o)) (JSPKey key:rest) = walkJSPath (KM.lookup (K.fromText key) o) rest
-    walkJSPath (Just (JSON.Array ar)) (JSPIdx idx:rest) = walkJSPath (ar V.!? idx) rest
-    walkJSPath (Just (JSON.Array ar)) [JSPFilter (EqualsCond txt)] = findFirstMatch (==) txt ar
-    walkJSPath (Just (JSON.Array ar)) [JSPFilter (NotEqualsCond txt)] = findFirstMatch (/=) txt ar
-    walkJSPath (Just (JSON.Array ar)) [JSPFilter (StartsWithCond txt)] = findFirstMatch T.isPrefixOf txt ar
-    walkJSPath (Just (JSON.Array ar)) [JSPFilter (EndsWithCond txt)] = findFirstMatch T.isSuffixOf txt ar
-    walkJSPath (Just (JSON.Array ar)) [JSPFilter (ContainsCond txt)] = findFirstMatch T.isInfixOf txt ar
-    walkJSPath _                      _                 = Nothing
-
-    findFirstMatch matchWith pattern = foldr checkMatch Nothing
-      where
-        checkMatch (JSON.String txt) acc
-            | pattern `matchWith` txt = Just $ JSON.String txt
-            | otherwise = acc
-        checkMatch _ acc = acc
-
-    unquoted :: JSON.Value -> BS.ByteString
-    unquoted (JSON.String t) = encodeUtf8 t
-    unquoted v               = LBS.toStrict $ JSON.encode v
--- impossible case - just added to please -Wincomplete-patterns
-parseClaims _ _ = return AuthResult { authClaims = KM.empty, authRole = mempty }
-
--- | 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
-  conf <- getConfig appState
+  conf@AppConfig{..} <- getConfig appState
   time <- getTime appState
 
   let token  = Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
-      parseJwt = runExceptT $ parseToken conf token time >>= parseClaims conf
+      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 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 }
-
-    (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