Makes JWT generation possible in RPC endpoints

Fixes SET execution to execute in separate statements as Hasql uses
prepared statements we need to send 1 commend per statement.

postgrest.cabal

@@ -53,6 +53,8 @@ executable postgrest
                      , mtl
                      , cassava
                      , jwt
+                     , lens
+                     , lens-aeson >= 1.0.0.5
                      , parsec
                      , errors
                      , bifunctors
@@ -66,34 +68,50 @@ library
 
   default-language:    Haskell2010
   default-extensions:  OverloadedStrings, ScopedTypeVariables, QuasiQuotes
-  build-depends:       base >=4.6 && <5
+  build-depends:       HTTP
-                     , hasql, hasql-backend
+                     , MissingH
-                     , hasql-postgres
-                     , warp, wai
-                     , wai-extra, wai-cors
-                     , wai-middleware-static
-                     , HTTP, convertible, http-types
-                     , case-insensitive
-                     , scientific, time
-                     , aeson, network
-                     , bytestring, text, split, string-conversions
-                     , stringsearch
-                     , containers, unordered-containers
-                     , optparse-applicative
-                     , regex-base, regex-tdfa
                      , Ranged-sets
-                     , transformers, MissingH
+                     , aeson
-                     , bcrypt, base64-string
+                     , base >=4.6 && <5
-                     , network-uri
+                     , base64-string
-                     , resource-pool
+                     , bcrypt
-                     , blaze-builder
-                     , vector
-                     , mtl
-                     , cassava
-                     , jwt
-                     , parsec
-                     , errors
                      , bifunctors
+                     , blaze-builder
+                     , bytestring
+                     , case-insensitive
+                     , cassava
+                     , containers
+                     , convertible
+                     , errors
+                     , hasql
+                     , hasql-backend
+                     , hasql-postgres
+                     , http-types
+                     , jwt
+                     , lens
+                     , lens-aeson >= 1.0.0.5
+                     , mtl
+                     , network
+                     , network-uri
+                     , optparse-applicative
+                     , parsec
+                     , regex-base
+                     , regex-tdfa
+                     , resource-pool
+                     , scientific
+                     , split
+                     , string-conversions
+                     , stringsearch
+                     , text
+                     , time
+                     , transformers
+                     , unordered-containers
+                     , vector
+                     , wai
+                     , wai-cors
+                     , wai-extra
+                     , wai-middleware-static
+                     , warp
 
   Other-Modules:       Paths_postgrest
   Exposed-Modules:     PostgREST.App
@@ -163,6 +181,8 @@ Test-Suite spec
                      , process
                      , heredoc
                      , jwt
+                     , lens
+                     , lens-aeson >= 1.0.0.5
                      , parsec
                      , errors
                      , bifunctors

src/PostgREST/App.hs

@@ -40,6 +40,7 @@ import           Network.Wai.Internal      (Response (..))
 import           Network.Wai.Parse         (parseHttpAccept)
 
 import           Data.Aeson
+import           Data.Aeson.Types (emptyArray)
 import           Data.Monoid
 import qualified Data.Vector               as V
 import qualified Hasql                     as H
@@ -53,6 +54,7 @@ import           PostgREST.PgStructure
 import           PostgREST.QueryBuilder
 import           PostgREST.RangeQuery
 import           PostgREST.Types
+import           PostgREST.Auth (tokenJWT)
 
 import           Prelude
 
@@ -172,9 +174,11 @@ app dbstructure conf reqBody req =
         then do
           let call = B.Stmt "select " V.empty True <>
                 asJson (callProc qi $ fromMaybe M.empty (decode reqBody))
-          body :: Maybe (Identity Text) <- H.maybeEx call
+          bodyJson :: Maybe (Identity Value) <- H.maybeEx call
           return $ responseLBS status200 [jsonH]
-            (cs $ fromMaybe "[]" $ runIdentity <$> body)
+                (if hasPrefer "return=jwt"
+                   then ("{\"token\":\"" <> (cs $ tokenJWT jwtSecret $ fromMaybe "[]" $ runIdentity <$> bodyJson) <> "\"}")
+                   else (cs $ encode $ fromMaybe emptyArray $ runIdentity <$> bodyJson))
         else return $ responseLBS status404 [] ""
 
       -- check that proc exists

src/PostgREST/Auth.hs

@@ -1,53 +1,23 @@
 {-# LANGUAGE FlexibleContexts #-}
 module PostgREST.Auth (
-    DbRole
+  setRole
-  , LoginAttempt (..)
-  , setRole
   , setJWTEnv
+  , tokenJWT
   ) where
 
 import           Control.Applicative
-import           Control.Monad           (mzero)
-
 import           Data.Aeson
 import           Data.Map (fromList, toList)
+import           Data.Maybe (fromMaybe)
 import           Data.Monoid
 import           Data.String.Conversions (cs)
 import           Data.Text (Text)
 import           PostgREST.PgQuery       (pgFmtLit)
-import           Prelude 
+import           Prelude
 import qualified Web.JWT                 as JWT
-
+import qualified Data.HashMap.Lazy       as HashMap
-
+import Data.Aeson.Lens
-
+import Control.Lens.Operators
-data AuthUser = AuthUser {
-    userId   :: String
-  , userPass :: String
-  , userRole :: Maybe String
-  } deriving (Show)
-
-instance FromJSON AuthUser where
-  parseJSON (Object v) = AuthUser <$>
-                         v .: "id" <*>
-                         v .: "pass" <*>
-                         v .:? "role"
-  parseJSON _ = mzero
-
-instance ToJSON AuthUser where
-  toJSON u = object [
-      "id" .= userId u
-    , "pass" .= userPass u
-    , "role" .= userRole u ]
-
-type DbRole = Text
-type UserId = Text
-
-data LoginAttempt =
-    NoCredentials
-  | MalformedAuth
-  | LoginFailed
-  | LoginSuccess DbRole UserId
-  deriving (Eq, Show)
 
 setJWTEnv :: Text -> Text -> Maybe [Text]
 setJWTEnv secret input = setDBEnv $ jwtClaims secret input
@@ -57,9 +27,9 @@ setDBEnv maybeClaims =
   (map setVar . toList) <$> maybeClaims
   where
     setVar ("role", String val) = setRole val
-    setVar (key, String val) = "set local postgrest.claims" <> key <> " = " <> pgFmtLit val <> ";"
+    setVar (k, String val) = "set local postgrest.claims." <> k <> " = " <> pgFmtLit val <> ";"
-    setVar (key, Bool val) = "set local postgrest.claims" <> key <> " = " <> showText val <> ";"
+    setVar (k, Bool val) = "set local postgrest.claims." <> k <> " = " <> showText val <> ";"
-    setVar (key, Number val) = "set local postgrest.claims" <> key <> " = " <> showText val <> ";"
+    setVar (k, Number val) = "set local postgrest.claims." <> k <> " = " <> showText val <> ";"
     setVar _ = ""
     showText :: Show a => a -> Text
     showText = cs . show
@@ -73,9 +43,10 @@ jwtClaims secret input = claims
     claims = JWT.unregisteredClaims <$> JWT.claims <$> decoded
     decoded = JWT.decodeAndVerifySignature (JWT.secret secret) input
 
-tokenJWT :: Text -> Text -> Text -> Text
+tokenJWT :: Text -> Value -> Text
-tokenJWT secret uid role = JWT.encodeSigned JWT.HS256 (JWT.secret secret) claimsSet
+tokenJWT secret claims = JWT.encodeSigned JWT.HS256 (JWT.secret secret) claimsSet
   where
     claimsSet = JWT.def {
-      JWT.unregisteredClaims = Data.Map.fromList [("id", String uid), ("role", String role)]
+      JWT.unregisteredClaims = Data.Map.fromList claimsList
     }
+    claimsList =  fromMaybe [] $ HashMap.toList <$> (claims ^? nth 0 . _Object)

src/PostgREST/Middleware.hs

@@ -35,7 +35,7 @@ runWithClaims :: forall s. AppConfig ->
                  (Request -> H.Tx P.Postgres s Response) ->
                  Request -> H.Tx P.Postgres s Response
 runWithClaims conf app req = do
-    H.unitEx $ B.Stmt env V.empty True
+    mapM_ H.unitEx $ stmt <$> env
     app req
  where
    hdrs = requestHeaders req
@@ -46,7 +46,8 @@ runWithClaims conf app req = do
      case split (==' ') (cs auth) of
        ("Bearer" : jwt : _) -> fromMaybe [] (setJWTEnv jwtSecret jwt)
        _ -> []
-   env = concat $ setRole anon : jwtEnv
+   env = setRole anon : jwtEnv
+   stmt = (flip $ flip B.Stmt V.empty) True
 
 redirectInsecure :: Application -> Application
 redirectInsecure app req respond = do