diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -19,6 +19,11 @@ 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) @@ -47,25 +52,35 @@ <> " / 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 listenChannel pgSettings + multi <- newHasqlBroadcaster shutdown listenChannel pgSettings + getTime <- mkGetTime runSettings appSettings $ - postgresWsMiddleware listenChannel (configJwtSecret conf) pool multi $ + 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 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.6.1.1 +version: 0.7.0.0 synopsis: Middleware to map LISTEN/NOTIFY messages to Websockets description: Please see README.md homepage: https://github.com/diogob/postgres-websockets#readme @@ -38,12 +38,13 @@ , hasql >= 1.4.1 , hasql-notifications >= 0.1.0.0 && < 0.2 , either >= 5.0.1.1 && < 5.1 - , stm-containers + , stm-containers >= 1.1.0.2 && < 1.2 , stm >= 2.5.0.0 && < 2.6 , retry >= 0.8.1.0 && < 0.9 , stringsearch >= 0.3.6.6 && < 0.4 , time >= 1.8.0.2 && < 1.9 , contravariant >= 1.5.2 && < 1.6 + , alarmclock >= 0.7.0.2 && < 0.8 default-language: Haskell2010 default-extensions: OverloadedStrings, NoImplicitPrelude, LambdaCase @@ -59,16 +60,17 @@ , hasql-pool >= 0.4 , warp >= 3.2 && < 4 , postgres-websockets - , protolude >= 0.2 - , base64-bytestring - , bytestring - , text + , 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 - , wai-app-static - , http-types + , 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 @@ -80,20 +82,20 @@ , ClaimsSpec , HasqlBroadcastSpec build-depends: base - , protolude >= 0.2 + , protolude >= 0.2.3 , postgres-websockets - , containers - , hspec - , hspec-wai - , hspec-wai-json - , aeson - , hasql - , hasql-pool + , hspec >= 2.7.1 && < 2.8 + , hspec-wai >= 0.9.2 && < 0.10 + , hspec-wai-json >= 0.9.2 && < 0.10 + , aeson >= 1.4.6.0 && < 1.5 + , hasql >= 0.19 + , hasql-pool >= 0.4 , hasql-notifications >= 0.1.0.0 && < 0.2 - , http-types + , http-types >= 0.9 + , time >= 1.8.0.2 && < 1.9 , unordered-containers >= 0.2 - , wai-extra - , stm + , wai-extra >= 3.0.29 && < 3.1 + , stm >= 2.5.0.0 && < 2.6 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 @@ -22,7 +22,9 @@ 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 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 @@ -38,19 +40,21 @@ 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 :: Text -> ByteString -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application +postgresWsMiddleware :: IO UTCTime -> Text -> ByteString -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application postgresWsMiddleware = WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp where - compose = (.) . (.) . (.) . (.) + 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 :: Text -> ByteString -> H.Pool -> Multiplexer -> WS.ServerApp -wsApp dbChannel secret pool multi pendingConn = - validateClaims requestChannel secret (toS jwtToken) >>= either rejectRequest forkSessions +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) @@ -68,17 +72,23 @@ -- 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 + 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 (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 sendNotifications + when (hasWrite mode) $ + let sendNotifications = void . H.notifyPool pool dbChannel . toS + in notifySession validClaims (toS ch) conn getTime sendNotifications - waitForever <- newEmptyMVar - void $ takeMVar waitForever + 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 @@ -86,9 +96,10 @@ notifySession :: A.Object -> Text -> WS.Connection + -> IO UTCTime -> (ByteString -> IO ()) -> IO () -notifySession claimsToSend ch wsCon send = +notifySession claimsToSend ch wsCon getTime send = withAsync (forever relayData) wait where relayData = jsonMsgWithTime >>= send @@ -102,5 +113,5 @@ claimsWithChannel = M.insert "channel" (A.String ch) claimsToSend claimsWithTime :: IO (M.HashMap Text A.Value) claimsWithTime = do - time <- getPOSIXTime - return $ M.insert "message_delivered_at" (A.Number $ fromRational $ toRational time) claimsWithChannel + time <- utcTimeToPOSIXSeconds <$> getTime + return $ M.insert "message_delivered_at" (A.Number $ realToFrac time) claimsWithChannel diff --git a/src/PostgresWebsockets/Claims.hs b/src/PostgresWebsockets/Claims.hs --- a/src/PostgresWebsockets/Claims.hs +++ b/src/PostgresWebsockets/Claims.hs @@ -10,23 +10,27 @@ 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 +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 Value +type Claims = M.HashMap Text JSON.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 = +validateClaims :: Maybe ByteString -> ByteString -> LByteString -> UTCTime -> IO (Either Text ConnectionInfo) +validateClaims requestChannel secret jwtToken time = runExceptT $ do - cl <- liftIO $ jwtClaims (parseJWK secret) jwtToken + 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' @@ -35,7 +39,7 @@ 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 (JSON.String s) -> pure $ encodeUtf8 s Just _ -> throwError "claim is not string value" Nothing -> nonExistingClaim defaultVal name @@ -53,20 +57,20 @@ -} data JWTAttempt = JWTInvalid JWTError | JWTMissingSecret - | JWTClaims (M.HashMap Text Value) + | JWTClaims (M.HashMap Text JSON.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) +jwtClaims :: UTCTime -> JWK -> LByteString -> IO JWTAttempt +jwtClaims _ _ "" = return $ JWTClaims M.empty +jwtClaims time jwk payload = do + let config = defaultJWTValidationSettings (const True) eJwt <- runExceptT $ do jwt <- decodeCompact payload - verifyClaims validation secret jwt + verifyClaimsAt config jwk time jwt return $ case eJwt of Left e -> JWTInvalid e Right jwt -> JWTClaims . claims2map $ jwt @@ -75,10 +79,10 @@ Internal helper used to turn JWT ClaimSet into something easier to work with -} -claims2map :: ClaimsSet -> M.HashMap Text Value -claims2map = val2map . toJSON +claims2map :: ClaimsSet -> M.HashMap Text JSON.Value +claims2map = val2map . JSON.toJSON where - val2map (Object o) = o + val2map (JSON.Object o) = o val2map _ = M.empty {-| @@ -96,4 +100,4 @@ parseJWK :: ByteString -> JWK parseJWK str = - fromMaybe (hs256jwk str) (decode (toS str) :: Maybe JWK) + fromMaybe (hs256jwk str) (JSON.decode (toS str) :: Maybe JWK) diff --git a/src/PostgresWebsockets/HasqlBroadcast.hs b/src/PostgresWebsockets/HasqlBroadcast.hs --- a/src/PostgresWebsockets/HasqlBroadcast.hs +++ b/src/PostgresWebsockets/HasqlBroadcast.hs @@ -26,19 +26,19 @@ {- | 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 :: Text -> ByteString -> IO Multiplexer -newHasqlBroadcaster ch = newHasqlBroadcasterForConnection . tryUntilConnected +newHasqlBroadcaster :: IO () -> Text -> ByteString -> IO Multiplexer +newHasqlBroadcaster onConnectionFailure ch = newHasqlBroadcasterForConnection . tryUntilConnected where - newHasqlBroadcasterForConnection = newHasqlBroadcasterForChannel ch + newHasqlBroadcasterForConnection = newHasqlBroadcasterForChannel onConnectionFailure ch {- | 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 :: Text -> ByteString -> IO (Either ByteString Multiplexer) -newHasqlBroadcasterOrError ch = +newHasqlBroadcasterOrError :: IO () -> Text -> ByteString -> IO (Either ByteString Multiplexer) +newHasqlBroadcasterOrError onConnectionFailure ch = acquire >=> (sequence . mapBoth show (newHasqlBroadcasterForConnection . return)) where - newHasqlBroadcasterForConnection = newHasqlBroadcasterForChannel ch + newHasqlBroadcasterForConnection = newHasqlBroadcasterForChannel onConnectionFailure ch tryUntilConnected :: ByteString -> IO Connection tryUntilConnected = @@ -78,13 +78,12 @@ @ -} -newHasqlBroadcasterForChannel :: Text -> IO Connection -> IO Multiplexer -newHasqlBroadcasterForChannel ch getCon = do - multi <- newMultiplexer openProducer closeProducer +newHasqlBroadcasterForChannel :: IO () -> Text -> IO Connection -> IO Multiplexer +newHasqlBroadcasterForChannel onConnectionFailure ch getCon = do + multi <- newMultiplexer openProducer $ const onConnectionFailure void $ relayMessagesForever multi return multi where - closeProducer _ = putErrLn "Broadcaster is dead" toMsg :: ByteString -> ByteString -> Message toMsg c m = case decode (toS m) of Just v -> Message (channelDef c v) m diff --git a/test/ClaimsSpec.hs b/test/ClaimsSpec.hs --- a/test/ClaimsSpec.hs +++ b/test/ClaimsSpec.hs @@ -5,13 +5,19 @@ import qualified Data.HashMap.Strict as M import Test.Hspec import Data.Aeson (Value (..) ) - +import Data.Time.Clock import PostgresWebsockets.Claims spec :: Spec spec = - describe "validate claims" - $ it "should succeed using a matching token" - $ validateClaims Nothing "reallyreallyreallyreallyverysafe" - "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0In0.1d4s-at2kWj8OSabHZHTbNh1dENF7NWy_r0ED3Rwf58" + 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 + time <- getCurrentTime + validateClaims Nothing "reallyreallyreallyreallyverysafe" + "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0In0.1d4s-at2kWj8OSabHZHTbNh1dENF7NWy_r0ED3Rwf58" time `shouldReturn` Right ("test", "r", M.fromList[("mode",String "r"),("channel",String "test")]) diff --git a/test/HasqlBroadcastSpec.hs b/test/HasqlBroadcastSpec.hs --- a/test/HasqlBroadcastSpec.hs +++ b/test/HasqlBroadcastSpec.hs @@ -15,7 +15,7 @@ <$> acquire connStr it "relay messages sent to the appropriate database channel" $ do - multi <- either (panic .show) id <$> newHasqlBroadcasterOrError "postgres-websockets" "postgres://localhost/postgres_ws_test" + multi <- either (panic .show) id <$> newHasqlBroadcasterOrError (pure ()) "postgres-websockets" "postgres://localhost/postgres_ws_test" msg <- liftIO newEmptyMVar onMessage multi "test" $ putMVar msg