diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -2,7 +2,6 @@ import Protolude import PostgRESTWS -import PostgRESTWS.Database (toPgIdentifier) import Config (AppConfig (..), PgVersion (..), minimumPgVersion, @@ -64,7 +63,7 @@ multi <- newHasqlBroadcaster pgSettings runSettings appSettings $ - postgrestWsMiddleware (toPgIdentifier . toS <$> configAuditChannel conf) (configJwtSecret conf) getTime pool multi $ + postgrestWsMiddleware (toS <$> configAuditChannel conf) (configJwtSecret conf) getTime pool multi $ logStdout $ staticApp $ defaultFileServerSettings $ toS $ configPath conf loadSecretFile :: AppConfig -> IO AppConfig diff --git a/postgrest-ws.cabal b/postgrest-ws.cabal --- a/postgrest-ws.cabal +++ b/postgrest-ws.cabal @@ -1,5 +1,5 @@ name: postgrest-ws -version: 0.3.3.0 +version: 0.4.0.0 synopsis: PostgREST extension to map LISTEN/NOTIFY messages to Websockets description: Please see README.md homepage: https://github.com/diogob/postgrest-ws#readme @@ -46,6 +46,7 @@ , stm , retry , stringsearch + , contravariant default-language: Haskell2010 default-extensions: OverloadedStrings, NoImplicitPrelude diff --git a/src/PostgRESTWS.hs b/src/PostgRESTWS.hs --- a/src/PostgRESTWS.hs +++ b/src/PostgRESTWS.hs @@ -19,6 +19,7 @@ import qualified Data.Text.Encoding.Error as T import Data.Time.Clock.POSIX (POSIXTime) +import qualified Data.HashMap.Strict as M import qualified Data.Aeson as A import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as BL @@ -37,7 +38,7 @@ instance A.ToJSON Message -- | Given a secret, a function to fetch the system time, a Hasql Pool and a Multiplexer this will give you a WAI middleware. -postgrestWsMiddleware :: Maybe PgIdentifier -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application +postgrestWsMiddleware :: Maybe ByteString -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application postgrestWsMiddleware = WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp where @@ -47,7 +48,7 @@ -- when the websocket is closed a ConnectionClosed Exception is triggered -- this kills all children and frees resources for us -wsApp :: Maybe PgIdentifier -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> WS.ServerApp +wsApp :: Maybe ByteString -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> WS.ServerApp wsApp mAuditChannel secret getTime pool multi pendingConn = getTime >>= forkSessionsWhenTokenIsValid . validateClaims secret jwtToken where @@ -57,7 +58,7 @@ rejectRequest = WS.rejectRequest pendingConn . encodeUtf8 -- the first char in path is '/' the rest is the token jwtToken = decodeUtf8 $ BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn - + notifySessionWithTime = notifySession getTime forkSessions (channel, mode, validClaims) = do -- role claim defaults to anon if not specified in jwt -- We should accept only after verifying JWT @@ -69,13 +70,12 @@ onMessage multi channel $ WS.sendTextData conn . B.payload when (hasWrite mode) $ - let channelName = toPgIdentifier channel - sendNotifications = void . case mAuditChannel of - Nothing -> notifyPool pool channelName + let sendNotifications = void . case mAuditChannel of + Nothing -> notifyPool pool channel Just auditChannel -> \mesg -> - notifyPool pool channelName mesg >> + notifyPool pool channel mesg >> notifyPool pool auditChannel mesg - in notifySession validClaims conn sendNotifications + in notifySessionWithTime validClaims conn sendNotifications waitForever <- newEmptyMVar void $ takeMVar waitForever @@ -83,13 +83,23 @@ -- Having both channel and claims as parameters seem redundant -- But it allows the function to ignore the claims structure and the source -- of the channel, so all claims decoding can be coded in the caller -notifySession :: A.Object +notifySession :: IO POSIXTime + -> A.Object -> WS.Connection -> (ByteString -> IO ()) -> IO () -notifySession claimsToSend wsCon send = +notifySession getTime claimsToSend wsCon send = withAsync (forever relayData) wait where - relayData = WS.receiveData wsCon >>= (void . send . jsonMsg) + relayData = jsonMsgWithTime >>= send + + jsonMsgWithTime = liftA2 jsonMsg claimsWithTime (WS.receiveData wsCon) + -- we need to decode the bytestring to re-encode valid JSON for the notification - jsonMsg = BL.toStrict . A.encode . Message claimsToSend . decodeUtf8With T.lenientDecode + jsonMsg :: M.HashMap Text A.Value -> ByteString -> ByteString + jsonMsg cl = BL.toStrict . A.encode . Message cl . decodeUtf8With T.lenientDecode + + claimsWithTime :: IO (M.HashMap Text A.Value) + claimsWithTime = do + time <- getTime + return $ M.insert "message_sent_at" (A.Number $ fromRational $ toRational time) claimsToSend diff --git a/src/PostgRESTWS/Database.hs b/src/PostgRESTWS/Database.hs --- a/src/PostgRESTWS/Database.hs +++ b/src/PostgRESTWS/Database.hs @@ -12,13 +12,18 @@ import Protolude import Hasql.Pool (Pool, UsageError, use) -import Hasql.Session (sql, run) +import Hasql.Session (sql, run, query) import qualified Hasql.Session as S +import Hasql.Query (statement) import Hasql.Connection (Connection, withLibPQConnection) +import qualified Hasql.Decoders as HD +import qualified Hasql.Encoders as HE import qualified Database.PostgreSQL.LibPQ as PQ import Data.Either.Combinators import qualified Data.ByteString.Char8 as B import Data.ByteString.Search (replace) +import Data.Functor.Contravariant (contramap) + newtype Error = NotifyError Text -- | A wrapped bytestring that represents a properly escaped and quoted PostgreSQL identifier @@ -38,12 +43,14 @@ strictlyReplaceQuotes = toS . replace "\"" ("\"\"" :: ByteString) -- | Given a Hasql Pool, a channel and a message sends a notify command to the database -notifyPool :: Pool -> PgIdentifier -> ByteString -> IO (Either Error ()) +notifyPool :: Pool -> ByteString -> ByteString -> IO (Either Error ()) notifyPool pool channel mesg = - mapError <$> use pool (sql ("NOTIFY " <> fromPgIdentifier channel <> ", '" <> mesg <> "'")) + mapError <$> use pool (query (toS channel, toS mesg) callStatement) where mapError :: Either UsageError () -> Either Error () mapError = mapLeft (NotifyError . show) + callStatement = statement ("SELECT pg_notify" <> "($1, $2)") encoder HD.unit False + encoder = contramap fst (HE.value HE.text) <> contramap snd (HE.value HE.text) -- | Given a Hasql Connection, a channel and a message sends a notify command to the database notify :: Connection -> PgIdentifier -> ByteString -> IO (Either Error ())