Add role-claim-key config value

src/PostgREST/Auth.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
 {-|
 Module      : PostgREST.Auth
 Description : PostgREST authorization functions.
@@ -19,9 +20,11 @@ module PostgREST.Auth (
   ) where
 
 import           Control.Lens.Operators
-import           Data.Aeson             (Value (..), decode, toJSON)
+import qualified Data.Aeson             as JSON
 import qualified Data.HashMap.Strict    as M
-import           Data.Time.Clock           (UTCTime)
+import           Data.Time.Clock        (UTCTime)
+import           Data.Vector            as V
+import           PostgREST.Types
 import           Protolude
 
 import qualified Crypto.JOSE.Types      as JOSE.Types
@@ -32,16 +35,16 @@ import           Crypto.JWT
 -}
 data JWTAttempt = JWTInvalid JWTError
                 | JWTMissingSecret
-                | JWTClaims (M.HashMap Text Value)
+                | JWTClaims (M.HashMap Text JSON.Value)
                 deriving (Eq, Show)
 
 {-|
   Receives the JWT secret and audience (from config) and a JWT and returns a map
   of JWT claims.
 -}
-jwtClaims :: Maybe JWK -> Maybe StringOrURI -> LByteString -> UTCTime -> IO JWTAttempt
-jwtClaims _ _ "" _ = return $ JWTClaims M.empty
-jwtClaims secret audience payload time =
+jwtClaims :: Maybe JWK -> Maybe StringOrURI -> LByteString -> UTCTime -> Maybe JSPath -> IO JWTAttempt
+jwtClaims _ _ "" _ _ = return $ JWTClaims M.empty
+jwtClaims secret audience payload time jspath =
   case secret of
     Nothing -> return JWTMissingSecret
     Just s -> do
@@ -51,7 +54,26 @@ jwtClaims secret audience payload time =
         verifyClaimsAt validation s time jwt
       return $ case eJwt of
         Left e    -> JWTInvalid e
-        Right jwt -> JWTClaims . claims2map $ jwt
+        Right jwt -> JWTClaims $ claims2map jwt jspath
+
+{-|
+  Turn JWT ClaimSet into something easier to work with,
+  also here the jspath is applied to put the "role" in the map
+-}
+claims2map :: ClaimsSet -> Maybe JSPath -> M.HashMap Text JSON.Value
+claims2map claims jspath = (\case
+    val@(JSON.Object o) ->
+      let role = maybe M.empty (M.singleton "role") $
+                 walkJSPath (Just val) =<< jspath in
+      M.delete "role" o `M.union` role -- mutating the map
+    _ -> M.empty
+  ) $ JSON.toJSON claims
+
+walkJSPath :: Maybe JSON.Value -> JSPath -> Maybe JSON.Value
+walkJSPath x                      []                = x
+walkJSPath (Just (JSON.Object o)) (JSPKey key:rest) = walkJSPath (M.lookup key o) rest
+walkJSPath (Just (JSON.Array ar)) (JSPIdx idx:rest) = walkJSPath (ar V.!? idx) rest
+walkJSPath _                      _                 = Nothing
 
 {-|
   Whether a response from jwtClaims contains a role claim
@@ -60,19 +82,9 @@ containsRole :: JWTAttempt -> Bool
 containsRole (JWTClaims claims) = M.member "role" claims
 containsRole _                  = False
 
-{-|
-  Internal helper used to turn JWT ClaimSet into something
-  easier to work with
--}
-claims2map :: ClaimsSet -> M.HashMap Text Value
-claims2map = val2map . toJSON
- where
-  val2map (Object o) = o
-  val2map _          = M.empty
-
 parseJWK :: ByteString -> JWK
 parseJWK str =
-  fromMaybe (hs256jwk str) (decode (toS str) :: Maybe JWK)
+  fromMaybe (hs256jwk str) (JSON.decode (toS str) :: Maybe JWK)
 
 {-|
   Internal helper to generate HMAC-SHA256. When the jwt key in the