From eae5857d0e5eb9839f0353d8b30c9bbf6ea308c5 Mon Sep 17 00:00:00 2001
From: Joe Nelson <cred+github@begriffs.com>
Date: Fri, 15 Apr 2016 07:30:36 -0700
Subject: [PATCH] Set role only once, and set it before other GUC vars (#560)

* Set role only once, and set it before other GUC vars

Fixes #559

* Unify role/claim logic in claimsToSQL

Suggested by @diogob
---
 CHANGELOG.md                |  2 ++
 src/PostgREST/Auth.hs       | 33 ++++++++++++++-------------------
 src/PostgREST/Middleware.hs | 13 +++++--------
 3 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cbad5209e..e57d5f6df 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Fixed
 
+- Prevent role from being changed twice - @begriffs
+
 ## [0.3.1.1] - 2016-03-28
 
 ### Fixed
diff --git a/src/PostgREST/Auth.hs b/src/PostgREST/Auth.hs
index e456271b5..3d5918350 100644
--- a/src/PostgREST/Auth.hs
+++ b/src/PostgREST/Auth.hs
@@ -12,8 +12,7 @@ In the test suite there is an example of simple login function that can be used
 very simple authentication system inside the PostgreSQL database.
 -}
 module PostgREST.Auth (
-    setRole
-  , claimsToSQL
+    claimsToSQL
   , jwtClaims
   , tokenJWT
   ) where
@@ -25,7 +24,7 @@ import           Data.Aeson.Types        (parseMaybe, emptyObject, emptyArray)
 import qualified Data.ByteString         as BS
 import qualified Data.Vector             as V
 import qualified Data.HashMap.Strict     as M
-import           Data.Maybe              (fromMaybe)
+import           Data.Maybe              (fromMaybe, maybeToList)
 import           Data.Monoid             ((<>))
 import           Data.String.Conversions (cs)
 import           Data.Text               (Text)
@@ -34,19 +33,20 @@ import           PostgREST.QueryBuilder  (pgFmtIdent, pgFmtLit, unquoted)
 import qualified Web.JWT                 as JWT
 
 {-|
-  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.
+  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.
 -}
 claimsToSQL :: M.HashMap Text Value -> [BS.ByteString]
-claimsToSQL = map setVar . M.toList
-  where
-    setVar ("role", String val) = setRole val
-    setVar (k, val) = "set local " <> cs (pgFmtIdent $ "postgrest.claims." <> k)
-                      <> " = " <> cs (valueToVariable val) <> ";"
-    valueToVariable = pgFmtLit . unquoted
+claimsToSQL claims = roleStmts <> varStmts
+ where
+  roleStmts = maybeToList $
+    (\r -> "set local role " <> r <> ";") . cs . valueToVariable <$> M.lookup "role" claims
+  varStmts = map setVar $ M.toList (M.delete "role" claims)
+  setVar (k, val) = "set local " <> cs (pgFmtIdent $ "postgrest.claims." <> k)
+                    <> " = " <> cs (valueToVariable val) <> ";"
+  valueToVariable = pgFmtLit . unquoted
 
 {-|
   Receives the JWT secret (from config) and a JWT and
@@ -70,11 +70,6 @@ jwtClaims secret input time =
   value2map (Object o) = o
   value2map _          = M.empty
 
-{-| Receives the name of a role and returns a SET ROLE statement -}
-setRole :: Text -> BS.ByteString
-setRole r = "set local role " <> cs (pgFmtLit r) <> ";"
-
-
 {-|
   Receives the JWT secret (from config) and a JWT and a JSON value
   and returns a signed JWT.
diff --git a/src/PostgREST/Middleware.hs b/src/PostgREST/Middleware.hs
index 2bc791031..cd76fa2bb 100644
--- a/src/PostgREST/Middleware.hs
+++ b/src/PostgREST/Middleware.hs
@@ -3,8 +3,7 @@
 
 module PostgREST.Middleware where
 
-import           Control.Monad                 (unless)
-import qualified Data.ByteString               as BS
+import           Data.Aeson                    (Value (..))
 import qualified Data.HashMap.Strict           as M
 import           Data.Maybe                    (fromMaybe)
 import           Data.String.Conversions       (cs)
@@ -21,7 +20,7 @@ import           Network.Wai.Middleware.Gzip   (def, gzip)
 import           Network.Wai.Middleware.Static (only, staticPolicy)
 
 import           PostgREST.ApiRequest          (pickContentType)
-import           PostgREST.Auth                (setRole, jwtClaims, claimsToSQL)
+import           PostgREST.Auth                (jwtClaims, claimsToSQL)
 import           PostgREST.Config              (AppConfig (..), corsPolicy)
 import           PostgREST.Error               (errResponse)
 
@@ -31,7 +30,6 @@ runWithClaims :: AppConfig -> NominalDiffTime ->
                  (Request -> H.Transaction Response) ->
                  Request -> H.Transaction Response
 runWithClaims conf time app req = do
-    H.sql setAnon
     let tokenStr = case split (== ' ') (cs auth) of
           ("Bearer" : t : _) -> t
           _                  -> ""
@@ -42,15 +40,14 @@ runWithClaims conf time app req = do
         if M.null claims && not (null tokenStr)
           then clientErr "Invalid JWT"
           else do
-            let cmdBatch = mconcat $ claimsToSQL claims
-            unless (BS.null cmdBatch) (H.sql cmdBatch)
+            -- role claim defaults to anon if not specified in jwt
+            H.sql . mconcat . claimsToSQL $ M.union claims (M.singleton "role" anon)
             app req
   where
     hdrs = requestHeaders req
     jwtSecret = configJwtSecret conf
     auth = fromMaybe "" $ lookup hAuthorization hdrs
-    anon = cs $ configAnonRole conf
-    setAnon = setRole anon
+    anon = String . cs $ configAnonRole conf
     clientErr = return . errResponse status400
 
 unsupportedAccept :: Application -> Application