postgres-websockets (empty) → 0.4.2.0
raw patch · 15 files changed
+1006/−0 lines, 15 filesdep +aesondep +ansi-wl-pprintdep +basesetup-changed
Dependencies added: aeson, ansi-wl-pprint, base, base64-bytestring, bytestring, configurator, containers, contravariant, either, hasql, hasql-pool, heredoc, hspec, hspec-wai, hspec-wai-json, http-types, jose, jwt, lens, lens-aeson, optparse-applicative, postgres-websockets, postgresql-libpq, protolude, retry, stm, stm-containers, stringsearch, text, time, transformers, unordered-containers, wai, wai-app-static, wai-extra, wai-websockets, warp, websockets
Files
- LICENSE +20/−0
- Setup.hs +2/−0
- app/Config.hs +134/−0
- app/Main.hs +89/−0
- postgres-websockets.cabal +106/−0
- src/PostgresWebsockets.hs +108/−0
- src/PostgresWebsockets/Broadcast.hs +117/−0
- src/PostgresWebsockets/Claims.hs +99/−0
- src/PostgresWebsockets/Database.hs +98/−0
- src/PostgresWebsockets/HasqlBroadcast.hs +95/−0
- test/BroadcastSpec.hs +46/−0
- test/ClaimsSpec.hs +17/−0
- test/DatabaseSpec.hs +26/−0
- test/HasqlBroadcastSpec.hs +48/−0
- test/Spec.hs +1/−0
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2014 Joe Nelson++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ app/Config.hs view
@@ -0,0 +1,134 @@+{-|+Module : Config+Description : Manages PostgresWebsockets configuration options.++This module provides a helper function to read the command line+arguments using the optparse-applicative and the AppConfig type to store+them. It also can be used to define other middleware configuration that+may be delegated to some sort of external configuration.++It currently includes a hardcoded CORS policy but this could easly be+turned in configurable behaviour if needed.++Other hardcoded options such as the minimum version number also belong here.+-}+module Config ( prettyVersion+ , readOptions+ , minimumPgVersion+ , PgVersion (..)+ , AppConfig (..)+ )+ where++import System.IO.Error (IOError)+import Control.Applicative+import qualified Data.Configurator as C+import qualified Data.Configurator.Types as C+import Data.Monoid+import Data.Text (intercalate, lines)+import Data.Text.Encoding (encodeUtf8)+import Data.Version (versionBranch)+import Options.Applicative hiding (str)+import Paths_postgrest_ws (version)+import Text.Heredoc+import Text.PrettyPrint.ANSI.Leijen hiding ((<>))+import qualified Text.PrettyPrint.ANSI.Leijen as L+import Protolude hiding (intercalate, (<>))++-- | Config file settings for the server+data AppConfig = AppConfig {+ configDatabase :: Text+ , configPath :: Text+ , configHost :: Text+ , configPort :: Int+ , configAuditChannel :: Maybe Text+ , configJwtSecret :: ByteString+ , configJwtSecretIsBase64 :: Bool+ , configPool :: Int+ }++-- | User friendly version number+prettyVersion :: Text+prettyVersion = intercalate "." $ map show $ versionBranch version++-- | Function to read and parse options from the command line+readOptions :: IO AppConfig+readOptions = do+ -- First read the config file path from command line+ cfgPath <- customExecParser parserPrefs opts+ -- Now read the actual config file+ conf <- catch+ (C.load [C.Required cfgPath])+ configNotfoundHint++ handle missingKeyHint $ do+ -- db ----------------+ cDbUri <- C.require conf "db-uri"+ cPool <- C.lookupDefault 10 conf "db-pool"+ -- server ------------+ 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 cAuditC (encodeUtf8 cJwtSec) cJwtB64 cPool++ where+ opts = info (helper <*> pathParser) $+ fullDesc+ <> progDesc (+ "PostgREST "+ <> toS prettyVersion+ <> " / create a REST API to an existing Postgres database"+ )+ <> footerDoc (Just $+ text "Example Config File:"+ L.<> nest 2 (hardline L.<> exampleCfg)+ )++ parserPrefs = prefs showHelpOnError++ configNotfoundHint :: IOError -> IO a+ configNotfoundHint e = die $ "Cannot open config file:\n\t" <> show e++ missingKeyHint :: C.KeyError -> IO a+ missingKeyHint (C.KeyError n) =+ die $+ "Required config parameter \"" <> n <> "\" is missing or of wrong type.\n" <>+ "Documentation for configuration options available at\n" <>+ "\thttp://postgrest.com/en/v0.4/admin.html#configuration\n\n" <>+ "Try the --example-config option to see how to configure PostgREST."++ exampleCfg :: Doc+ exampleCfg = vsep . map (text . toS) . lines $+ [str|db-uri = "postgres://user:pass@localhost:5432/dbname"+ |db-pool = 10+ |+ |server-root = "./client-example"+ |server-host = "*4"+ |server-port = 3000+ |+ |## choose a secret to enable JWT auth+ |## (use "@filename" to load from separate file)+ |# jwt-secret = "foo"+ |# secret-is-base64 = false+ |]+++pathParser :: Parser FilePath+pathParser =+ strArgument $+ metavar "FILENAME" <>+ help "Path to configuration file"++data PgVersion = PgVersion {+ pgvNum :: Int32+, pgvName :: Text+}++-- | Tells the minimum PostgreSQL version required by this version of PostgresWebsockets+minimumPgVersion :: PgVersion+minimumPgVersion = PgVersion 90300 "9.3"
+ app/Main.hs view
@@ -0,0 +1,89 @@+module Main where++import Protolude hiding (replace)+import PostgresWebsockets+import Config (AppConfig (..),+ PgVersion (..),+ minimumPgVersion,+ prettyVersion,+ readOptions)++import qualified Data.ByteString as BS+import qualified Data.ByteString.Base64 as B64+import Data.String (IsString (..))+import Data.Text (pack, replace, strip, stripPrefix)+import Data.Text.Encoding (encodeUtf8, decodeUtf8)+import qualified Hasql.Query as H+import qualified Hasql.Session as H+import qualified Hasql.Decoders as HD+import qualified Hasql.Encoders as HE+import qualified Hasql.Pool as P+import Network.Wai.Application.Static++import Network.Wai.Handler.Warp+import Network.Wai.Middleware.RequestLogger (logStdout)+import System.IO (BufferMode (..),+ hSetBuffering)++isServerVersionSupported :: H.Session Bool+isServerVersionSupported = do+ ver <- H.query () pgVersion+ return $ ver >= pgvNum minimumPgVersion+ where+ pgVersion =+ H.statement "SELECT current_setting('server_version_num')::integer"+ HE.unit (HD.singleRow $ HD.value HD.int4) False++main :: IO ()+main = do+ hSetBuffering stdout LineBuffering+ hSetBuffering stdin LineBuffering+ hSetBuffering stderr NoBuffering++ conf <- loadSecretFile =<< readOptions+ let host = configHost conf+ port = configPort conf+ pgSettings = toS (configDatabase conf)+ appSettings = setHost ((fromString . toS) host)+ . setPort port+ . setServerName (toS $ "postgrest/" <> prettyVersion)+ . setTimeout 3600+ $ defaultSettings++ putStrLn $ ("Listening on port " :: Text) <> show (configPort conf)++ pool <- P.acquire (configPool conf, 10, pgSettings)+ multi <- newHasqlBroadcaster pgSettings++ runSettings appSettings $+ postgrestWsMiddleware (toS <$> configAuditChannel conf) (configJwtSecret conf) pool multi $+ logStdout $ staticApp $ defaultFileServerSettings $ toS $ configPath conf++loadSecretFile :: AppConfig -> IO AppConfig+loadSecretFile conf = extractAndTransform secret+ where+ secret = decodeUtf8 $ configJwtSecret conf+ isB64 = configJwtSecretIsBase64 conf++ extractAndTransform :: Text -> IO AppConfig+ extractAndTransform s =+ fmap setSecret $ transformString isB64 =<<+ case stripPrefix "@" s of+ Nothing -> return . encodeUtf8 $ s+ Just filename -> chomp <$> BS.readFile (toS filename)+ where+ chomp bs = fromMaybe bs (BS.stripSuffix "\n" bs)++ -- Turns the Base64url encoded JWT into Base64+ transformString :: Bool -> ByteString -> IO ByteString+ transformString False t = return t+ transformString True t =+ case B64.decode $ encodeUtf8 $ strip $ replaceUrlChars $ decodeUtf8 t of+ Left errMsg -> panic $ pack errMsg+ Right bs -> return bs++ setSecret bs = conf {configJwtSecret = bs}++ -- replace: Replace every occurrence of one substring with another+ replaceUrlChars =+ replace "_" "/" . replace "-" "+" . replace "." "="
+ postgres-websockets.cabal view
@@ -0,0 +1,106 @@+name: postgres-websockets+version: 0.4.2.0+synopsis: PostgREST extension to map LISTEN/NOTIFY messages to Websockets+description: Please see README.md+homepage: https://github.com/diogob/postgres-websockets#readme+license: BSD3+license-file: LICENSE+author: Diogo Biazus+maintainer: diogo@biazus.me+copyright: 2016 Diogo Biazus+category: Web+build-type: Simple+-- extra-source-files:+cabal-version: >=1.10++library+ hs-source-dirs: src+ ghc-options: -Wall+ exposed-modules: PostgresWebsockets+ , PostgresWebsockets.Broadcast+ , PostgresWebsockets.Database+ , PostgresWebsockets.HasqlBroadcast+ , PostgresWebsockets.Claims+ build-depends: base >= 4.7 && < 5+ , hasql-pool >= 0.4 && < 0.5+ , text >= 1.2 && < 2+ , wai >= 3.2 && < 4+ , websockets >= 0.9 && < 0.11+ , wai-websockets >= 3.0 && < 4+ , http-types >= 0.9+ , bytestring >= 0.10+ , postgresql-libpq+ , lens+ , lens-aeson+ , jose >= 0.6+ , unordered-containers >= 0.2+ , postgresql-libpq >= 0.9 && < 1.0+ , aeson >= 0.11+ , protolude >= 0.2+ , jwt >= 0.7.2 && < 0.8+ , hasql >= 0.19+ , either+ , stm-containers+ , stm+ , retry+ , stringsearch+ , time+ , contravariant+ default-language: Haskell2010+ default-extensions: OverloadedStrings, NoImplicitPrelude, LambdaCase++executable postgres-websockets+ hs-source-dirs: app+ main-is: Main.hs+ other-modules: Config+ ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N+ build-depends: base >= 4.7 && < 5+ , transformers >= 0.4 && < 0.6+ , hasql >= 0.19+ , hasql-pool >= 0.4+ , warp >= 3.2 && < 4+ , postgres-websockets+ , protolude >= 0.2+ , base64-bytestring+ , bytestring+ , configurator+ , optparse-applicative+ , text+ , time+ , wai+ , wai-extra+ , wai-app-static+ , heredoc+ , ansi-wl-pprint+ default-language: Haskell2010+ default-extensions: OverloadedStrings, NoImplicitPrelude, QuasiQuotes++test-suite postgres-websockets-test+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Spec.hs+ other-modules: BroadcastSpec+ , ClaimsSpec+ , DatabaseSpec+ , HasqlBroadcastSpec+ build-depends: base+ , protolude >= 0.2+ , postgres-websockets+ , containers+ , hspec+ , hspec-wai+ , hspec-wai-json+ , aeson+ , hasql+ , hasql-pool+ , http-types+ , unordered-containers >= 0.2+ , wai-extra+ , stm+ ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N+ default-language: Haskell2010+ default-extensions: OverloadedStrings, NoImplicitPrelude++source-repository head+ type: git+ location: https://github.com/diogob/postgres-websockets
+ src/PostgresWebsockets.hs view
@@ -0,0 +1,108 @@+{-| PostgresWebsockets Middleware, composing this allows postgrest to create+ websockets connections that will communicate with the database through LISTEN/NOTIFY channels.+-}+{-# LANGUAGE DeriveGeneric #-}++module PostgresWebsockets+ ( postgrestWsMiddleware+ -- * Re-exports+ , newHasqlBroadcaster+ , newHasqlBroadcasterOrError+ ) where++import qualified Hasql.Pool as H+import qualified Network.Wai as Wai+import qualified Network.Wai.Handler.WebSockets as WS+import qualified Network.WebSockets as WS+import Protolude++import qualified Data.Aeson as A+import qualified Data.ByteString.Char8 as BS+import qualified Data.ByteString.Lazy as BL+import qualified Data.HashMap.Strict as M+import qualified Data.Text.Encoding.Error as T+import Data.Time.Clock.POSIX (getPOSIXTime)+import PostgresWebsockets.Broadcast (Multiplexer, onMessage)+import qualified PostgresWebsockets.Broadcast as B+import PostgresWebsockets.Claims+import PostgresWebsockets.Database+import PostgresWebsockets.HasqlBroadcast (newHasqlBroadcaster,+ newHasqlBroadcasterOrError)++data Message = Message+ { claims :: A.Object+ , payload :: Text+ } deriving (Show, Eq, Generic)++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 ByteString -> ByteString -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application+postgrestWsMiddleware =+ WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp+ where+ compose = (.) . (.) . (.) . (.)++-- private functions++-- when the websocket is closed a ConnectionClosed Exception is triggered+-- this kills all children and frees resources for us+wsApp :: Maybe ByteString -> ByteString -> H.Pool -> Multiplexer -> WS.ServerApp+wsApp mAuditChannel secret pool multi pendingConn =+ validateClaims requestChannel secret (toS jwtToken) >>= either rejectRequest forkSessions+ where+ hasRead m = m == ("r" :: ByteString) || m == ("rw" :: ByteString)+ hasWrite m = m == ("w" :: ByteString) || m == ("rw" :: ByteString)+ rejectRequest = WS.rejectRequest pendingConn . encodeUtf8+ -- the URI has one of the two formats - /:jwt or /:channel/:jwt + pathElements = BS.split '/' $ BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn+ jwtToken + | length pathElements > 1 = headDef "" $ tailSafe pathElements+ | length pathElements <= 1 = headDef "" pathElements+ requestChannel+ | length pathElements > 1 = Just $ headDef "" pathElements+ | length pathElements <= 1 = Nothing+ notifySessionWithTime = notifySession+ forkSessions (channel, mode, validClaims) = do+ -- role claim defaults to anon if not specified in jwt+ -- We should accept only after verifying JWT+ conn <- WS.acceptRequest pendingConn+ -- Fork a pinging thread to ensure browser connections stay alive+ WS.forkPingThread conn 30++ when (hasRead mode) $+ onMessage multi channel $ WS.sendTextData conn . B.payload++ when (hasWrite mode) $+ let sendNotifications = void . case mAuditChannel of+ Nothing -> notifyPool pool channel+ Just auditChannel -> \mesg ->+ notifyPool pool channel mesg >>+ notifyPool pool auditChannel mesg+ in notifySessionWithTime 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 :: A.Object+ -> WS.Connection+ -> (ByteString -> IO ())+ -> IO ()+notifySession claimsToSend wsCon send =+ withAsync (forever relayData) wait+ where+ 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 :: 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 <- getPOSIXTime+ return $ M.insert "message_delivered_at" (A.Number $ fromRational $ toRational time) claimsToSend
+ src/PostgresWebsockets/Broadcast.hs view
@@ -0,0 +1,117 @@+{-| PostgresWebsockets functions to broadcast messages to several listening clients+ This module provides a type called Multiplexer.+ The multiplexer contains a map of channels and a producer thread.++ This module avoids any database implementation details, it is used by HasqlBroadcast where+ the database logic is combined.+-}+module PostgresWebsockets.Broadcast ( Multiplexer (src)+ , Message (..)+ , SourceCommands (..)+ , newMultiplexer+ , onMessage+ , relayMessages+ , relayMessagesForever+ , openChannelProducer+ , closeChannelProducer+ -- * Re-exports+ , readTQueue+ , writeTQueue+ , readTChan+ ) where++import Protolude+import qualified STMContainers.Map as M+import Control.Concurrent.STM.TChan+import Control.Concurrent.STM.TQueue++import GHC.Show++data SourceCommands = Open ByteString | Close ByteString deriving (Show)+data Message = Message { channel :: ByteString+ , payload :: ByteString+ } deriving (Eq, Show)++data Multiplexer = Multiplexer { channels :: M.Map ByteString Channel+ , src :: ThreadId+ , commands :: TQueue SourceCommands+ , messages :: TQueue Message+ }++instance Show Multiplexer where+ show Multiplexer{} = "Multiplexer"++data Channel = Channel { broadcast :: TChan Message+ , listeners :: Integer+ , close :: STM ()+ }++-- | Open a multiplexer's channel+openChannelProducer :: Multiplexer -> ByteString -> STM ()+openChannelProducer multi ch = writeTQueue (commands multi) (Open ch)++-- | Close a multiplexer's channel+closeChannelProducer :: Multiplexer -> ByteString -> STM ()+closeChannelProducer multi chan = writeTQueue (commands multi) (Close chan)++-- | Opens a thread that relays messages from the producer thread to the channels forever+relayMessagesForever :: Multiplexer -> IO ThreadId+relayMessagesForever = forkIO . forever . relayMessages++-- | Reads the messages from the producer and relays them to the active listeners in their respective channels.+relayMessages :: Multiplexer -> IO ()+relayMessages multi =+ atomically $ do+ m <- readTQueue (messages multi)+ mChannel <- M.lookup (channel m) (channels multi)+ case mChannel of+ Nothing -> return ()+ Just c -> writeTChan (broadcast c) m++newMultiplexer :: (TQueue SourceCommands -> TQueue Message -> IO a)+ -> (Either SomeException a -> IO ())+ -> IO Multiplexer+newMultiplexer openProducer closeProducer = do+ cmds <- newTQueueIO+ msgs <- newTQueueIO+ m <- liftA2 Multiplexer M.newIO (forkFinally (openProducer cmds msgs) closeProducer)+ return $ m cmds msgs++openChannel :: Multiplexer -> ByteString -> STM Channel+openChannel multi chan = do+ c <- newBroadcastTChan+ let newChannel = Channel{ broadcast = c+ , listeners = 0+ , close = closeChannelProducer multi chan+ }+ M.insert newChannel chan (channels multi)+ openChannelProducer multi chan+ return newChannel++{- | Adds a listener to a certain multiplexer's channel.+ The listener must be a function that takes a 'TChan Message' and perform any IO action.+ All listeners run in their own thread.+ 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 -> (Message -> IO()) -> IO ()+onMessage multi chan action = do+ listener <- atomically $ openChannelWhenNotFound >>= addListener+ void $ forkFinally (forever (atomically (readTChan listener) >>= action)) disposeListener+ where+ disposeListener _ = atomically $ do+ mC <- M.lookup chan (channels multi)+ let c = fromMaybe (panic $ "trying to remove listener from non existing channel: " <> toS chan) mC+ M.delete chan (channels multi)+ if listeners c - 1 > 0+ then M.insert Channel{ broadcast = broadcast c, listeners = listeners c - 1, close = close c} chan (channels multi)+ else closeChannelProducer multi chan+ openChannelWhenNotFound =+ M.lookup chan (channels multi) >>= \case+ Nothing -> openChannel multi chan+ Just ch -> return ch+ addListener ch = do+ M.delete chan (channels multi)+ let newChannel = Channel{ broadcast = broadcast ch, listeners = listeners ch + 1, close = close ch}+ M.insert newChannel chan (channels multi)+ dupTChan $ broadcast newChannel
+ src/PostgresWebsockets/Claims.hs view
@@ -0,0 +1,99 @@+{-| This module provides the JWT claims validation. Since websockets and+ listening connections in the database tend to be resource intensive+ (not to mention stateful) we need claims authorizing a specific channel and+ mode of operation.+-}+module PostgresWebsockets.Claims+ ( validateClaims+ ) where++import Control.Lens+import qualified Crypto.JOSE.Types as JOSE.Types+import Crypto.JWT+import Data.Aeson (Value (..), decode, toJSON)+import qualified Data.HashMap.Strict as M+import Protolude+++type Claims = M.HashMap Text Value+type ConnectionInfo = (ByteString, ByteString, Claims)++{-| Given a secret, a token and a timestamp it validates the claims and returns+ either an error message or a triple containing channel, mode and claims hashmap.+-}+validateClaims :: Maybe ByteString -> ByteString -> LByteString -> IO (Either Text ConnectionInfo)+validateClaims requestChannel secret jwtToken =+ runExceptT $ do+ cl <- liftIO $ jwtClaims (parseJWK secret) jwtToken+ cl' <- case cl of+ JWTClaims c -> pure c+ _ -> throwError "Error"+ channel <- claimAsJSON requestChannel "channel" cl'+ mode <- claimAsJSON Nothing "mode" cl'+ pure (channel, mode, cl')++ where+ claimAsJSON :: Maybe ByteString -> Text -> Claims -> ExceptT Text IO ByteString+ claimAsJSON defaultVal name cl = case M.lookup name cl of+ Just (String s) -> pure $ encodeUtf8 s+ Just _ -> throwError "claim is not string value"+ Nothing -> nonExistingClaim defaultVal name++ nonExistingClaim :: Maybe ByteString -> Text -> ExceptT Text IO ByteString+ nonExistingClaim Nothing name = throwError (name <> " not in claims")+ nonExistingClaim (Just defaultVal) _ = pure defaultVal++{- Private functions and types copied from postgrest++ This code duplication will be short lived since postgrest will migrate towards jose+ Then this library will use jose's verifyClaims and error types.+-}+{-|+ Possible situations encountered with client JWTs+-}+data JWTAttempt = JWTInvalid JWTError+ | JWTMissingSecret+ | JWTClaims (M.HashMap Text Value)+ deriving Eq++{-|+ Receives the JWT secret (from config) and a JWT and returns a map+ of JWT claims.+-}+jwtClaims :: JWK -> LByteString -> IO JWTAttempt+jwtClaims _ "" = return $ JWTClaims M.empty+jwtClaims secret payload = do+ let validation = defaultJWTValidationSettings (const True)+ eJwt <- runExceptT $ do+ jwt <- decodeCompact payload+ verifyClaims validation secret jwt+ return $ case eJwt of+ Left e -> JWTInvalid e+ Right jwt -> JWTClaims . claims2map $ jwt++{-|+ Internal helper used to turn JWT ClaimSet into something+ easier to work with+-}+claims2map :: ClaimsSet -> M.HashMap Text Value+claims2map = val2map . toJSON+ where+ val2map (Object o) = o+ val2map _ = M.empty++{-|+ Internal helper to generate HMAC-SHA256. When the jwt key in the+ config file is a simple string rather than a JWK object, we'll+ apply this function to it.+-}+hs256jwk :: ByteString -> JWK+hs256jwk key =+ fromKeyMaterial km+ & jwkUse .~ Just Sig+ & jwkAlg .~ (Just $ JWSAlg HS256)+ where+ km = OctKeyMaterial (OctKeyParameters (JOSE.Types.Base64Octets key))++parseJWK :: ByteString -> JWK+parseJWK str =+ fromMaybe (hs256jwk str) (decode (toS str) :: Maybe JWK)
+ src/PostgresWebsockets/Database.hs view
@@ -0,0 +1,98 @@+{-| This module encapsulates knowledge about the SQL commands and the Hasql interface.+-}+module PostgresWebsockets.Database+ ( notifyPool+ , notify+ , listen+ , unlisten+ , waitForNotifications+ , PgIdentifier+ , toPgIdentifier+ ) where++import Protolude hiding (replace)+import Hasql.Pool (Pool, UsageError, use)+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+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 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 ())+notify con channel mesg =+ 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 -> PgIdentifier -> IO ()+listen con channel =+ void $ withLibPQConnection con execListen+ where+ 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 -> PgIdentifier -> IO ()+unlisten con channel =+ void $ withLibPQConnection con execListen+ where+ 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.++ The message handler passed as first argument needs two parameters channel and payload.+-}++waitForNotifications :: (ByteString -> ByteString -> IO()) -> Connection -> IO ()+waitForNotifications sendNotification con =+ withLibPQConnection con $ void . forkIO . forever . pqFetch+ where+ pqFetch pqCon = do+ mNotification <- PQ.notifies pqCon+ case mNotification of+ Nothing -> do+ mfd <- PQ.socket pqCon+ case mfd of+ Nothing -> panic "Error checking for PostgreSQL notifications"+ Just fd -> do+ void $ threadWaitRead fd+ void $ PQ.consumeInput pqCon+ Just notification ->+ sendNotification (PQ.notifyRelname notification) (PQ.notifyExtra notification)
+ src/PostgresWebsockets/HasqlBroadcast.hs view
@@ -0,0 +1,95 @@+{-| Uses Broadcast module adding database as a source producer.+ This module provides a function to produce a 'Multiplexer' from a Hasql 'Connection'.+ The producer issues a LISTEN command upon Open commands and UNLISTEN upon Close.+-}+module PostgresWebsockets.HasqlBroadcast+ ( newHasqlBroadcaster+ , newHasqlBroadcasterOrError+ -- re-export+ , acquire+ , relayMessages+ , relayMessagesForever+ ) where++import Protolude++import Hasql.Connection+import Data.Either.Combinators (mapBoth)+import Data.Function (id)+import Control.Retry (RetryStatus, retrying, capDelay, exponentialBackoff)++import PostgresWebsockets.Database+import PostgresWebsockets.Broadcast++{- | Returns a multiplexer from a connection URI, keeps trying to connect in case there is any error.+ This function also spawns a thread that keeps relaying the messages from the database to the multiplexer's listeners+-}+newHasqlBroadcaster :: ByteString -> IO Multiplexer+newHasqlBroadcaster = newHasqlBroadcasterForConnection . tryUntilConnected++{- | Returns a multiplexer from a connection URI or an error message on the left case+ This function also spawns a thread that keeps relaying the messages from the database to the multiplexer's listeners+-}+newHasqlBroadcasterOrError :: ByteString -> IO (Either ByteString Multiplexer)+newHasqlBroadcasterOrError =+ acquire >=> (sequence . mapBoth show (newHasqlBroadcasterForConnection . return))++tryUntilConnected :: ByteString -> IO Connection+tryUntilConnected =+ fmap (either (panic "Failure on connection retry") id) . retryConnection+ where+ retryConnection conStr = retrying retryPolicy shouldRetry (const $ acquire conStr)+ maxDelayInMicroseconds = 32000000+ firstDelayInMicroseconds = 1000000+ retryPolicy = capDelay maxDelayInMicroseconds $ exponentialBackoff firstDelayInMicroseconds+ shouldRetry :: RetryStatus -> Either ConnectionError Connection -> IO Bool+ shouldRetry _ con =+ case con of+ Left err -> do+ putErrLn $ "Error connecting notification listener to database: " <> show err+ return True+ _ -> return False++{- | Returns a multiplexer from an IO Connection, listen for different database notification channels using the connection produced.++ This function also spawns a thread that keeps relaying the messages from the database to the multiplexer's listeners++ To listen on channels *chat*++ @+ import Protolude+ import PostgresWebsockets.HasqlBroadcast+ import PostgresWebsockets.Broadcast+ import Hasql.Connection++ main = do+ conOrError <- H.acquire "postgres://localhost/test_database"+ let con = either (panic . show) id conOrError :: Connection+ multi <- newHasqlBroadcaster con++ onMessage multi "chat" (\ch ->+ forever $ fmap print (atomically $ readTChan ch)+ @++-}+newHasqlBroadcasterForConnection :: IO Connection -> IO Multiplexer+newHasqlBroadcasterForConnection getCon = do+ multi <- newMultiplexer openProducer closeProducer+ void $ relayMessagesForever multi+ return multi+ where+ closeProducer _ = putErrLn "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+++putErrLn :: Text -> IO ()+putErrLn = hPutStrLn stderr
+ test/BroadcastSpec.hs view
@@ -0,0 +1,46 @@+module BroadcastSpec (spec) where++import Protolude+import Control.Concurrent.STM.TQueue++import Test.Hspec++import PostgresWebsockets.Broadcast++spec :: Spec+spec = do+ describe "newMultiplexer" $+ it "opens a separate thread for a producer function" $ do+ output <- newTQueueIO :: IO (TQueue ThreadId)++ void $ liftIO $ newMultiplexer (\_ _-> do+ tid <- myThreadId+ atomically $ writeTQueue output tid+ ) (\_ -> return ())++ outMsg <- atomically $ readTQueue output+ myThreadId `shouldNotReturn` outMsg+ describe "relayMessages" $+ it "relays a single message from producer to 1 listener on 1 test channel" $ do+ output <- newTQueueIO :: IO (TQueue Message)+ multi <- liftIO $ newMultiplexer (\_ msgs->+ atomically $ writeTQueue msgs (Message "test" "payload")) (\_ -> return ())+ void $ onMessage multi "test" $ atomically . writeTQueue output++ liftIO $ relayMessages multi++ outMsg <- atomically $ readTQueue output+ outMsg `shouldBe` Message "test" "payload"+ describe "onMessage" $+ it "sends an open command to the producer" $ do+ output <- newTQueueIO :: IO (TQueue Text)+ multi <- liftIO $ newMultiplexer (\cmds _->+ atomically $ do+ cmd <- readTQueue cmds+ writeTQueue output (show cmd :: Text)+ ) (\_ -> return ())++ void $ onMessage multi "new channel" (\_ -> return ())++ outMsg <- atomically $ readTQueue output+ outMsg `shouldBe` "Open \"new channel\""
+ test/ClaimsSpec.hs view
@@ -0,0 +1,17 @@+module ClaimsSpec (spec) where++import Protolude++import qualified Data.HashMap.Strict as M+import Test.Hspec+import Data.Aeson (Value (..) )++import PostgresWebsockets.Claims++spec :: Spec+spec =+ describe "validate claims"+ $ it "should succeed using a matching token"+ $ validateClaims Nothing "reallyreallyreallyreallyverysafe"+ "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0In0.1d4s-at2kWj8OSabHZHTbNh1dENF7NWy_r0ED3Rwf58"+ `shouldReturn` Right ("test", "r", M.fromList[("mode",String "r"),("channel",String "test")])
+ test/DatabaseSpec.hs view
@@ -0,0 +1,26 @@+module DatabaseSpec (spec) where++import Protolude++import Data.Function (id)+import qualified Hasql.Connection as H+import Test.Hspec++import PostgresWebsockets.Database++spec :: Spec+spec =+ describe "waitForNotifications" $+ it "does not block the connection ann trigger action upon notification" $ do+ conOrError <- H.acquire "postgres://localhost/postgrest_test"+ let con = either (panic . show) id conOrError :: H.Connection+ notification <- liftIO newEmptyMVar++ waitForNotifications (curry $ putMVar notification) con+ listen con $ toPgIdentifier "test"++ conOrError2 <- H.acquire "postgres://localhost/postgrest_test"+ let con2 = either (panic . show) id conOrError2 :: H.Connection+ void $ notify con2 (toPgIdentifier "test") "hello there"++ readMVar notification `shouldReturn` ("test", "hello there")
+ test/HasqlBroadcastSpec.hs view
@@ -0,0 +1,48 @@+module HasqlBroadcastSpec (spec) where++import Protolude++import Data.Function (id)+import qualified Hasql.Query as H+import qualified Hasql.Session as H+import qualified Hasql.Decoders as HD+import qualified Hasql.Encoders as HE++import Test.Hspec++import PostgresWebsockets.Broadcast+import PostgresWebsockets.HasqlBroadcast++spec :: Spec+spec = describe "newHasqlBroadcaster" $ do+ let booleanQueryShouldReturn con query expected =+ either (panic . show) id+ <$> H.run query con+ `shouldReturn` expected+ newConnection connStr =+ either (panic . show) id+ <$> acquire connStr++ it "start listening on a database connection as we send an Open command" $ do+ con <- newConnection "postgres://localhost/postgrest_test"+ multi <- liftIO $ newHasqlBroadcaster "postgres://localhost/postgrest_test"++ atomically $ openChannelProducer multi "test channel"+ threadDelay 1000000++ 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++ it "stops listening on a database connection as we send a Close command" $ do+ con <- newConnection "postgres://localhost/postgrest_test"+ multi <- liftIO $ newHasqlBroadcaster "postgres://localhost/postgrest_test"++ atomically $ closeChannelProducer multi "test channel"+ threadDelay 1000000++ 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
+ test/Spec.hs view
@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}