diff --git a/app/Config.hs b/app/Config.hs --- a/app/Config.hs +++ b/app/Config.hs @@ -43,6 +43,7 @@ , configPath :: Text , configHost :: Text , configPort :: Int + , configAuditChannel :: Maybe Text , configJwtSecret :: B.ByteString , configJwtSecretIsBase64 :: Bool @@ -72,11 +73,12 @@ cPath <- C.require conf "server-root" cHost <- C.lookupDefault "*4" conf "server-host" cPort <- C.lookupDefault 3000 conf "server-port" + cAuditC <- C.lookup conf "audit-channel" -- jwt --------------- cJwtSec <- C.require conf "jwt-secret" cJwtB64 <- C.lookupDefault False conf "secret-is-base64" - return $ AppConfig cDbUri cPath cHost cPort (encodeUtf8 cJwtSec) cJwtB64 cPool + return $ AppConfig cDbUri cPath cHost cPort cAuditC (encodeUtf8 cJwtSec) cJwtB64 cPool where opts = info (helper <*> pathParser) $ diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -2,6 +2,7 @@ import Protolude import PostgRESTWS +import PostgRESTWS.Database (toPgIdentifier) import Config (AppConfig (..), PgVersion (..), minimumPgVersion, @@ -63,7 +64,7 @@ multi <- newHasqlBroadcaster pgSettings runSettings appSettings $ - postgrestWsMiddleware (configJwtSecret conf) getTime pool multi $ + postgrestWsMiddleware (toPgIdentifier . 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.2.0 +version: 0.3.3.0 synopsis: PostgREST extension to map LISTEN/NOTIFY messages to Websockets description: Please see README.md homepage: https://github.com/diogob/postgrest-ws#readme @@ -45,6 +45,7 @@ , stm-containers , stm , retry + , stringsearch 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 @@ -25,7 +25,7 @@ import PostgRESTWS.Claims import PostgRESTWS.Database -import PostgRESTWS.Broadcast (Multiplexer, onMessage, readTChan) +import PostgRESTWS.Broadcast (Multiplexer, onMessage) import PostgRESTWS.HasqlBroadcast (newHasqlBroadcaster, newHasqlBroadcasterOrError) import qualified PostgRESTWS.Broadcast as B @@ -37,18 +37,18 @@ 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 :: ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application +postgrestWsMiddleware :: Maybe PgIdentifier -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application postgrestWsMiddleware = WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp where - compose = (.) . (.) . (.) . (.) + compose = (.) . (.) . (.) . (.) . (.) -- private functions -- when the websocket is closed a ConnectionClosed Exception is triggered -- this kills all children and frees resources for us -wsApp :: ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> WS.ServerApp -wsApp secret getTime pqCon multi pendingConn = +wsApp :: Maybe PgIdentifier -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> WS.ServerApp +wsApp mAuditChannel secret getTime pool multi pendingConn = getTime >>= forkSessionsWhenTokenIsValid . validateClaims secret jwtToken where forkSessionsWhenTokenIsValid = either rejectRequest forkSessions @@ -57,6 +57,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 + forkSessions (channel, mode, validClaims) = do -- role claim defaults to anon if not specified in jwt -- We should accept only after verifying JWT @@ -65,31 +66,30 @@ WS.forkPingThread conn 30 when (hasRead mode) $ - onMessage multi channel (\ch -> - forever $ atomically (readTChan ch) >>= WS.sendTextData conn . B.payload) + onMessage multi channel $ WS.sendTextData conn . B.payload - notifySessionFinished <- if hasWrite mode - then forkAndWait $ forever $ notifySession channel validClaims pqCon conn - else newMVar () - takeMVar notifySessionFinished + when (hasWrite mode) $ + let channelName = toPgIdentifier channel + sendNotifications = void . case mAuditChannel of + Nothing -> notifyPool pool channelName + Just auditChannel -> \mesg -> + notifyPool pool channelName mesg >> + notifyPool pool auditChannel mesg + in notifySession validClaims conn sendNotifications + waitForever <- newEmptyMVar + void $ takeMVar waitForever + -- 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 :: BS.ByteString - -> A.Object - -> H.Pool - -> WS.Connection - -> IO () -notifySession channel claimsToSend pool wsCon = - WS.receiveData wsCon >>= (void . send . jsonMsg) +notifySession :: A.Object + -> WS.Connection + -> (ByteString -> IO ()) + -> IO () +notifySession claimsToSend wsCon send = + withAsync (forever relayData) wait where - send = notifyPool pool channel + relayData = WS.receiveData wsCon >>= (void . send . jsonMsg) -- we need to decode the bytestring to re-encode valid JSON for the notification jsonMsg = BL.toStrict . A.encode . Message claimsToSend . decodeUtf8With T.lenientDecode - -forkAndWait :: IO () -> IO (MVar ()) -forkAndWait io = do - mvar <- newEmptyMVar - void $ forkFinally io (\_ -> putMVar mvar ()) - return mvar diff --git a/src/PostgRESTWS/Broadcast.hs b/src/PostgRESTWS/Broadcast.hs --- a/src/PostgRESTWS/Broadcast.hs +++ b/src/PostgRESTWS/Broadcast.hs @@ -94,10 +94,10 @@ The first listener will open the channel, when a listener dies it will check if there acquire any others and close the channel when that's the case. -} -onMessage :: Multiplexer -> ByteString -> (TChan Message -> IO()) -> IO () +onMessage :: Multiplexer -> ByteString -> (Message -> IO()) -> IO () onMessage multi chan action = do listener <- atomically $ openChannelWhenNotFound >>= addListener - void $ forkFinally (action listener) disposeListener + void $ forkFinally (forever (atomically (readTChan listener) >>= action)) disposeListener where disposeListener _ = atomically $ do mC <- M.lookup chan (channels multi) diff --git a/src/PostgRESTWS/Claims.hs b/src/PostgRESTWS/Claims.hs --- a/src/PostgRESTWS/Claims.hs +++ b/src/PostgRESTWS/Claims.hs @@ -14,7 +14,6 @@ import Data.Time.Clock.POSIX (POSIXTime) import Control.Lens import Data.Aeson.Lens -import Data.Maybe (fromJust) import Data.Time.Clock (NominalDiffTime) import qualified Web.JWT as JWT @@ -67,10 +66,11 @@ jwtClaims :: JWT.Secret -> Text -> NominalDiffTime -> JWTAttempt jwtClaims _ "" _ = JWTClaims M.empty jwtClaims secret jwt time = - case isExpired <$> mClaims of - Just True -> JWTExpired + case mClaims of Nothing -> JWTInvalid - Just False -> JWTClaims $ value2map $ fromJust mClaims + Just cl -> if isExpired cl + then JWTExpired + else JWTClaims $ value2map cl where mClaims = toJSON . JWT.claims <$> JWT.decodeAndVerifySignature secret jwt isExpired claims = diff --git a/src/PostgRESTWS/Database.hs b/src/PostgRESTWS/Database.hs --- a/src/PostgRESTWS/Database.hs +++ b/src/PostgRESTWS/Database.hs @@ -6,6 +6,8 @@ , listen , unlisten , waitForNotifications + , PgIdentifier + , toPgIdentifier ) where import Protolude @@ -15,38 +17,55 @@ import Hasql.Connection (Connection, withLibPQConnection) import qualified Database.PostgreSQL.LibPQ as PQ import Data.Either.Combinators - +import qualified Data.ByteString.Char8 as B +import Data.ByteString.Search (replace) newtype Error = NotifyError Text +-- | A wrapped bytestring that represents a properly escaped and quoted PostgreSQL identifier +newtype PgIdentifier = PgIdentifier ByteString deriving (Show) + +-- | Given a PgIdentifier returns the wrapped bytestring +fromPgIdentifier :: PgIdentifier -> ByteString +fromPgIdentifier (PgIdentifier bs) = bs + +-- | Given a bytestring returns a properly quoted and escaped PgIdentifier +toPgIdentifier :: ByteString -> PgIdentifier +toPgIdentifier x = PgIdentifier $ "\"" <> strictlyReplaceQuotes (trimNullChars x) <> "\"" + where + trimNullChars :: ByteString -> ByteString + trimNullChars = B.takeWhile (/= '\x0') + strictlyReplaceQuotes :: ByteString -> ByteString + strictlyReplaceQuotes = toS . replace "\"" ("\"\"" :: ByteString) + -- | Given a Hasql Pool, a channel and a message sends a notify command to the database -notifyPool :: Pool -> ByteString -> ByteString -> IO (Either Error ()) +notifyPool :: Pool -> PgIdentifier -> ByteString -> IO (Either Error ()) notifyPool pool channel mesg = - mapError <$> use pool (sql ("NOTIFY " <> channel <> ", '" <> mesg <> "'")) + mapError <$> use pool (sql ("NOTIFY " <> fromPgIdentifier channel <> ", '" <> mesg <> "'")) where mapError :: Either UsageError () -> Either Error () mapError = mapLeft (NotifyError . show) -- | Given a Hasql Connection, a channel and a message sends a notify command to the database -notify :: Connection -> ByteString -> ByteString -> IO (Either Error ()) +notify :: Connection -> PgIdentifier -> ByteString -> IO (Either Error ()) notify con channel mesg = - mapError <$> run (sql ("NOTIFY " <> channel <> ", '" <> mesg <> "'")) con + mapError <$> run (sql ("NOTIFY " <> fromPgIdentifier channel <> ", '" <> mesg <> "'")) con where mapError :: Either S.Error () -> Either Error () mapError = mapLeft (NotifyError . show) -- | Given a Hasql Connection and a channel sends a listen command to the database -listen :: Connection -> ByteString -> IO () +listen :: Connection -> PgIdentifier -> IO () listen con channel = void $ withLibPQConnection con execListen where - execListen pqCon = void $ PQ.exec pqCon $ "LISTEN " <> channel + execListen pqCon = void $ PQ.exec pqCon $ "LISTEN " <> fromPgIdentifier channel -- | Given a Hasql Connection and a channel sends a unlisten command to the database -unlisten :: Connection -> ByteString -> IO () +unlisten :: Connection -> PgIdentifier -> IO () unlisten con channel = void $ withLibPQConnection con execListen where - execListen pqCon = void $ PQ.exec pqCon $ "UNLISTEN " <> channel + execListen pqCon = void $ PQ.exec pqCon $ "UNLISTEN " <> fromPgIdentifier channel {- | Given a function that handles notifications and a Hasql connection forks a thread that listens on the database connection and calls the handler everytime a message arrives. diff --git a/src/PostgRESTWS/HasqlBroadcast.hs b/src/PostgRESTWS/HasqlBroadcast.hs --- a/src/PostgRESTWS/HasqlBroadcast.hs +++ b/src/PostgRESTWS/HasqlBroadcast.hs @@ -75,17 +75,18 @@ -} newHasqlBroadcasterForConnection :: IO Connection -> IO Multiplexer newHasqlBroadcasterForConnection getCon = do - multi <- - newMultiplexer (\cmds msgs-> do - con <- getCon - waitForNotifications - (\c m-> atomically $ writeTQueue msgs $ Message c m) - con - forever $ do - cmd <- atomically $ readTQueue cmds - case cmd of - Open ch -> listen con ch - Close ch -> unlisten con ch - ) (\_ -> hPutStrLn stderr "Broadcaster is dead") + multi <- newMultiplexer openProducer closeProducer void $ relayMessagesForever multi return multi + where + closeProducer _ = hPutStrLn stderr "Broadcaster is dead" + openProducer cmds msgs = do + con <- getCon + waitForNotifications + (\c m-> atomically $ writeTQueue msgs $ Message c m) + con + forever $ do + cmd <- atomically $ readTQueue cmds + case cmd of + Open ch -> listen con $ toPgIdentifier ch + Close ch -> unlisten con $ toPgIdentifier ch diff --git a/test/BroadcastSpec.hs b/test/BroadcastSpec.hs --- a/test/BroadcastSpec.hs +++ b/test/BroadcastSpec.hs @@ -1,7 +1,6 @@ module BroadcastSpec (spec) where import Protolude -import Control.Concurrent.STM.TChan import Control.Concurrent.STM.TQueue import Test.Hspec @@ -26,8 +25,7 @@ output <- newTQueueIO :: IO (TQueue Message) multi <- liftIO $ newMultiplexer (\_ msgs-> atomically $ writeTQueue msgs (Message "test" "payload")) (\_ -> return ()) - void $ onMessage multi "test" (\ch -> - atomically $ readTChan ch >>= writeTQueue output) + void $ onMessage multi "test" $ atomically . writeTQueue output liftIO $ relayMessages multi diff --git a/test/DatabaseSpec.hs b/test/DatabaseSpec.hs --- a/test/DatabaseSpec.hs +++ b/test/DatabaseSpec.hs @@ -17,10 +17,10 @@ notification <- liftIO newEmptyMVar waitForNotifications (curry $ putMVar notification) con - listen con "test" + listen con $ toPgIdentifier "test" conOrError2 <- H.acquire "postgres://localhost/postgrest_test" let con2 = either (panic . show) id conOrError2 :: H.Connection - void $ notify con2 "test" "hello there" + void $ notify con2 (toPgIdentifier "test") "hello there" readMVar notification `shouldReturn` ("test", "hello there") diff --git a/test/HasqlBroadcastSpec.hs b/test/HasqlBroadcastSpec.hs --- a/test/HasqlBroadcastSpec.hs +++ b/test/HasqlBroadcastSpec.hs @@ -27,9 +27,10 @@ con <- newConnection "postgres://localhost/postgrest_test" multi <- liftIO $ newHasqlBroadcaster "postgres://localhost/postgrest_test" - atomically $ openChannelProducer multi "test" + atomically $ openChannelProducer multi "test channel" + threadDelay 1000000 - let statement = H.statement "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE query ~* 'LISTEN \"test\"')" + let statement = H.statement "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE query ~* 'LISTEN \\\"test channel\\\"')" HE.unit (HD.singleRow $ HD.value HD.bool) False query = H.query () statement booleanQueryShouldReturn con query True @@ -38,9 +39,10 @@ con <- newConnection "postgres://localhost/postgrest_test" multi <- liftIO $ newHasqlBroadcaster "postgres://localhost/postgrest_test" - atomically $ closeChannelProducer multi "test" + atomically $ closeChannelProducer multi "test channel" + threadDelay 1000000 - let statement = H.statement "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE query ~* 'UNLISTEN \"test\"')" + let statement = H.statement "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE query ~* 'UNLISTEN \\\"test channel\\\"')" HE.unit (HD.singleRow $ HD.value HD.bool) False query = H.query () statement booleanQueryShouldReturn con query True