Make HTTP headers available as GUCs #800 (#849)

src/PostgREST/ApiRequest.hs

@@ -27,7 +27,7 @@ import           Control.Arrow             ((***))
 import qualified Data.Text                 as T
 import qualified Data.Vector               as V
 import           Network.HTTP.Base         (urlEncodeVars)
-import           Network.HTTP.Types.Header (hAuthorization)
+import           Network.HTTP.Types.Header (hAuthorization, hCookie)
 import           Network.HTTP.Types.URI    (parseSimpleQuery)
 import           Network.Wai               (Request (..))
 import           Network.Wai.Parse         (parseHttpAccept)
@@ -40,6 +40,8 @@ import           PostgREST.Types           ( QualifiedIdentifier (..)
                                            , ApiRequestError(..)
                                            , toMime)
 import           Data.Ranged.Ranges        (Range(..), rangeIntersection, emptyRange)
+import qualified Data.CaseInsensitive      as CI
+import           Web.Cookie                (parseCookiesText)
 
 type RequestBody = BL.ByteString
 
@@ -92,6 +94,10 @@ data ApiRequest = ApiRequest {
   , iCanonicalQS :: ByteString
   -- | JSON Web Token
   , iJWT :: Text
+  -- | HTTP request headers
+  , iHeaders :: [(Text, Text)]
+  -- | Request Cookies
+  , iCookies :: [(Text, Text)]
   }
 
 -- | Examines HTTP request and translates it into user intent.
@@ -119,6 +125,8 @@ userApiRequest schema req reqBody
         . parseSimpleQuery
         $ rawQueryString req
       , iJWT = tokenStr
+      , iHeaders = [ (toS $ CI.foldedCase k, toS v) | (k,v) <- hdrs, k /= hAuthorization, k /= hCookie]
+      , iCookies = fromMaybe [] $ parseCookiesText <$> lookupHeader "Cookie"
       }
  where
   isTargetingProc = fromMaybe False $ (== "rpc") <$> listToMaybe path

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 (
-    claimsToSQL
-  , containsRole
+    containsRole
   , jwtClaims
   , tokenJWT
   , JWTAttempt(..)
@@ -28,24 +27,8 @@ import qualified Data.Vector             as V
 import qualified Data.HashMap.Strict     as M
 import           Data.Maybe              (fromJust)
 import           Data.Time.Clock         (NominalDiffTime)
-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.
--}
-claimsToSQL :: M.HashMap Text Value -> [ByteString]
-claimsToSQL claims = roleStmts <> varStmts
- where
-  roleStmts = maybeToList $
-    (\r -> "set local role " <> r <> ";") . toS . valueToVariable <$> M.lookup "role" claims
-  varStmts = map setVar $ M.toList (M.delete "role" claims)
-  setVar (k, val) = "set local " <> toS (pgFmtIdent $ "request.jwt.claim." <> k)
-                    <> " = " <> toS (valueToVariable val) <> ";"
-  valueToVariable = pgFmtLit . unquoted
 
 {-|
   Possible situations encountered with client JWTs

src/PostgREST/Middleware.hs

@@ -1,5 +1,6 @@
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE FlexibleContexts #-}
 
 module PostgREST.Middleware where
 
@@ -15,10 +16,11 @@ import           Network.Wai.Middleware.Gzip   (def, gzip)
 import           Network.Wai.Middleware.Static (only, staticPolicy)
 
 import           PostgREST.ApiRequest          (ApiRequest(..))
-import           PostgREST.Auth                (claimsToSQL, JWTAttempt(..))
+import           PostgREST.Auth                (JWTAttempt(..))
 import           PostgREST.Config              (AppConfig (..), corsPolicy)
 import           PostgREST.Error               (simpleError)
 import           PostgREST.Types               (ContentType (..), toHeader)
+import           PostgREST.QueryBuilder        (pgFmtLit, unquoted, pgFmtEnvVar)
 
 import           Protolude                     hiding (concat, null)
 
@@ -31,14 +33,20 @@ runWithClaims conf eClaims app req =
     JWTInvalid -> return $ unauthed "JWT invalid"
     JWTMissingSecret -> return $ simpleError status500 "Server lacks JWT secret"
     JWTClaims claims -> do
-      -- role claim defaults to anon if not specified in jwt
-      let setClaims = claimsToSQL (M.union claims (M.singleton "role" anon))
-      H.sql $ mconcat setClaims
+      H.sql $ toS.mconcat $ setRoleSql ++ claimsSql ++ headersSql ++ cookiesSql
       mapM_ H.sql customReqCheck
       app req
+      where
+        headersSql = map (pgFmtEnvVar "request.header.") $ iHeaders req
+        cookiesSql = map (pgFmtEnvVar "request.cookie.") $ iCookies req
+        claimsSql = map (pgFmtEnvVar "request.jwt.claim.") [(c,unquoted v) | (c,v) <- M.toList claimsWithRole]
+        setRoleSql = maybeToList $
+          (\r -> "set local role " <> r <> ";") . toS . pgFmtLit . unquoted <$> M.lookup "role" claimsWithRole
+        -- role claim defaults to anon if not specified in jwt
+        claimsWithRole = M.union claims (M.singleton "role" anon)
+        anon = String . toS $ configAnonRole conf
+        customReqCheck = (\f -> "select " <> toS f <> "();") <$> configReqCheck conf
   where
-    anon = String . toS $ configAnonRole conf
-    customReqCheck = (\f -> "select " <> toS f <> "();") <$> configReqCheck conf
     unauthed message = responseLBS unauthorized401
       [ toHeader CTApplicationJSON
       , ( "WWW-Authenticate"

src/PostgREST/QueryBuilder.hs

@@ -24,6 +24,7 @@ module PostgREST.QueryBuilder (
   , sourceCTEName
   , unquoted
   , ResultsWithCount
+  , pgFmtEnvVar
   ) where
 
 import qualified Hasql.Query             as H
@@ -35,7 +36,7 @@ import qualified Data.Aeson              as JSON
 import           PostgREST.RangeQuery    (NonnegRange, rangeLimit, rangeOffset, allRange)
 import           Data.Functor.Contravariant (contramap)
 import qualified Data.HashMap.Strict     as HM
-import           Data.Maybe              
+import           Data.Maybe
 import           Data.Text               (intercalate, unwords, replace, isInfixOf, toLower, split)
 import qualified Data.Text as T          (map, takeWhile, null)
 import qualified Data.Text.Encoding as T
@@ -300,7 +301,7 @@ requestToQuery schema _ (DbMutate (Insert mainTbl (PayloadJSON rows) returnings)
           if T.null colsString
             then if V.null rows then ["SELECT null WHERE false"] else ["DEFAULT VALUES"]
             else ["SELECT", colsString, "FROM json_populate_recordset(null::" , fromQi qi, ", $1)"]
-        ret = if null returnings 
+        ret = if null returnings
                   then ""
                   else unwords [" RETURNING ", intercalate ", " (map (pgFmtColumn qi) returnings)]
 requestToQuery schema _ (DbMutate (Update mainTbl (PayloadJSON rows) conditions returnings)) =
@@ -358,7 +359,7 @@ asJsonF = "coalesce(array_to_json(array_agg(row_to_json(_postgrest_t))), '[]')::
 asJsonSingleF :: SqlFragment --TODO! unsafe when the query actually returns multiple rows, used only on inserting and returning single element
 asJsonSingleF = "coalesce(string_agg(row_to_json(_postgrest_t)::text, ','), '')::character varying "
 
-asBinaryF :: FieldName -> SqlFragment 
+asBinaryF :: FieldName -> SqlFragment
 asBinaryF fieldName = "coalesce(string_agg(_postgrest_t." <> pgFmtIdent fieldName <> ", ''), '')"
 
 locationF :: [Text] -> SqlFragment
@@ -484,5 +485,9 @@ pgFmtAs (Just xx) Nothing = case lastMay xx of
   Nothing -> ""
 pgFmtAs _ (Just alias) = " AS " <> pgFmtIdent alias
 
+pgFmtEnvVar :: Text -> (Text, Text) -> SqlFragment
+pgFmtEnvVar prefix (k, v) =
+  "set local " <> pgFmtIdent (prefix <> k) <> " = " <> pgFmtLit v <> ";"
+
 trimNullChars :: Text -> Text
 trimNullChars = T.takeWhile (/= '\x0')