Eliminates SET role duplication and changes Auth module interface

src/PostgREST/Auth.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE FlexibleContexts #-}
 {-|
 Module      : PostgREST.Auth
 Description : PostgREST authorization functions.
@@ -10,10 +11,10 @@ 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 FlexibleContexts #-}
 module PostgREST.Auth (
   setRole
-  , setJWTEnv
+  , claimsToSQL
+  , jwtClaims
   , tokenJWT
   ) where
 
@@ -32,22 +33,28 @@ import qualified Web.JWT                 as JWT
 import qualified Data.HashMap.Lazy       as H
 
 {-|
-  Receives the JWT secret (from config) and a JWT and
-  returns a list of PostgreSQL statements to set the claims
-  as user defined GUCs.
-  Except if we have a claim called role, this one is mapped to
-  a SET ROLE statement.
+  Receives a map of JWT claims and returns a list
+  of PostgreSQL statements to set the claims as user defined GUCs.
+  Except if we have a claim called role,
+  this one is mapped to a SET ROLE statement.
   In case there is any problem decoding the JWT it returns Nothing.
 -}
-setJWTEnv :: Text -> Text -> Maybe [Text]
-setJWTEnv secret input = setDBEnv jwtClaims
+claimsToSQL :: JWT.ClaimsMap -> [Text]
+claimsToSQL = map setVar . toList
   where
-    setDBEnv maybeClaims = (map setVar . toList) <$> maybeClaims
     setVar ("role", String val) = setRole val
     setVar (k, val) = "set local postgrest.claims." <> pgFmtIdent k <>
                       " = " <> valueToVariable val <> ";"
     valueToVariable = pgFmtLit . unquoted
-    jwtClaims = JWT.unregisteredClaims <$> JWT.claims <$> decoded
+
+{-|
+  Receives the JWT secret (from config) and a JWT and
+  returns a map of JWT claims
+  In case there is any problem decoding the JWT it returns Nothing.
+-}
+jwtClaims :: Text -> Text -> Maybe JWT.ClaimsMap
+jwtClaims secret input = JWT.unregisteredClaims <$> JWT.claims <$> decoded
+  where
     decoded = JWT.decodeAndVerifySignature (JWT.secret secret) input
 
 -- | Receives the name of a role and returns a SET ROLE statement

src/PostgREST/Middleware.hs

@@ -26,13 +26,14 @@ import           Network.Wai.Middleware.Gzip   (def, gzip)
 import           Network.Wai.Middleware.Static (only, staticPolicy)
 
 import           PostgREST.App                 (contentTypeForAccept)
-import           PostgREST.Auth                (setRole, setJWTEnv)
+import           PostgREST.Auth                (setRole, jwtClaims, claimsToSQL)
 import           PostgREST.Config              (AppConfig (..), corsPolicy)
 
 import           Prelude hiding(concat)
 
 import qualified Data.Vector             as V
 import qualified Hasql.Backend           as B
+import qualified Data.Map.Lazy           as M
 
 runWithClaims :: forall s. AppConfig ->
                  (Request -> H.Tx P.Postgres s Response) ->
@@ -41,16 +42,20 @@ runWithClaims conf app req = do
     mapM_ H.unitEx $ stmt <$> env
     app req
  where
+   stmt = (flip $ flip B.Stmt V.empty) True
    hdrs = requestHeaders req
    jwtSecret = (cs $ configJwtSecret conf) :: Text
    auth = fromMaybe "" $ lookup hAuthorization hdrs
    anon = cs $ configAnonRole conf
-   jwtEnv =
+   claims =
+     fromMaybe (M.fromList []) $
      case split (==' ') (cs auth) of
-       ("Bearer" : jwt : _) -> fromMaybe [] (setJWTEnv jwtSecret jwt)
-       _ -> []
-   env = setRole anon : jwtEnv
-   stmt = (flip $ flip B.Stmt V.empty) True
+       ("Bearer" : jwt : _) -> jwtClaims jwtSecret jwt
+       _ -> Nothing
+   env = if M.member "role" claims
+            then jwtEnv
+            else setRole anon : jwtEnv
+   jwtEnv = claimsToSQL claims
 
 redirectInsecure :: Application -> Application
 redirectInsecure app req respond = do