add: config jwt-schema-claim-key for schema selection in JWT

It follows the same JSPath grammar as `jwt-role-claim-key`. If
the schema is specified in JWT claims, it overides the `Accept-Profile`
and `Content-Profile` headers.

Signed-off-by: Taimoor Zaeem <taimoorzaeem@gmail.com>

Author: taimoorzaeem

CHANGELOG.md

@@ -14,6 +14,7 @@ All notable changes to this project will be documented in this file. From versio
 - Add `Vary` header to responses by @develop7 in #4609
 - Add config `db-timezone-enabled` for optional querying of timezones by @taimoorzaeem in #4751
 - Log when the pool is released during schema cache reload on `log-level=debug` by @mkleczek in #4668
+- Add config `jwt-schema-claim-key` for schema selection in JWT by @taimoorzaeem in #4608
 
 ### Changed
 

docs/postgrest.dict

@@ -92,6 +92,7 @@ lte
 macOS
 misprediction
 multi
+multitenancy
 namespace
 namespaced
 Nanos
@@ -200,4 +201,4 @@ webuser
 wfts
 www
 debouncing
-deduplicates
+deduplicates

docs/references/api/schemas.rst

@@ -41,6 +41,8 @@ Only the selected schema gets added to the `search_path <https://www.postgresql.
 
    These headers are based on the "Content Negotiation by Profile" spec: https://www.w3.org/TR/dx-prof-conneg
 
+.. _profile_headers:
+
 GET/HEAD
 ~~~~~~~~
 

docs/references/auth.rst

@@ -226,6 +226,33 @@ JWT Role Extraction
 
 A JSPath DSL that specifies the location of the :code:`role` key in the JWT claims. It's configured by :ref:`jwt-role-claim-key`. This can be used to consume a JWT provided by a third party service like Auth0, Okta, Microsoft Entra or Keycloak.
 
+.. code::
+
+    # {"postgrest":{"roles": ["other", "author"]}}
+    # the DSL accepts characters that are alphanumerical or one of "_$@" as keys
+    jwt-role-claim-key = ".postgrest.roles[1]"
+
+See :ref:`jspath_dsl_grammar` for more details on how to specify the location.
+
+.. _jwt_schema_extract:
+
+JWT Schema Extraction
+---------------------
+
+Schema can be specified in JWT claims. It is configured by :ref:`jwt-schema-claim-key`. This feature can be used for JWT-driven schema-based multitenancy. It allows fully hidden schema selection without exposing the schema in :ref:`profile headers <profile_headers>`. The schema specified in JWT takes precedence over profile headers.
+
+.. code::
+
+    # {"postgrest":{"schema": "jwt-tenant"}}
+    jwt-schema-claim-key = ".postgrest.schema"
+
+See :ref:`jspath_dsl_grammar` for more details on how to specify the location.
+
+.. _jspath_dsl_grammar:
+
+JSPath DSL Grammar
+~~~~~~~~~~~~~~~~~~
+
 The DSL follows the `JSONPath <https://goessner.net/articles/JsonPath/>`_ expression grammar with extended string comparison operators. Supported operators are:
 
 - ``==`` selects the first array element that exactly matches the right operand
@@ -234,7 +261,7 @@ The DSL follows the `JSONPath <https://goessner.net/articles/JsonPath/>`_ expres
 - ``==^`` selects the first array element that ends with the right operand
 - ``*==`` selects the first array element that contains the right operand
 
-The selected role value can also be sliced using the slice operator ``[a:b]``. It is similar to `slice operator in python <https://docs.python.org/3/library/functions.html#slice>`_. Negative index values are also supported. The syntax is as:
+The selected value can also be sliced using the slice operator ``[a:b]``. It is similar to `slice operator in python <https://docs.python.org/3/library/functions.html#slice>`_. Negative index values are also supported. The syntax is as:
 
 - ``[a:b]`` take slice from index ``a`` up to ``b``
 - ``[a:]`` take slice from index ``a`` to end
@@ -249,10 +276,6 @@ Usage examples:
 
   .. code:: bash
 
-    # {"postgrest":{"roles": ["other", "author"]}}
-    # the DSL accepts characters that are alphanumerical or one of "_$@" as keys
-    jwt-role-claim-key = ".postgrest.roles[1]"
-
     # {"https://www.example.com/role": { "key": "author" }}
     # non-alphanumerical characters can go inside quotes(escaped in the config value)
     jwt-role-claim-key = ".\"https://www.example.com/role\".key"
@@ -271,6 +294,8 @@ Usage examples:
     jwt-role-claim-key = ".postgrest.wlcg[0][1:]"
     jwt-role-claim-key = ".postgrest.wlcg[1][1:-1]"
 
+These examples also apply to :ref:`jwt-schema-claim-key`.
+
 .. note::
 
   The string comparison operators are implemented as a custom extension to the JSPath and does not strictly follow the `RFC 9535 <https://www.rfc-editor.org/rfc/rfc9535.html>`_.

docs/references/configuration.rst

@@ -664,6 +664,21 @@ jwt-role-claim-key
 
   See :ref:`jwt_role_extract` on how to specify key paths and usage examples.
 
+.. _jwt-schema-claim-key:
+
+jwt-schema-claim-key
+--------------------
+
+  =============== =================================
+  **Type**        String
+  **Default**     .schema
+  **Reloadable**  Y
+  **Environment** PGRST_JWT_SCHEMA_CLAIM_KEY
+  **In-Database** pgrst.jwt_schema_claim_key
+  =============== =================================
+
+  See :ref:`jwt_schema_extract` on how to specify key paths and usage examples.
+
 .. _jwt-secret:
 
 jwt-secret

src/PostgREST/ApiRequest.hs

@@ -30,6 +30,7 @@ import PostgREST.ApiRequest.Types        (Action (..), DbAction (..),
                                           InvokeMethod (..),
                                           Mutation (..), Payload (..),
                                           RequestBody, Resource (..))
+import PostgREST.Auth.Types              (AuthResult (..))
 import PostgREST.Config                  (AppConfig (..),
                                           OpenAPIMode (..))
 import PostgREST.Config.Database         (TimezoneNames)
@@ -76,10 +77,10 @@ data ApiRequest = ApiRequest {
   }
 
 -- | Examines HTTP request and translates it into user intent.
-userApiRequest :: AppConfig -> Preferences.Preferences -> Request -> RequestBody -> Either ApiRequestError ApiRequest
-userApiRequest conf prefs req reqBody = do
+userApiRequest :: AppConfig -> Preferences.Preferences -> Request -> AuthResult -> RequestBody -> Either ApiRequestError ApiRequest
+userApiRequest conf prefs req authResult reqBody = do
   resource <- getResource conf $ pathInfo req
-  (schema, negotiatedByProfile) <- getSchema conf hdrs method
+  (schema, negotiatedByProfile) <- getSchema conf hdrs method authResult
   act <- getAction resource schema method
   qPrms <- first QueryParamError $ QueryParams.parse (actIsInvokeSafe act) $ rawQueryString req
   (topLevelRange, ranges) <- getRanges method qPrms hdrs
@@ -151,14 +152,25 @@ getAction resource schema method =
   where
     qi = QualifiedIdentifier schema
 
-
-getSchema :: AppConfig -> RequestHeaders -> ByteString -> Either ApiRequestError (Schema, Bool)
-getSchema AppConfig{configDbSchemas} hdrs method = do
-  case profile of
-    Just p | p `notElem` configDbSchemas -> Left $ UnacceptableSchema p $ toList configDbSchemas
-           | otherwise                   -> Right (p, True)
-    Nothing -> Right (defaultSchema, length configDbSchemas /= 1) -- if we have many schemas, assume the default schema was negotiated
+-- |
+-- We get schema in the following order:
+-- 1. Check JWT for a schema claim
+-- 2. If no schema claim, then check "Accept-Profile" and "Content-Profile" headers
+-- 3. If profile headers not sent, then default to first schema in "db-schemas"
+getSchema :: AppConfig -> RequestHeaders -> ByteString -> AuthResult -> Either ApiRequestError (Schema, Bool)
+getSchema AppConfig{configDbSchemas} hdrs method AuthResult{authSchema} = do
+  case authSchema of
+    Just s -> checkSchemaAcceptable (decodeUtf8 s) False
+    Nothing  ->
+      case profile of
+        Just p  -> checkSchemaAcceptable p True
+        Nothing -> Right (defaultSchema, length configDbSchemas /= 1) -- if we have many schemas, assume the default schema was negotiated
   where
+    checkSchemaAcceptable :: Text -> Bool -> Either ApiRequestError (Schema,Bool)
+    checkSchemaAcceptable schema isNegotiated
+      | schema `notElem` configDbSchemas = Left $ UnacceptableSchema schema $ toList configDbSchemas
+      | otherwise = Right (schema, isNegotiated)
+
     defaultSchema = NonEmptyList.head configDbSchemas
     profile = case method of
       -- POST/PATCH/PUT/DELETE don't use the same header as per the spec

src/PostgREST/App.hs

@@ -177,7 +177,7 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache authResult@AuthRe
       timezones = dbTimezones sCache
       prefs = ApiRequest.userPreferences conf req timezones
 
-  (parseTime, apiReq@ApiRequest{..}) <- withTiming $ liftEither . mapLeft Error.ApiRequestErr $ ApiRequest.userApiRequest conf prefs req body
+  (parseTime, apiReq@ApiRequest{..}) <- withTiming $ liftEither . mapLeft Error.ApiRequestErr $ ApiRequest.userApiRequest conf prefs req authResult body
   (planTime, plan)                   <- withTiming $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
 
   let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest

src/PostgREST/Auth/Jwt.hs

@@ -8,8 +8,8 @@ This module provides functions to deal with JWT parsing and validation (http://j
 {-# LANGUAGE FlexibleContexts      #-}
 {-# LANGUAGE ImpredicativeTypes    #-}
 {-# LANGUAGE LambdaCase            #-}
-{-# LANGUAGE NamedFieldPuns        #-}
 {-# LANGUAGE QuantifiedConstraints #-}
+{-# LANGUAGE RecordWildCards       #-}
 
 module PostgREST.Auth.Jwt
   ( parseAndDecodeClaims
@@ -110,14 +110,16 @@ parseToken secret tkn = do
       jwtDecodeError _                    = JwtDecodeErr UnreachableDecodeError
 
 parseClaims :: (MonadError Error m, MonadIO m) => AppConfig -> UTCTime -> JSON.Object -> m AuthResult
-parseClaims cfg@AppConfig{configJwtRoleClaimKey, configDbAnonRole} time mclaims = do
+parseClaims cfg@AppConfig{..} time mclaims = do
   validateClaims time (audMatchesCfg cfg) mclaims
   -- role defaults to anon if not specified in jwt
   role <- liftEither . maybeToRight (JwtErr JwtTokenRequired) $
     unquoted <$> walkJSPath (Just $ JSON.Object mclaims) configJwtRoleClaimKey <|> configDbAnonRole
+  let schema = unquoted <$> walkJSPath (Just $ JSON.Object mclaims) configJwtSchemaClaimKey
   pure AuthResult
            { authClaims = mclaims
            , authRole = role
+           , authSchema = schema
            }
   where
     unquoted :: JSON.Value -> BS.ByteString

src/PostgREST/Auth/Types.hs

@@ -6,10 +6,13 @@ import qualified Data.Aeson        as JSON
 import qualified Data.Aeson.KeyMap as KM
 import qualified Data.ByteString   as BS
 
+import Protolude
+
 -- |
 -- Parse and store result for JWT Claims. Can be accessed in
 -- db through GUCs (for RLS etc)
 data AuthResult = AuthResult
   { authClaims :: KM.KeyMap JSON.Value
   , authRole   :: BS.ByteString
+  , authSchema :: Maybe BS.ByteString
   }

src/PostgREST/Config.hs

@@ -61,7 +61,8 @@ import PostgREST.Config.Database         (RoleIsolationLvl,
                                           RoleSettings)
 import PostgREST.Config.JSPath           (FilterExp (..), JSPath,
                                           JSPathExp (..), dumpJSPath,
-                                          pRoleClaimKey)
+                                          pRoleClaimKey,
+                                          pSchemaClaimKey)
 import PostgREST.Config.Proxy            (Proxy (..),
                                           isMalformedProxyUri, toURI)
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..),
@@ -102,6 +103,7 @@ data AppConfig = AppConfig
   , configJWKS                     :: Maybe JwkSet
   , configJwtAudience              :: Maybe Text
   , configJwtRoleClaimKey          :: JSPath
+  , configJwtSchemaClaimKey        :: JSPath
   , configJwtSecret                :: Maybe BS.ByteString
   , configJwtSecretIsBase64        :: Bool
   , configJwtCacheMaxEntries       :: Int
@@ -187,6 +189,7 @@ toText conf =
       ,("db-uri",                    q . configDbUri)
       ,("jwt-aud",                   q . fromMaybe mempty . configJwtAudience)
       ,("jwt-role-claim-key",        q . T.intercalate mempty . fmap dumpJSPath . configJwtRoleClaimKey)
+      ,("jwt-schema-claim-key",      q . T.intercalate mempty . fmap dumpJSPath . configJwtSchemaClaimKey)
       ,("jwt-secret",                q . T.decodeUtf8 . showJwtSecret)
       ,("jwt-secret-is-base64",          T.toLower . show . configJwtSecretIsBase64)
       ,("jwt-cache-max-entries",         show . configJwtCacheMaxEntries)
@@ -300,6 +303,7 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
     <*> pure Nothing
     <*> optStringOrURI "jwt-aud"
     <*> parseRoleClaimKey "jwt-role-claim-key" "role-claim-key"
+    <*> parseSchemaClaimKey "jwt-schema-claim-key"
     <*> (fmap encodeUtf8 <$> optString "jwt-secret")
     <*> (fromMaybe False <$> optWithAlias
           (optBool "jwt-secret-is-base64")
@@ -421,6 +425,12 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
         Nothing  -> pure [JSPKey "role"]
         Just rck -> either (fail . show) pure $ pRoleClaimKey rck
 
+    parseSchemaClaimKey :: C.Key -> C.Parser C.Config JSPath
+    parseSchemaClaimKey k =
+      optString k >>= \case
+        Nothing  -> pure [JSPKey "schema"]
+        Just sck -> either (fail . show) pure $ pSchemaClaimKey sck
+
     parseCORSAllowedOrigins k =
       optString k >>= \case
         Nothing   -> pure Nothing
@@ -741,6 +751,9 @@ exampleConfigFile = S.unlines
   , ""
   , "## Jspath to the role claim key"
   ,  "jwt-role-claim-key = \".role\""
+  , ""
+  , "## Jspath to the schema claim key"
+  ,  "jwt-schema-claim-key = \".schema\""
   ,  ""
   ,  "## Choose a secret, JSON Web Key (or set) to enable JWT auth"
   ,  "## (use \"@filename\" to load from separate file)"