diff --git a/app/Config.hs b/app/Config.hs --- a/app/Config.hs +++ b/app/Config.hs @@ -31,7 +31,7 @@ import Options.Applicative hiding (str) import Paths_postgrest_ws (version) import Text.Heredoc -import Text.PrettyPrint.ANSI.Leijen hiding ((<>), (<$>)) +import Text.PrettyPrint.ANSI.Leijen hiding ((<>)) import qualified Text.PrettyPrint.ANSI.Leijen as L import Protolude hiding (intercalate, (<>)) diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -13,7 +13,6 @@ import Data.String (IsString (..)) import Data.Text (pack, replace, strip, stripPrefix) import Data.Text.Encoding (encodeUtf8, decodeUtf8) -import Data.Time.Clock.POSIX (getPOSIXTime) import qualified Hasql.Query as H import qualified Hasql.Session as H import qualified Hasql.Decoders as HD @@ -57,7 +56,7 @@ multi <- newHasqlBroadcaster pgSettings runSettings appSettings $ - postgrestWsMiddleware (toS <$> configAuditChannel conf) (configJwtSecret conf) getPOSIXTime pool multi $ + postgrestWsMiddleware (toS <$> configAuditChannel conf) (configJwtSecret conf) pool multi $ logStdout $ staticApp $ defaultFileServerSettings $ toS $ configPath conf loadSecretFile :: AppConfig -> IO AppConfig diff --git a/postgrest-ws.cabal b/postgrest-ws.cabal --- a/postgrest-ws.cabal +++ b/postgrest-ws.cabal @@ -1,5 +1,5 @@ name: postgrest-ws -version: 0.4.1.0 +version: 0.4.2.0 synopsis: PostgREST extension to map LISTEN/NOTIFY messages to Websockets description: Please see README.md homepage: https://github.com/diogob/postgrest-ws#readme @@ -47,7 +47,7 @@ , time , contravariant default-language: Haskell2010 - default-extensions: OverloadedStrings, NoImplicitPrelude + default-extensions: OverloadedStrings, NoImplicitPrelude, LambdaCase executable postgrest-ws hs-source-dirs: app @@ -59,10 +59,7 @@ , hasql >= 0.19 , hasql-pool >= 0.4 , warp >= 3.2 && < 4 - , unix >= 2.7 && < 3 - , jwt >= 0.7 && < 1 , postgrest-ws - , postgresql-libpq >= 0.9 && < 1.0 , protolude >= 0.2 , base64-bytestring , bytestring @@ -74,9 +71,7 @@ , wai-extra , wai-app-static , heredoc - , auto-update , ansi-wl-pprint - , http-types default-language: Haskell2010 default-extensions: OverloadedStrings, NoImplicitPrelude, QuasiQuotes diff --git a/src/PostgRESTWS.hs b/src/PostgRESTWS.hs --- a/src/PostgRESTWS.hs +++ b/src/PostgRESTWS.hs @@ -17,12 +17,11 @@ import Protolude import qualified Data.Aeson as A -import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BS import qualified Data.ByteString.Lazy as BL import qualified Data.HashMap.Strict as M import qualified Data.Text.Encoding.Error as T -import Data.Time.Clock.POSIX (POSIXTime) - +import Data.Time.Clock.POSIX (getPOSIXTime) import PostgRESTWS.Broadcast (Multiplexer, onMessage) import qualified PostgRESTWS.Broadcast as B import PostgRESTWS.Claims @@ -38,26 +37,32 @@ instance A.ToJSON Message -- | Given a secret, a function to fetch the system time, a Hasql Pool and a Multiplexer this will give you a WAI middleware. -postgrestWsMiddleware :: Maybe ByteString -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application +postgrestWsMiddleware :: Maybe ByteString -> ByteString -> H.Pool -> Multiplexer -> Wai.Application -> Wai.Application postgrestWsMiddleware = WS.websocketsOr WS.defaultConnectionOptions `compose` wsApp where - compose = (.) . (.) . (.) . (.) . (.) + compose = (.) . (.) . (.) . (.) -- private functions -- when the websocket is closed a ConnectionClosed Exception is triggered -- this kills all children and frees resources for us -wsApp :: Maybe ByteString -> ByteString -> IO POSIXTime -> H.Pool -> Multiplexer -> WS.ServerApp -wsApp mAuditChannel secret getTime pool multi pendingConn = do - validateClaims secret (toS jwtToken) >>= either rejectRequest forkSessions +wsApp :: Maybe ByteString -> ByteString -> H.Pool -> Multiplexer -> WS.ServerApp +wsApp mAuditChannel secret pool multi pendingConn = + validateClaims requestChannel secret (toS jwtToken) >>= either rejectRequest forkSessions where hasRead m = m == ("r" :: ByteString) || m == ("rw" :: ByteString) hasWrite m = m == ("w" :: ByteString) || m == ("rw" :: ByteString) rejectRequest = WS.rejectRequest pendingConn . encodeUtf8 - -- the first char in path is '/' the rest is the token - jwtToken = BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn - notifySessionWithTime = notifySession getTime + -- the URI has one of the two formats - /:jwt or /:channel/:jwt + pathElements = BS.split '/' $ BS.drop 1 $ WS.requestPath $ WS.pendingRequest pendingConn + jwtToken + | length pathElements > 1 = headDef "" $ tailSafe pathElements + | length pathElements <= 1 = headDef "" pathElements + requestChannel + | length pathElements > 1 = Just $ headDef "" pathElements + | length pathElements <= 1 = Nothing + notifySessionWithTime = notifySession forkSessions (channel, mode, validClaims) = do -- role claim defaults to anon if not specified in jwt -- We should accept only after verifying JWT @@ -82,12 +87,11 @@ -- 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 :: IO POSIXTime - -> A.Object +notifySession :: A.Object -> WS.Connection -> (ByteString -> IO ()) -> IO () -notifySession getTime claimsToSend wsCon send = +notifySession claimsToSend wsCon send = withAsync (forever relayData) wait where relayData = jsonMsgWithTime >>= send @@ -100,5 +104,5 @@ claimsWithTime :: IO (M.HashMap Text A.Value) claimsWithTime = do - time <- getTime + time <- getPOSIXTime return $ M.insert "message_delivered_at" (A.Number $ fromRational $ toRational time) claimsToSend diff --git a/src/PostgRESTWS/Broadcast.hs b/src/PostgRESTWS/Broadcast.hs --- a/src/PostgRESTWS/Broadcast.hs +++ b/src/PostgRESTWS/Broadcast.hs @@ -107,10 +107,9 @@ then M.insert Channel{ broadcast = broadcast c, listeners = listeners c - 1, close = close c} chan (channels multi) else closeChannelProducer multi chan openChannelWhenNotFound = - M.lookup chan (channels multi) >>= \mC -> - case mC of - Nothing -> openChannel multi chan - Just ch -> return ch + M.lookup chan (channels multi) >>= \case + Nothing -> openChannel multi chan + Just ch -> return ch addListener ch = do M.delete chan (channels multi) let newChannel = Channel{ broadcast = broadcast ch, listeners = listeners ch + 1, close = close ch} diff --git a/src/PostgRESTWS/Claims.hs b/src/PostgRESTWS/Claims.hs --- a/src/PostgRESTWS/Claims.hs +++ b/src/PostgRESTWS/Claims.hs @@ -21,23 +21,27 @@ {-| 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 :: ByteString -> LByteString -> IO (Either Text ConnectionInfo) -validateClaims secret jwtToken = +validateClaims :: Maybe ByteString -> ByteString -> LByteString -> IO (Either Text ConnectionInfo) +validateClaims requestChannel secret jwtToken = runExceptT $ do cl <- liftIO $ jwtClaims (parseJWK secret) jwtToken cl' <- case cl of JWTClaims c -> pure c _ -> throwError "Error" - channel <- claimAsJSON "channel" cl' - mode <- claimAsJSON "mode" cl' + channel <- claimAsJSON requestChannel "channel" cl' + mode <- claimAsJSON Nothing "mode" cl' pure (channel, mode, cl') where - claimAsJSON :: Text -> Claims -> ExceptT Text IO ByteString - claimAsJSON name cl = case M.lookup name cl of + claimAsJSON :: Maybe ByteString -> Text -> Claims -> ExceptT Text IO ByteString + claimAsJSON defaultVal name cl = case M.lookup name cl of Just (String s) -> pure $ encodeUtf8 s Just _ -> throwError "claim is not string value" - Nothing -> throwError (name <> " not in claims") + Nothing -> nonExistingClaim defaultVal name + + nonExistingClaim :: Maybe ByteString -> Text -> ExceptT Text IO ByteString + nonExistingClaim Nothing name = throwError (name <> " not in claims") + nonExistingClaim (Just defaultVal) _ = pure defaultVal {- Private functions and types copied from postgrest diff --git a/test/ClaimsSpec.hs b/test/ClaimsSpec.hs --- a/test/ClaimsSpec.hs +++ b/test/ClaimsSpec.hs @@ -12,6 +12,6 @@ spec = describe "validate claims" $ it "should succeed using a matching token" - $ validateClaims "reallyreallyreallyreallyverysafe" + $ validateClaims Nothing "reallyreallyreallyreallyverysafe" "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJtb2RlIjoiciIsImNoYW5uZWwiOiJ0ZXN0In0.1d4s-at2kWj8OSabHZHTbNh1dENF7NWy_r0ED3Rwf58" `shouldReturn` Right ("test", "r", M.fromList[("mode",String "r"),("channel",String "test")])