packages feed

postgres-websockets 0.7.0.0 → 0.8.0.0

raw patch · 12 files changed

+546/−353 lines, 12 filesdep +asyncdep +lens-aesondep +networkdep −transformersdep ~hasqldep ~hasql-pooldep ~lens

Dependencies added: async, lens-aeson, network

Dependencies removed: transformers

Dependency ranges changed: hasql, hasql-pool, lens, websockets

Files

− app/Config.hs 
@@ -1,65 +0,0 @@-{-|-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 Env-import           Data.Text                   (intercalate)-import           Data.Version                (versionBranch)-import           Paths_postgres_websockets   (version)-import           Protolude hiding            (intercalate, (<>), optional)---- | Config file settings for the server-data AppConfig = AppConfig {-    configDatabase          :: Text-  , configPath              :: Maybe Text-  , configHost              :: Text-  , configPort              :: Int-  , configListenChannel     :: 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 environment-readOptions :: IO AppConfig-readOptions =-    Env.parse (header "You need to configure some environment variables to start the service.") $-      AppConfig <$> var (str <=< nonempty) "PGWS_DB_URI"  (help "String to connect to PostgreSQL")-                <*> optional (var str "PGWS_ROOT_PATH" (help "Root path to serve static files, unset to disable."))-                <*> var str "PGWS_HOST" (def "*4" <> helpDef show <> help "Address the server will listen for websocket connections")-                <*> var auto "PGWS_PORT" (def 3000 <> helpDef show <> help "Port the server will listen for websocket connections")-                <*> var str "PGWS_LISTEN_CHANNEL" (def "postgres-websockets-listener" <> helpDef show <> help "Master channel used in the database to send or read messages in any notification channel")-                <*> var str "PGWS_JWT_SECRET" (help "Secret used to sign JWT tokens used to open communications channels")-                <*> var auto "PGWS_JWT_SECRET_BASE64" (def False <> helpDef show <> help "Indicate whether the JWT secret should be decoded from a base64 encoded string")-                <*> var auto "PGWS_POOL_SIZE" (def 10 <> helpDef show <> help "How many connection to the database should be used by the connection pool")--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
@@ -1,45 +1,9 @@ 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.Statement                      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 Data.Time.Clock (UTCTime, getCurrentTime)-import Control.AutoUpdate       ( defaultUpdateSettings-                                , mkAutoUpdate-                                , updateAction-                                )--import           Network.Wai (Application, responseLBS)-import           Network.HTTP.Types (status200)-import           Network.Wai.Handler.Warp-import           Network.Wai.Middleware.RequestLogger (logStdout)-import           System.IO                            (BufferMode (..),-                                                       hSetBuffering)+import Protolude+import PostgresWebsockets -isServerVersionSupported :: H.Session Bool-isServerVersionSupported = do-  ver <- H.statement () pgVersion-  return $ ver >= pgvNum minimumPgVersion- where-  pgVersion =-    H.Statement "SELECT current_setting('server_version_num')::integer"-      HE.noParams (HD.singleRow $ HD.column $ HD.nonNullable HD.int4) False+import System.IO (BufferMode (..), hSetBuffering)  main :: IO () main = do@@ -51,67 +15,5 @@                <> prettyVersion                <> " / Connects websockets to PostgreSQL asynchronous notifications." -  conf <- loadSecretFile =<< readOptions-  shutdownSignal <- newEmptyMVar-  let host = configHost conf-      port = configPort conf-      listenChannel = toS $ configListenChannel conf-      pgSettings = toS (configDatabase conf)-      waitForShutdown cl = void $ forkIO (takeMVar shutdownSignal >> cl >> die "Shutting server down...")--      appSettings = setHost ((fromString . toS) host)-                  . setPort port-                  . setServerName (toS $ "postgres-websockets/" <> prettyVersion)-                  . setTimeout 3600-                  . setInstallShutdownHandler waitForShutdown-                  . setGracefulShutdownTimeout (Just 5)-                  $ defaultSettings--  putStrLn $ ("Listening on port " :: Text) <> show (configPort conf)--  let shutdown = putErrLn ("Broadcaster connection is dead" :: Text) >> putMVar shutdownSignal ()-  pool <- P.acquire (configPool conf, 10, pgSettings)-  multi <- newHasqlBroadcaster shutdown listenChannel pgSettings-  getTime <- mkGetTime--  runSettings appSettings $-    postgresWsMiddleware getTime listenChannel (configJwtSecret conf) pool multi $-    logStdout $ maybe dummyApp staticApp' (configPath conf)--  where-    mkGetTime :: IO (IO UTCTime)-    mkGetTime = mkAutoUpdate defaultUpdateSettings {updateAction = getCurrentTime}-    staticApp' :: Text -> Application-    staticApp' = staticApp . defaultFileServerSettings . toS-    dummyApp :: Application-    dummyApp _ respond =-        respond $ responseLBS status200 [("Content-Type", "text/plain")] "Hello, Web!"--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 "." "="+  conf <- loadConfig+  void $ serve conf
postgres-websockets.cabal view
@@ -1,5 +1,5 @@ name:                postgres-websockets-version:             0.7.0.0+version:             0.8.0.0 synopsis:            Middleware to map LISTEN/NOTIFY messages to Websockets description:         Please see README.md homepage:            https://github.com/diogob/postgres-websockets#readme@@ -20,6 +20,11 @@                      , PostgresWebsockets.Broadcast                      , PostgresWebsockets.HasqlBroadcast                      , PostgresWebsockets.Claims+                     , PostgresWebsockets.Config++  other-modules:       Paths_postgres_websockets+                     , PostgresWebsockets.Server+                     , PostgresWebsockets.Middleware   build-depends:       base >= 4.7 && < 5                      , hasql-pool >= 0.5 && < 0.6                      , text >= 1.2 && < 1.3@@ -45,32 +50,26 @@                      , time >= 1.8.0.2 && < 1.9                      , contravariant >= 1.5.2 && < 1.6                      , alarmclock >= 0.7.0.2 && < 0.8+                     , async >= 2.2.0 && < 2.3+                     , envparse >= 0.4.1+                     , base64-bytestring >= 1.0.0.3 && < 1.1+                     , bytestring >= 0.10+                     , warp >= 3.2 && < 4+                     , wai-extra >= 3.0.29 && < 3.1+                     , wai-app-static >= 3.1.7.1 && < 3.2+                     , auto-update >= 0.1.6 && < 0.2+   default-language:    Haskell2010-  default-extensions: OverloadedStrings, NoImplicitPrelude, LambdaCase+  default-extensions: OverloadedStrings, NoImplicitPrelude, LambdaCase, RecordWildCards  executable postgres-websockets   hs-source-dirs:      app   main-is:             Main.hs-  other-modules:       Config-                     , Paths_postgres_websockets+  other-modules:   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.3-                     , base64-bytestring >= 1.0.0.3 && < 1.1-                     , bytestring >= 0.10-                     , text >= 1.2 && < 1.3-                     , time >= 1.8.0.2 && < 1.9-                     , wai >= 3.2 && < 4-                     , wai-extra >= 3.0.29 && < 3.1-                     , wai-app-static >= 3.1.7.1 && < 3.2-                     , http-types >= 0.9-                     , envparse >= 0.4.1-                     , auto-update >= 0.1.6 && < 0.2   default-language:    Haskell2010   default-extensions: OverloadedStrings, NoImplicitPrelude, QuasiQuotes @@ -81,6 +80,7 @@   other-modules:       BroadcastSpec                      , ClaimsSpec                      , HasqlBroadcastSpec+                     , ServerSpec   build-depends:       base                      , protolude >= 0.2.3                      , postgres-websockets@@ -96,6 +96,10 @@                      , unordered-containers >= 0.2                      , wai-extra >= 3.0.29 && < 3.1                      , stm >= 2.5.0.0 && < 2.6+                     , websockets >= 0.12.7.0 && < 0.13+                     , network >= 2.8.0.1 && < 2.9+                     , lens >= 4.17.1 && < 4.18+                     , lens-aeson   ghc-options:         -Wall -threaded -rtsopts -with-rtsopts=-N   default-language:    Haskell2010   default-extensions: OverloadedStrings, NoImplicitPrelude
src/PostgresWebsockets.hs view
@@ -1,117 +1,16 @@-{-| PostgresWebsockets Middleware, composing this allows postgrest to create-    websockets connections that will communicate with the database through LISTEN/NOTIFY channels.--}-{-# LANGUAGE DeriveGeneric #-}+{-|+Module      : PostgresWebsockets+Description : PostgresWebsockets main library interface. +These are all function necessary to start a fully functionaing service.+-} module PostgresWebsockets-  ( postgresWsMiddleware-  -- * Re-exports-  , newHasqlBroadcaster-  , newHasqlBroadcasterOrError+  ( prettyVersion+  , loadConfig+  , serve+  , postgresWsMiddleware   ) where -import qualified Hasql.Pool                     as H-import qualified Hasql.Notifications            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 (UTCTime)-import           Data.Time.Clock.POSIX          (utcTimeToPOSIXSeconds, posixSecondsToUTCTime)-import           Control.Concurrent.AlarmClock (newAlarmClock, setAlarm)-import           PostgresWebsockets.Broadcast          (Multiplexer, onMessage)-import qualified PostgresWebsockets.Broadcast          as B-import           PostgresWebsockets.Claims-import           PostgresWebsockets.HasqlBroadcast     (newHasqlBroadcaster,-                                                 newHasqlBroadcasterOrError)--data Message = Message-  { claims  :: A.Object-  , channel :: Text-  , 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.-postgresWsMiddleware :: IO UTCTime -> Text -> ByteString -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application-postgresWsMiddleware =-  WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp-  where-    compose = (.) . (.) . (.) . (.) . (.)---- private functions-jwtExpirationStatusCode :: Word16-jwtExpirationStatusCode = 3001---- when the websocket is closed a ConnectionClosed Exception is triggered--- this kills all children and frees resources for us-wsApp :: IO UTCTime -> Text -> ByteString -> H.Pool -> Multiplexer -> WS.ServerApp-wsApp getTime dbChannel secret pool multi pendingConn =-  getTime >>= 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-    forkSessions (ch, 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.withPingThread conn 30 (pure ()) $ do-            case M.lookup "exp" validClaims of-              Just (A.Number expClaim) -> do-                connectionExpirer <- newAlarmClock $ const (WS.sendCloseCode conn jwtExpirationStatusCode ("JWT expired" :: ByteString))-                setAlarm connectionExpirer (posixSecondsToUTCTime $ realToFrac expClaim)-              Just _ -> pure ()-              Nothing -> pure ()--            when (hasRead mode) $-              onMessage multi ch $ WS.sendTextData conn . B.payload--            when (hasWrite mode) $-              let sendNotifications = void . H.notifyPool pool dbChannel . toS-              in notifySession validClaims (toS ch) conn getTime 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-              -> Text-              -> WS.Connection-              -> IO UTCTime-              -> (ByteString -> IO ())-              -> IO ()-notifySession claimsToSend ch wsCon getTime 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 ch . decodeUtf8With T.lenientDecode--    claimsWithChannel = M.insert "channel" (A.String ch) claimsToSend-    claimsWithTime :: IO (M.HashMap Text A.Value)-    claimsWithTime = do-      time <- utcTimeToPOSIXSeconds <$> getTime-      return $ M.insert "message_delivered_at" (A.Number $ realToFrac time) claimsWithChannel+import           PostgresWebsockets.Middleware+import           PostgresWebsockets.Server+import           PostgresWebsockets.Config
src/PostgresWebsockets/Broadcast.hs view
@@ -1,9 +1,13 @@-{-| 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.+{-| +Module      : PostgresWebsockets.Broadcast+Description : Distribute messages from one producer to several consumers. -    This module avoids any database implementation details, it is used by HasqlBroadcast where-    the database logic is combined.+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 (..)
src/PostgresWebsockets/Claims.hs view
@@ -1,10 +1,14 @@-{-| 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+Description : Parse and validate JWT to open postgres-websockets channels.++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+  ( ConnectionInfo,validateClaims   ) where  import           Control.Lens@@ -12,51 +16,63 @@ import           Crypto.JWT import qualified Data.HashMap.Strict as M import           Protolude+import Data.List import Data.Time.Clock (UTCTime)-import Data.String (String, fromString) import qualified Data.Aeson as JSON-import qualified Data.Aeson.Types as JSON - type Claims = M.HashMap Text JSON.Value-type ConnectionInfo = (ByteString, ByteString, Claims)+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 -> UTCTime -> IO (Either Text ConnectionInfo)-validateClaims requestChannel secret jwtToken time =-  runExceptT $ do-    cl <- liftIO $ jwtClaims time (parseJWK secret) jwtToken-    cl' <- case cl of-      JWTClaims c -> pure c-      JWTInvalid JWTExpired -> throwError "Token expired"-      _ -> 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 (JSON.String s) -> pure $ encodeUtf8 s-      Just _ -> throwError "claim is not string value"-      Nothing -> nonExistingClaim defaultVal name+validateClaims+  :: Maybe ByteString+  -> ByteString+  -> LByteString+  -> UTCTime+  -> IO (Either Text ConnectionInfo)+validateClaims requestChannel secret jwtToken time = runExceptT $ do+  cl  <- liftIO $ jwtClaims time (parseJWK secret) jwtToken+  cl' <- case cl of+    JWTClaims  c          -> pure c+    JWTInvalid JWTExpired -> throwError "Token expired"+    JWTInvalid err -> throwError $ "Error: " <> show err+  channels  <-  let chs = claimAsJSONList "channels" cl' in pure $ case claimAsJSON "channel" cl' of+    Just c ->  case chs of+      Just cs ->  nub (c : cs)+      Nothing ->  [c]+    Nothing -> fromMaybe [] chs+  mode <-+    let md = claimAsJSON "mode" cl'+    in case md of+          Just m  -> pure m+          Nothing -> throwError "Missing mode"+  requestedAllowedChannels <- case (requestChannel, length channels) of+    (Just rc, 0) -> pure [rc]+    (Just rc, _) -> pure $ filter (== rc) channels+    (Nothing, _) -> pure channels+  validChannels <- if null requestedAllowedChannels then throwError "No allowed channels" else pure requestedAllowedChannels+  pure (validChannels, mode, cl') -    nonExistingClaim :: Maybe ByteString -> Text -> ExceptT Text IO ByteString-    nonExistingClaim Nothing name = throwError (name <> " not in claims")-    nonExistingClaim (Just defaultVal) _ = pure defaultVal+ where+  claimAsJSON :: Text -> Claims -> Maybe ByteString+  claimAsJSON name cl = case M.lookup name cl of+    Just (JSON.String s) -> Just $ encodeUtf8 s+    _ -> Nothing -{- Private functions and types copied from postgrest+  claimAsJSONList :: Text -> Claims -> Maybe [ByteString]+  claimAsJSONList name cl = case M.lookup name cl of+    Just channelsJson ->+      case JSON.fromJSON channelsJson :: JSON.Result [Text] of+        JSON.Success channelsList -> Just $ encodeUtf8 <$> channelsList+        _ -> Nothing+    Nothing -> Nothing -   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 JSON.Value)                 deriving Eq @@ -66,11 +82,11 @@ -} jwtClaims :: UTCTime -> JWK -> LByteString -> IO JWTAttempt jwtClaims _ _ "" = return $ JWTClaims M.empty-jwtClaims time jwk payload = do+jwtClaims time jwk' payload = do   let config = defaultJWTValidationSettings (const True)   eJwt <- runExceptT $ do     jwt <- decodeCompact payload-    verifyClaimsAt config jwk time jwt+    verifyClaimsAt config jwk' time jwt   return $ case eJwt of     Left e    -> JWTInvalid e     Right jwt -> JWTClaims . claims2map $ jwt
+ src/PostgresWebsockets/Config.hs view
@@ -0,0 +1,102 @@+{-|+Module      : PostgresWebsockets.Config+Description : Manages PostgresWebsockets configuration options.++This module provides a helper function to read the command line+arguments using  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.+-}+module PostgresWebsockets.Config +        ( prettyVersion+        , loadConfig+        , warpSettings+        , AppConfig (..)+        ) where++import           Env+import           Data.Text                   (intercalate, pack, replace, strip, stripPrefix)+import           Data.Version                (versionBranch)+import           Paths_postgres_websockets   (version)+import           Protolude hiding            (intercalate, (<>), optional, replace)+import           Data.String (IsString(..))+import           Network.Wai.Handler.Warp+import qualified Data.ByteString                      as BS+import qualified Data.ByteString.Base64               as B64++-- | Config file settings for the server+data AppConfig = AppConfig {+    configDatabase          :: Text+  , configPath              :: Maybe Text+  , configHost              :: Text+  , configPort              :: Int+  , configListenChannel     :: Text+  , configJwtSecret         :: ByteString+  , configJwtSecretIsBase64 :: Bool+  , configPool              :: Int+  }++-- | User friendly version number+prettyVersion :: Text+prettyVersion = intercalate "." $ map show $ versionBranch version++-- | Load all postgres-websockets config from Environment variables. This can be used to use just the middleware or to feed into warpSettings+loadConfig :: IO AppConfig+loadConfig = readOptions >>= loadSecretFile++-- | Given a shutdown handler and an AppConfig builds a Warp Settings to start a stand-alone server+warpSettings :: (IO () -> IO ()) -> AppConfig -> Settings+warpSettings waitForShutdown AppConfig{..} =+      setHost (fromString $ toS configHost)+                  . setPort configPort+                  . setServerName (toS $ "postgres-websockets/" <> prettyVersion)+                  . setTimeout 3600+                  . setInstallShutdownHandler waitForShutdown+                  . setGracefulShutdownTimeout (Just 5)+                  $ defaultSettings+++-- private++-- | Function to read and parse options from the environment+readOptions :: IO AppConfig+readOptions =+    Env.parse (header "You need to configure some environment variables to start the service.") $+      AppConfig <$> var (str <=< nonempty) "PGWS_DB_URI"  (help "String to connect to PostgreSQL")+                <*> optional (var str "PGWS_ROOT_PATH" (help "Root path to serve static files, unset to disable."))+                <*> var str "PGWS_HOST" (def "*4" <> helpDef show <> help "Address the server will listen for websocket connections")+                <*> var auto "PGWS_PORT" (def 3000 <> helpDef show <> help "Port the server will listen for websocket connections")+                <*> var str "PGWS_LISTEN_CHANNEL" (def "postgres-websockets-listener" <> helpDef show <> help "Master channel used in the database to send or read messages in any notification channel")+                <*> var str "PGWS_JWT_SECRET" (help "Secret used to sign JWT tokens used to open communications channels")+                <*> var auto "PGWS_JWT_SECRET_BASE64" (def False <> helpDef show <> help "Indicate whether the JWT secret should be decoded from a base64 encoded string")+                <*> var auto "PGWS_POOL_SIZE" (def 10 <> helpDef show <> help "How many connection to the database should be used by the connection pool")++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 "." "="+
src/PostgresWebsockets/HasqlBroadcast.hs view
@@ -1,6 +1,10 @@-{-| 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.Broadcast+Description : Build a Hasql.Notifications based producer 'Multiplexer'.++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@@ -76,7 +80,6 @@     onMessage multi "chat" (\ch ->       forever $ fmap print (atomically $ readTChan ch)    @- -} newHasqlBroadcasterForChannel :: IO () -> Text -> IO Connection -> IO Multiplexer newHasqlBroadcasterForChannel onConnectionFailure ch getCon = do
+ src/PostgresWebsockets/Middleware.hs view
@@ -0,0 +1,127 @@+{-|+Module      : PostgresWebsockets.Middleware+Description : PostgresWebsockets WAI middleware, add functionality to any WAI application.++Allow websockets connections that will communicate with the database through LISTEN/NOTIFY channels.+-}+{-# LANGUAGE DeriveGeneric #-}++module PostgresWebsockets.Middleware+  ( postgresWsMiddleware+  ) where++import qualified Hasql.Pool                     as H+import qualified Hasql.Notifications            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 (UTCTime)+import           Data.Time.Clock.POSIX          (utcTimeToPOSIXSeconds, posixSecondsToUTCTime)+import           Control.Concurrent.AlarmClock (newAlarmClock, setAlarm)+import           PostgresWebsockets.Broadcast          (Multiplexer, onMessage)+import qualified PostgresWebsockets.Broadcast          as B+import           PostgresWebsockets.Claims++data Message = Message+  { claims  :: A.Object+  , channel :: Text+  , 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.+postgresWsMiddleware :: IO UTCTime -> Text -> ByteString -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application+postgresWsMiddleware =+  WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp+  where+    compose = (.) . (.) . (.) . (.) . (.)++-- private functions+jwtExpirationStatusCode :: Word16+jwtExpirationStatusCode = 3001++-- when the websocket is closed a ConnectionClosed Exception is triggered+-- this kills all children and frees resources for us+wsApp :: IO UTCTime -> Text -> ByteString -> H.Pool -> Multiplexer -> WS.ServerApp+wsApp getTime dbChannel secret pool multi pendingConn =+  getTime >>= 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 :: Text -> IO ()+    rejectRequest msg = do+      putErrLn $ "Rejecting Request: " <> msg+      WS.rejectRequest pendingConn (toS msg)++    -- the URI has one of the two formats - /:jwt or /:channel/:jwt+    pathElements = BS.split '/' $ BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn+    jwtToken = +      case length pathElements `compare` 1 of+        GT -> headDef "" $ tailSafe pathElements+        _ -> headDef "" pathElements+    requestChannel =+      case length pathElements `compare` 1 of+        GT -> Just $ headDef "" pathElements+        _ -> Nothing+    forkSessions :: ConnectionInfo -> IO ()+    forkSessions (chs, mode, validClaims) = do+          -- We should accept only after verifying JWT+          conn <- WS.acceptRequest pendingConn+          -- Fork a pinging thread to ensure browser connections stay alive+          WS.withPingThread conn 30 (pure ()) $ do+            case M.lookup "exp" validClaims of+              Just (A.Number expClaim) -> do+                connectionExpirer <- newAlarmClock $ const (WS.sendCloseCode conn jwtExpirationStatusCode ("JWT expired" :: ByteString))+                setAlarm connectionExpirer (posixSecondsToUTCTime $ realToFrac expClaim)+              Just _ -> pure ()+              Nothing -> pure ()++            when (hasRead mode) $+              forM_ chs $ flip (onMessage multi) $ WS.sendTextData conn . B.payload++            when (hasWrite mode) $+              let sendNotifications = void . H.notifyPool pool dbChannel . toS+              in notifySession validClaims conn getTime sendNotifications chs++            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+              -> IO UTCTime+              -> (ByteString -> IO ())+              -> [ByteString]+              -> IO ()+notifySession claimsToSend wsCon getTime send chs =+  withAsync (forever relayData) wait+  where+    relayData = do +      msg <- WS.receiveData wsCon+      forM_ chs (relayChannelData msg . toS)++    relayChannelData msg ch = do+      claims' <- claimsWithTime ch+      send $ jsonMsg ch claims' msg++    -- we need to decode the bytestring to re-encode valid JSON for the notification+    jsonMsg :: Text -> M.HashMap Text A.Value -> ByteString -> ByteString+    jsonMsg ch cl = BL.toStrict . A.encode . Message cl ch . decodeUtf8With T.lenientDecode++    claimsWithTime :: Text -> IO (M.HashMap Text A.Value)+    claimsWithTime ch = do+      time <- utcTimeToPOSIXSeconds <$> getTime+      return $ M.insert "message_delivered_at" (A.Number $ realToFrac time) (claimsWithChannel ch)++    claimsWithChannel ch = M.insert "channel" (A.String ch) claimsToSend
+ src/PostgresWebsockets/Server.hs view
@@ -0,0 +1,53 @@+{-|+Module      : PostgresWebsockets.Server+Description : Functions to start a full stand-alone PostgresWebsockets server.+-}+module PostgresWebsockets.Server+        ( serve +        ) where++import           Protolude+import           PostgresWebsockets.Middleware+import           PostgresWebsockets.Config+import           PostgresWebsockets.HasqlBroadcast (newHasqlBroadcaster)++import qualified Hasql.Pool as P+import           Network.Wai.Application.Static+import           Data.Time.Clock (UTCTime, getCurrentTime)+import Control.AutoUpdate       ( defaultUpdateSettings+                                , mkAutoUpdate+                                , updateAction+                                )+import           Network.Wai (Application, responseLBS)+import           Network.HTTP.Types (status200)+import           Network.Wai.Handler.Warp+import           Network.Wai.Middleware.RequestLogger (logStdout)++-- | Start a stand-alone warp server using the parameters from AppConfig and a opening a database connection pool.+serve :: AppConfig -> IO ()+serve conf = do+  shutdownSignal <- newEmptyMVar+  let listenChannel = toS $ configListenChannel conf+      pgSettings = toS (configDatabase conf)+      waitForShutdown cl = void $ forkIO (takeMVar shutdownSignal >> cl >> die "Shutting server down...")+      appSettings = warpSettings waitForShutdown conf++  putStrLn $ ("Listening on port " :: Text) <> show (configPort conf)++  let shutdown = putErrLn ("Broadcaster connection is dead" :: Text) >> putMVar shutdownSignal ()+  pool <- P.acquire (configPool conf, 10, pgSettings)+  multi <- newHasqlBroadcaster shutdown listenChannel pgSettings+  getTime <- mkGetTime++  runSettings appSettings $+    postgresWsMiddleware getTime listenChannel (configJwtSecret conf) pool multi $+    logStdout $ maybe dummyApp staticApp' (configPath conf)++  where+    mkGetTime :: IO (IO UTCTime)+    mkGetTime = mkAutoUpdate defaultUpdateSettings {updateAction = getCurrentTime}+    staticApp' :: Text -> Application+    staticApp' = staticApp . defaultFileServerSettings . toS+    dummyApp :: Application+    dummyApp _ respond =+        respond $ responseLBS status200 [("Content-Type", "text/plain")] "Hello, Web!"
test/ClaimsSpec.hs view
@@ -4,20 +4,51 @@  import qualified Data.HashMap.Strict as M import           Test.Hspec-import           Data.Aeson          (Value (..) )+import           Data.Aeson          (Value (..),  toJSON) import Data.Time.Clock import           PostgresWebsockets.Claims +secret :: ByteString+secret = "reallyreallyreallyreallyverysafe"+ spec :: Spec spec =   describe "validate claims" $ do     it "should invalidate an expired token" $ do       time <- getCurrentTime-      validateClaims Nothing "reallyreallyreallyreallyverysafe"-                   "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0IiwiZXhwIjoxfQ.4rDYiMZFR2WHB7Eq4HMdvDP_BQZVtHIfyJgy0NshbHY" time-                   `shouldReturn` Left "Token expired"-    it "should succeed using a matching token" $ do+      validateClaims Nothing secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0IiwiZXhwIjoxfQ.4rDYiMZFR2WHB7Eq4HMdvDP_BQZVtHIfyJgy0NshbHY" time+        `shouldReturn` Left "Token expired"+    it "request any channel from a token that does not have channels or channel claims should succeed" $ do       time <- getCurrentTime-      validateClaims Nothing "reallyreallyreallyreallyverysafe"-                   "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0In0.1d4s-at2kWj8OSabHZHTbNh1dENF7NWy_r0ED3Rwf58" time-                   `shouldReturn` Right ("test", "r", M.fromList[("mode",String "r"),("channel",String "test")])+      validateClaims (Just (encodeUtf8 "test")) secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciJ9.jL5SsRFegNUlbBm8_okhHSujqLcKKZdDglfdqNl1_rY" time+        `shouldReturn` Right (["test"], "r", M.fromList[("mode",String "r")])+    it "requesting a channel that is set by and old style channel claim should work" $ do+      time <- getCurrentTime+      validateClaims (Just (encodeUtf8 "test")) secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0In0.1d4s-at2kWj8OSabHZHTbNh1dENF7NWy_r0ED3Rwf58" time+        `shouldReturn` Right (["test"], "r", M.fromList[("mode",String "r"),("channel",String "test")])+    it "no requesting channel should return all channels in the token" $ do+      time <- getCurrentTime+      validateClaims Nothing secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWxzIjpbInRlc3QiLCJhbm90aGVyIHRlc3QiXX0.b9N8J8tPOPIxxFj5WJ7sWrmcL8ib63i8eirsRZTM9N0" time+        `shouldReturn` Right (["test", "another test"], "r", M.fromList[("mode",String "r"),("channels",  toJSON["test"::Text, "another test"::Text] )        ])++    it "requesting a channel from the channels claim shoud return only the requested channel" $ do+      time <- getCurrentTime+      validateClaims (Just (encodeUtf8 "test")) secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWxzIjpbInRlc3QiLCJ0ZXN0MiJdfQ.MumdJ5FpFX4Z6SJD3qsygVF0r9vqxfqhj5J30O32N0k" time+        `shouldReturn` Right (["test"], "r", M.fromList[("mode",String "r"),("channels",  toJSON ["test"::Text, "test2"] )])++    it "requesting a channel not from the channels claim shoud error" $ do+      time <- getCurrentTime+      validateClaims (Just (encodeUtf8 "notAllowed")) secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWxzIjpbInRlc3QiLCJ0ZXN0MiJdfQ.MumdJ5FpFX4Z6SJD3qsygVF0r9vqxfqhj5J30O32N0k" time+        `shouldReturn` Left  "No allowed channels"++    it "requesting a channel with no mode fails" $ do+      time <- getCurrentTime+      validateClaims (Just (encodeUtf8 "test")) secret+        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJjaGFubmVscyI6WyJ0ZXN0IiwidGVzdDIiXX0.akC1PEYk2DEZtLP2XjC6qXOGZJejmPx49qv-VeEtQYQ" time+        `shouldReturn` Left "Missing mode"
+ test/ServerSpec.hs view
@@ -0,0 +1,117 @@+module ServerSpec (spec) where++import Protolude++import Test.Hspec+import PostgresWebsockets+import PostgresWebsockets.Config++import Control.Lens+import Data.Aeson.Lens++import qualified Network.WebSockets as WS+import           Network.Socket (withSocketsDo)++testServerConfig :: AppConfig+testServerConfig = AppConfig +                    { configDatabase = "postgres://localhost/postgres"+                    , configPath = Nothing+                    , configHost = "*"+                    , configPort = 8080+                    , configListenChannel = "postgres-websockets-test-channel"+                    , configJwtSecret = "reallyreallyreallyreallyverysafe"+                    , configJwtSecretIsBase64 = False+                    , configPool = 10+                    }++startTestServer :: IO ThreadId+startTestServer = do+    threadId <- forkIO $ serve testServerConfig+    threadDelay 1000+    pure threadId++withServer :: IO () -> IO ()+withServer action =+  bracket startTestServer+          killThread+          (const action)++sendWsData :: Text -> Text -> IO ()+sendWsData uri msg =+    withSocketsDo $ +        WS.runClient +            "localhost" +            (configPort testServerConfig) +            (toS uri) +            (`WS.sendTextData` msg)++testChannel :: Text+testChannel = "/test/eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoicncifQ.auy9z4-pqoVEAay9oMi1FuG7ux_C_9RQCH8-wZgej18"++secondaryChannel :: Text+secondaryChannel = "/secondary/eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoicncifQ.auy9z4-pqoVEAay9oMi1FuG7ux_C_9RQCH8-wZgej18"++testAndSecondaryChannel :: Text+testAndSecondaryChannel = "/eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoicnciLCJjaGFubmVscyI6WyJ0ZXN0Iiwic2Vjb25kYXJ5Il19.7tB2A9MhpY4tyqhfnHNy5FUYw4gwpKtL4UAHBXbNEz4"++waitForWsData :: Text -> IO (MVar ByteString)+waitForWsData uri = do+    msg <- newEmptyMVar+    void $ forkIO $+        withSocketsDo $ +            WS.runClient +                "localhost" +                (configPort testServerConfig) +                (toS uri) +                (\c -> do+                    m <- WS.receiveData c+                    putMVar msg m+                )+    threadDelay 1000+    pure msg++waitForMultipleWsData :: Int -> Text -> IO (MVar [ByteString])+waitForMultipleWsData messageCount uri = do+    msg <- newEmptyMVar+    void $ forkIO $+        withSocketsDo $ +            WS.runClient +                "localhost" +                (configPort testServerConfig) +                (toS uri) +                (\c -> do+                    m <- replicateM messageCount (WS.receiveData c)+                    putMVar msg m+                )+    threadDelay 1000+    pure msg++spec :: Spec+spec = around_ withServer $+            describe "serve" $ do+                it "should be able to send messages to test server" $+                    sendWsData testChannel "test data"+                it "should be able to receive messages from test server" $ do+                    msg <- waitForWsData testChannel+                    sendWsData testChannel "test data"+                    msgJson <- takeMVar msg+                    (msgJson ^? key "payload" . _String) `shouldBe` Just "test data"+                it "should be able to send messages to multiple channels in one shot" $ do+                    msg <- waitForWsData testChannel+                    secondaryMsg <- waitForWsData secondaryChannel+                    sendWsData testAndSecondaryChannel "test data"+                    msgJson <- takeMVar msg+                    secondaryMsgJson <- takeMVar secondaryMsg++                    (msgJson ^? key "payload" . _String) `shouldBe` Just "test data"+                    (msgJson ^? key "channel" . _String) `shouldBe` Just "test"+                    (secondaryMsgJson ^? key "payload" . _String) `shouldBe` Just "test data"+                    (secondaryMsgJson ^? key "channel" . _String) `shouldBe` Just "secondary"+                it "should be able to receive from multiple channels in one shot" $ do+                    msgs <- waitForMultipleWsData 2 testAndSecondaryChannel+                    sendWsData testAndSecondaryChannel "test data"+                    msgsJson <- takeMVar msgs++                    forM_ +                        msgsJson +                        (\msgJson -> (msgJson ^? key "payload" . _String) `shouldBe` Just "test data")