diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/app/Config.hs b/app/Config.hs new file mode 100644 --- /dev/null +++ b/app/Config.hs @@ -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" diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -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 "." "=" diff --git a/postgres-websockets.cabal b/postgres-websockets.cabal new file mode 100644 --- /dev/null +++ b/postgres-websockets.cabal @@ -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 diff --git a/src/PostgresWebsockets.hs b/src/PostgresWebsockets.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets.hs @@ -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 diff --git a/src/PostgresWebsockets/Broadcast.hs b/src/PostgresWebsockets/Broadcast.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/Broadcast.hs @@ -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 diff --git a/src/PostgresWebsockets/Claims.hs b/src/PostgresWebsockets/Claims.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/Claims.hs @@ -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) diff --git a/src/PostgresWebsockets/Database.hs b/src/PostgresWebsockets/Database.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/Database.hs @@ -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) diff --git a/src/PostgresWebsockets/HasqlBroadcast.hs b/src/PostgresWebsockets/HasqlBroadcast.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/HasqlBroadcast.hs @@ -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 diff --git a/test/BroadcastSpec.hs b/test/BroadcastSpec.hs new file mode 100644 --- /dev/null +++ b/test/BroadcastSpec.hs @@ -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\"" diff --git a/test/ClaimsSpec.hs b/test/ClaimsSpec.hs new file mode 100644 --- /dev/null +++ b/test/ClaimsSpec.hs @@ -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")]) diff --git a/test/DatabaseSpec.hs b/test/DatabaseSpec.hs new file mode 100644 --- /dev/null +++ b/test/DatabaseSpec.hs @@ -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") diff --git a/test/HasqlBroadcastSpec.hs b/test/HasqlBroadcastSpec.hs new file mode 100644 --- /dev/null +++ b/test/HasqlBroadcastSpec.hs @@ -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 diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}