diff --git a/app/Config.hs b/app/Config.hs deleted file mode 100644 --- a/app/Config.hs +++ /dev/null @@ -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" diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -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 diff --git a/postgres-websockets.cabal b/postgres-websockets.cabal --- a/postgres-websockets.cabal +++ b/postgres-websockets.cabal @@ -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 diff --git a/src/PostgresWebsockets.hs b/src/PostgresWebsockets.hs --- a/src/PostgresWebsockets.hs +++ b/src/PostgresWebsockets.hs @@ -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 diff --git a/src/PostgresWebsockets/Broadcast.hs b/src/PostgresWebsockets/Broadcast.hs --- a/src/PostgresWebsockets/Broadcast.hs +++ b/src/PostgresWebsockets/Broadcast.hs @@ -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 (..) diff --git a/src/PostgresWebsockets/Claims.hs b/src/PostgresWebsockets/Claims.hs --- a/src/PostgresWebsockets/Claims.hs +++ b/src/PostgresWebsockets/Claims.hs @@ -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 diff --git a/src/PostgresWebsockets/Config.hs b/src/PostgresWebsockets/Config.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/Config.hs @@ -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 "." "=" + diff --git a/src/PostgresWebsockets/HasqlBroadcast.hs b/src/PostgresWebsockets/HasqlBroadcast.hs --- a/src/PostgresWebsockets/HasqlBroadcast.hs +++ b/src/PostgresWebsockets/HasqlBroadcast.hs @@ -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 diff --git a/src/PostgresWebsockets/Middleware.hs b/src/PostgresWebsockets/Middleware.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/Middleware.hs @@ -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 diff --git a/src/PostgresWebsockets/Server.hs b/src/PostgresWebsockets/Server.hs new file mode 100644 --- /dev/null +++ b/src/PostgresWebsockets/Server.hs @@ -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!" diff --git a/test/ClaimsSpec.hs b/test/ClaimsSpec.hs --- a/test/ClaimsSpec.hs +++ b/test/ClaimsSpec.hs @@ -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" diff --git a/test/ServerSpec.hs b/test/ServerSpec.hs new file mode 100644 --- /dev/null +++ b/test/ServerSpec.hs @@ -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")