diff --git a/client/API.hs b/client/API.hs --- a/client/API.hs +++ b/client/API.hs @@ -1,139 +1,86 @@ --- | Capability record for the NATS client surface. +-- | Stable public surface for the core NATS client. +-- +-- Constructors are intentionally hidden. This lets clients, messages, and +-- subscriptions acquire new capabilities without changing their public +-- representation. module API - ( - Client (..) - , MsgView (..) + ( Client + , Message + , MsgView + , Subscription + , subscriptionSid + , NatsError (..) + , Subject + , Payload + , Headers , PublishOption , SubscribeOption + , RequestOption + , UnsubscribeOption + , PingOption + , FlushOption + , ResetOption + , CloseOption + , publish + , subscribe + , subscribeOnce + , request + , unsubscribe + , newInbox + , ping + , flush + , reset + , close + , subject + , sid + , replyTo + , payload + , headers , withSubscriptionExpiry , withQueueGroup - , withPayload - , withReplyCallback , withReplyTo , withHeaders + , withRequestTimeout + , withRequestHeaders ) where -import qualified Data.ByteString as BS -import Data.Time.Clock (NominalDiffTime) -import Lib.CallOption (CallOption) -import Publish.Config (PublishConfig) -import Subscription.Types (SubscribeConfig (..)) -import qualified Types.Msg as Msg -import Types.Msg (Headers, Payload, SID, Subject) - --- | Client capabilities for publishing, subscribing, and lifecycle control. -data Client = Client - { publish :: Subject -> [PublishOption] -> IO () - -- ^ Publish a message, optionally overriding publish options. - , subscribe :: Subject -> [SubscribeOption] -> (Maybe MsgView -> IO ()) -> IO SID - -- ^ Subscribe to a subject and handle delivered messages. - , request :: Subject -> [SubscribeOption] -> (Maybe MsgView -> IO ()) -> IO SID - -- ^ Subscribe with request semantics and auto-unsubscribe after a reply. - , unsubscribe :: SID -> IO () - -- ^ Unsubscribe from a subscription by SID. - , newInbox :: IO Subject - -- ^ Create a unique inbox subject for replies. - , ping :: IO () -> IO () - -- ^ Send a ping and run the callback when a pong arrives. - , flush :: IO () - -- ^ Flush buffered writes to the server. - , reset :: IO () - -- ^ Reset the client connection state. - , close :: IO () - -- ^ Close the client connection and release resources. - } - --- | MsgView represents a MSG in the NATS protocol. -data MsgView = MsgView - { -- | The subject of the message. - subject :: BS.ByteString - -- | The SID (subscription ID) of the message. - , sid :: BS.ByteString - -- | The replyTo subject, if any. - , replyTo :: Maybe BS.ByteString - -- | The payload of the message, if any. - , payload :: Maybe BS.ByteString - -- | Headers associated with the message, if any. - , headers :: Maybe [(BS.ByteString, BS.ByteString)] - } - deriving (Eq, Show) - -type PublishOption = CallOption PublishConfig - -type SubscribeOption = CallOption SubscribeConfig - --- | withSubscriptionExpiry sets the reply subscription expiry in seconds. --- Default: no expiry (reply subscriptions stay open until unsubscribe). --- --- __Examples:__ --- --- @ --- {-# LANGUAGE OverloadedStrings #-} --- --- subscribe client \"events.created\" [withSubscriptionExpiry 2] print --- @ -withSubscriptionExpiry :: NominalDiffTime -> SubscribeOption -withSubscriptionExpiry expirySeconds cfg = cfg { expiry = Just expirySeconds } - --- | withQueueGroup sets the queue group for a subscription. --- Default: no queue group. -withQueueGroup :: Subject -> SubscribeOption -withQueueGroup queueGroup cfg = cfg { subscribeQueueGroup = Just queueGroup } - --- | withPayload is used to set the payload for a publish operation. --- Default: no payload. --- --- __Examples:__ --- --- @ --- {-# LANGUAGE OverloadedStrings #-} --- --- publish client \"updates\" [withPayload \"hello\"] --- @ -withPayload :: Payload -> PublishOption -withPayload payload (_, callback, headers, replyTo') = - (Just payload, callback, headers, replyTo') - --- | withReplyCallback is used to set a callback for a reply to a publish operation. --- Default: no reply subscription; publishes are fire-and-forget. --- --- __Examples:__ --- --- @ --- {-# LANGUAGE OverloadedStrings #-} --- --- publish client \"service.echo\" [withReplyCallback print] --- @ -withReplyCallback :: (Maybe MsgView -> IO ()) -> PublishOption -withReplyCallback callback (payload, _, headers, replyTo') = - (payload, Just (callback . fmap (\msg -> MsgView - { subject = Msg.subject msg - , sid = Msg.sid msg - , replyTo = Msg.replyTo msg - , payload = Msg.payload msg - , headers = Msg.headers msg - })), headers, replyTo') - --- | withReplyTo sets an explicit reply subject for a publish operation. --- Default: no reply subject unless a reply callback is configured. --- --- This option only controls the outgoing reply subject. It does not create a --- subscription; callers that expect multiple replies should subscribe to the --- reply subject themselves. -withReplyTo :: Subject -> PublishOption -withReplyTo replySubject (payload, callback, headers, _) = - (payload, callback, headers, Just replySubject) - --- | withHeaders is used to set headers for a publish operation. --- Default: no headers. --- --- __Examples:__ --- --- @ --- {-# LANGUAGE OverloadedStrings #-} --- --- publish client \"updates\" [withHeaders [(\"source\", \"test\")]] --- @ -withHeaders :: Headers -> PublishOption -withHeaders headers (payload, callback, _, replyTo') = - (payload, callback, Just headers, replyTo') +import Client.API + ( Client + , CloseOption + , FlushOption + , Headers + , Message + , MsgView + , NatsError (..) + , Payload + , PingOption + , PublishOption + , RequestOption + , ResetOption + , Subject + , SubscribeOption + , Subscription + , UnsubscribeOption + , close + , flush + , headers + , newInbox + , payload + , ping + , publish + , replyTo + , request + , reset + , sid + , subject + , subscribe + , subscribeOnce + , subscriptionSid + , unsubscribe + , withHeaders + , withQueueGroup + , withReplyTo + , withRequestHeaders + , withRequestTimeout + , withSubscriptionExpiry + ) diff --git a/client/Client.hs b/client/Client.hs --- a/client/Client.hs +++ b/client/Client.hs @@ -1,483 +1,62 @@ -{-# LANGUAGE OverloadedStrings #-} - --- | High-level client implementation for NATS. +-- | Construct a NATS client satisfying the contract in "API". module Client - ( newClient + ( Server + , ServerConfigError (..) + , server + , serverWithDefaultPort + , serverHost + , serverPort + , connect + , newClient , ConfigOption , withConnectName , withEcho , withAuthToken + , withAuthTokenHandler , withUserPass + , withUserPassHandler , withNKey + , withNKeyHandler , withJWT + , withJWTHandlers + , withTLS , withTLSCert + , withTLSRootCA + , withTLSServerName + , withTLSInsecure , withMinimumLogLevel , withLogAction , withConnectionAttempts + , withConnectTimeoutMicros , withCallbackConcurrency + , withMessageLimit + , withPendingDeliveryLimits + , withErrorHandler , withBufferLimit , withExitAction , LogLevel (..) , LogEntry (..) , renderLogEntry , AuthTokenData + , AuthTokenHandler , UserPassData + , UserPassHandler , NKeyData + , NKeyPublicKey , JWTTokenData + , JWTHandler + , SignatureHandler + , AuthError (..) , TLSPublicKey , TLSPrivateKey , TLSCertData + , TLSConfig (..) , ClientExitReason (..) + , ServerError + , serverErrorReason + , ConnectError (..) + , ConnectAttemptError (..) + , ConnectFailure (..) ) where -import API (Client (..), MsgView (..)) -import qualified Auth.Jwt as AuthJwt -import qualified Auth.NKey as AuthNKey -import qualified Auth.None as AuthNone -import qualified Auth.Token as AuthToken -import Auth.Types - ( Auth - , AuthTokenData - , JWTTokenData - , NKeyData - , UserPassData - ) -import qualified Auth.UserPass as AuthUserPass -import Control.Concurrent (forkIO) -import Control.Concurrent.STM -import Control.Exception (SomeException, displayException) -import Control.Monad (void, when) -import qualified Data.ByteString as BS -import qualified Data.ByteString.Char8 as BC -import Engine (closeClient, resetClient, runEngine) -import Lib.CallOption (CallOption, applyCallOptions) -import Lib.Logger - ( LogEntry (..) - , LogLevel (..) - , LoggerConfig (..) - , MonadLogger (..) - , defaultLogger - , newLogContext - , renderLogEntry - ) -import Network.Connection (connectionApi) -import Network.ConnectionAPI (newConn) -import Parser.Attoparsec (parserApi) -import Pipeline.Broadcasting (broadcastingApi) -import Pipeline.Streaming (streamingApi) -import Publish (defaultPublishConfig) -import Publish.Config (PublishConfig) -import Queue.API (QueueItem (QueueItem)) -import Queue.TransactionalQueue (newQueue) -import State.Store - ( ClientState - , enqueue - , newClientState - , nextInbox - , nextSid - , pushPingAction - , readServerInfo - , readStatus - , runClient - , setConnectName - , waitForClosed - , waitForConnectionGenerationAfter - , waitForNotRunning - ) -import State.Types - ( ClientConfig (..) - , ClientExitReason (..) - , ClientStatus (..) - , TLSCertData - , TLSPrivateKey - , TLSPublicKey - ) -import Subscription.Store - ( SubscriptionStore - , awaitNoTrackedExpiries - , hasTrackedExpiries - , newSubscriptionStore - , register - , startExpiryWorker - , startWorkers - , unregister - ) -import Subscription.Types - ( SubscribeConfig (..) - , SubscriptionMeta (SubscriptionMeta) - ) -import qualified Types.Connect as Connect -import qualified Types.Info as Info -import qualified Types.Msg as Msg -import Types.Ping (Ping (..)) -import qualified Types.Pub as Pub -import qualified Types.Sub as Sub -import qualified Types.Unsub as Unsub -import Validators.Validators (validate) - -data ClientAuth = ClientAuthNone - | ClientAuthToken AuthTokenData - | ClientAuthUserPass UserPassData - | ClientAuthNKey NKeyData - | ClientAuthJWT JWTTokenData - -data ClientOptions = ClientOptions - { optionConnectConfig :: Connect.Connect - , optionAuth :: ClientAuth - , optionTlsCert :: Maybe TLSCertData - , optionLoggerConfig :: LoggerConfig - , optionConnectionAttempts :: Int - , optionCallbackConcurrency :: Int - , optionBufferLimit :: Int - , optionExitAction :: ClientExitReason -> IO () - , optionConnectOptions :: [(String, Int)] - } - -newClient :: [(String, Int)] -> [ConfigOption] -> IO Client -newClient servers configOptions = do - loggerConfig' <- defaultLogger - ctx <- newLogContext - let defaultOptions = applyCallOptions configOptions ClientOptions - { optionConnectConfig = defaultConnect - , optionAuth = ClientAuthNone - , optionTlsCert = Nothing - , optionLoggerConfig = loggerConfig' - , optionConnectionAttempts = 5 - , optionCallbackConcurrency = 1 - , optionBufferLimit = 4096 - , optionExitAction = const (pure ()) - , optionConnectOptions = servers - } - clientConfig = - ClientConfig - { connectionAttempts = optionConnectionAttempts defaultOptions - , callbackConcurrency = optionCallbackConcurrency defaultOptions - , bufferLimit = optionBufferLimit defaultOptions - , connectConfig = optionConnectConfig defaultOptions - , loggerConfig = optionLoggerConfig defaultOptions - , tlsCert = optionTlsCert defaultOptions - , exitAction = optionExitAction defaultOptions - , connectOptions = optionConnectOptions defaultOptions - } - configuredAuth = selectAuth (optionAuth defaultOptions) - - queue <- newQueue - conn <- newConn connectionApi - clientState <- newClientState clientConfig queue conn ctx - store <- newSubscriptionStore - - setConnectName clientState (Connect.name (optionConnectConfig defaultOptions)) - logStaticConfiguration clientState defaultOptions - - startWorkers - (callbackConcurrency clientConfig) - store - (do - waitForClosed clientState - awaitNoTrackedExpiries store) - (handleCallbackError clientState) - - startExpiryWorker store $ - shouldStopExpiryWorker clientState store - - void . forkIO $ - runEngine - connectionApi - streamingApi - broadcastingApi - parserApi - clientState - store - configuredAuth - - atomically $ - waitForConnectionGenerationAfter clientState 0 - `orElse` waitForClosed clientState - - pure Client - { publish = \subject publishOptions -> do - let cfg = applyCallOptions publishOptions defaultPublishConfig - publishClient clientState store subject cfg - , subscribe = \subject subscribeOptions callback -> do - let cfg = applyCallOptions subscribeOptions defaultSubscribeConfig - subscribeClient clientState store False subject cfg (toInternalCallback callback) - , request = \subject subscribeOptions callback -> do - let cfg = applyCallOptions subscribeOptions defaultSubscribeConfig - subscribeClient clientState store True subject cfg (toInternalCallback callback) - , unsubscribe = unsubscribeClient clientState store - , newInbox = nextInbox clientState - , ping = pingClient clientState - , flush = flushClient clientState - , reset = resetClient connectionApi clientState store - , close = closeClient connectionApi clientState store - } - -type ConfigOption = CallOption ClientOptions - -withConnectName :: BS.ByteString -> ConfigOption -withConnectName name config = - config - { optionConnectConfig = - (optionConnectConfig config) { Connect.name = Just name } - } - -withEcho :: Bool -> ConfigOption -withEcho enabled config = - config - { optionConnectConfig = - (optionConnectConfig config) { Connect.echo = Just enabled } - } - -withAuthToken :: AuthTokenData -> ConfigOption -withAuthToken token config = config { optionAuth = ClientAuthToken token } - -withUserPass :: UserPassData -> ConfigOption -withUserPass userPass config = config { optionAuth = ClientAuthUserPass userPass } - -withNKey :: NKeyData -> ConfigOption -withNKey nkey config = config { optionAuth = ClientAuthNKey nkey } - -withJWT :: JWTTokenData -> ConfigOption -withJWT jwt config = config { optionAuth = ClientAuthJWT jwt } - -withTLSCert :: TLSCertData -> ConfigOption -withTLSCert cert config = config { optionTlsCert = Just cert } - -withMinimumLogLevel :: LogLevel -> ConfigOption -withMinimumLogLevel minimumLogLevel config = - config - { optionLoggerConfig = - (optionLoggerConfig config) { minLogLevel = minimumLogLevel } - } - -withLogAction :: (LogEntry -> IO ()) -> ConfigOption -withLogAction logAction config = - config - { optionLoggerConfig = - (optionLoggerConfig config) { logFn = logAction } - } - -withConnectionAttempts :: Int -> ConfigOption -withConnectionAttempts attempts config = - config { optionConnectionAttempts = attempts } - -withCallbackConcurrency :: Int -> ConfigOption -withCallbackConcurrency concurrency config = - config { optionCallbackConcurrency = concurrency } - -withBufferLimit :: Int -> ConfigOption -withBufferLimit limit config = - config { optionBufferLimit = max 1 limit } - -withExitAction :: (ClientExitReason -> IO ()) -> ConfigOption -withExitAction action config = config { optionExitAction = action } - -defaultConnect :: Connect.Connect -defaultConnect = - Connect.Connect - { Connect.verbose = False - , Connect.pedantic = True - , Connect.tls_required = False - , Connect.auth_token = Nothing - , Connect.user = Nothing - , Connect.pass = Nothing - , Connect.name = Nothing - , Connect.lang = "haskell" - , Connect.version = "0.1.0" - , Connect.protocol = Nothing - , Connect.echo = Just True - , Connect.sig = Nothing - , Connect.jwt = Nothing - , Connect.nkey = Nothing - , Connect.no_responders = Just True - , Connect.headers = Just True - } - -defaultSubscribeConfig :: SubscribeConfig -defaultSubscribeConfig = SubscribeConfig Nothing Nothing - -selectAuth :: ClientAuth -> Auth -selectAuth authSelection = - case authSelection of - ClientAuthNone -> - AuthNone.auth - ClientAuthToken token -> - AuthToken.auth token - ClientAuthUserPass userPass -> - AuthUserPass.auth userPass - ClientAuthNKey seed -> - AuthNKey.auth seed - ClientAuthJWT creds -> - AuthJwt.auth creds - -logStaticConfiguration :: ClientState -> ClientOptions -> IO () -logStaticConfiguration client options = - runClient client $ do - case optionAuth options of - ClientAuthNone -> - logMessage Info "no authentication method provided" - ClientAuthToken _ -> - logMessage Info "using auth token" - ClientAuthUserPass (user, _) -> - logMessage Info ("using user/pass: " ++ show user) - ClientAuthNKey _ -> - logMessage Info "using nkey" - ClientAuthJWT _ -> - logMessage Info "using jwt" - case optionTlsCert options of - Nothing -> - pure () - Just _ -> - logMessage Info "using tls certificate" - -handleCallbackError :: ClientState -> SomeException -> IO () -handleCallbackError client err = - runClient client $ - logMessage Error ("callback failed: " ++ displayException err) - -shouldStopExpiryWorker :: ClientState -> SubscriptionStore -> IO Bool -shouldStopExpiryWorker client store = do - status <- readStatus client - tracked <- hasTrackedExpiries store - pure $ - case status of - Closed _ -> not tracked - _ -> False - -toInternalCallback :: (Maybe MsgView -> IO ()) -> Maybe Msg.Msg -> IO () -toInternalCallback callback = - callback . fmap toMsgView - -toMsgView :: Msg.Msg -> MsgView -toMsgView msg = - MsgView - { subject = Msg.subject msg - , sid = Msg.sid msg - , replyTo = Msg.replyTo msg - , payload = Msg.payload msg - , headers = Msg.headers msg - } - -publishClient :: ClientState -> SubscriptionStore -> Msg.Subject -> PublishConfig -> IO () -publishClient client store subject (payload, callback, headers, configuredReplyTo) = do - runClient client $ - logMessage Debug ("publishing to subject: " ++ show subject) - replyTo <- case callback of - Nothing -> pure configuredReplyTo - Just _ -> Just <$> maybe (nextInbox client) pure configuredReplyTo - let publishMessage = - Pub.Pub - { Pub.subject = subject - , Pub.payload = payload - , Pub.replyTo = replyTo - , Pub.headers = headers - } - shouldPublish <- canPublish client publishMessage - when shouldPublish $ do - case callback of - Nothing -> - pure () - Just replyCallback -> do - case replyTo of - Nothing -> - pure () - Just replyInbox -> do - sid <- nextSid client - let meta = - SubscriptionMeta replyInbox Nothing True - register store sid meta defaultSubscribeConfig replyCallback - enqueue client $ - QueueItem - Sub.Sub - { Sub.subject = replyInbox - , Sub.queueGroup = Nothing - , Sub.sid = sid - } - enqueue client $ - QueueItem - Unsub.Unsub - { Unsub.sid = sid - , Unsub.maxMsg = Just 1 - } - enqueue client (QueueItem publishMessage) - -canPublish :: ClientState -> Pub.Pub -> IO Bool -canPublish client publishMessage = - case validate publishMessage of - Left reason -> do - runClient client $ - logMessage Error ("dropping invalid publish: " ++ BC.unpack reason) - pure False - Right () -> do - serverInfo <- readServerInfo client - case serverInfo of - Just info - | Pub.messageSize publishMessage > Info.max_payload info -> do - runClient client $ - logMessage Error - ("dropping publish: payload size " - ++ show (Pub.messageSize publishMessage) - ++ " exceeds server max_payload " - ++ show (Info.max_payload info)) - pure False - _ -> - pure True - -subscribeClient :: ClientState -> SubscriptionStore -> Bool -> Msg.Subject -> SubscribeConfig -> (Maybe Msg.Msg -> IO ()) -> IO Msg.SID -subscribeClient client store isReply subject cfg callback = do - runClient client $ - logMessage Debug ("subscribing to subject: " ++ show subject) - sid <- nextSid client - let queueGroup = - subscribeQueueGroup cfg - subscriptionMessage = - Sub.Sub - { Sub.subject = subject - , Sub.queueGroup = queueGroup - , Sub.sid = sid - } - case validate subscriptionMessage of - Left reason -> - runClient client $ - logMessage Error ("dropping invalid subscription: " ++ BC.unpack reason) - Right () -> do - let meta = - SubscriptionMeta subject queueGroup isReply - register store sid meta cfg callback - enqueue client $ - QueueItem - subscriptionMessage - when isReply $ - enqueue client - (QueueItem - Unsub.Unsub - { Unsub.sid = sid - , Unsub.maxMsg = Just 1 - }) - pure sid - -unsubscribeClient :: ClientState -> SubscriptionStore -> Msg.SID -> IO () -unsubscribeClient client store sid = do - runClient client $ - logMessage Debug ("unsubscribing SID: " ++ show sid) - unregister store sid - enqueue client $ - QueueItem - Unsub.Unsub - { Unsub.sid = sid - , Unsub.maxMsg = Nothing - } - -pingClient :: ClientState -> IO () -> IO () -pingClient client action = do - runClient client $ - logMessage Debug "sending ping to server" - pushPingAction client action - enqueue client (QueueItem Ping) - -flushClient :: ClientState -> IO () -flushClient client = do - ponged <- newEmptyTMVarIO - pingClient client (atomically (void (tryPutTMVar ponged ()))) - atomically $ - readTMVar ponged `orElse` waitForNotRunning client +import Client.Implementation diff --git a/client/Client/Implementation.hs b/client/Client/Implementation.hs new file mode 100644 --- /dev/null +++ b/client/Client/Implementation.hs @@ -0,0 +1,824 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | High-level client implementation for NATS. +module Client.Implementation + ( Server + , ServerConfigError (..) + , server + , serverWithDefaultPort + , serverHost + , serverPort + , connect + , newClient + , ConfigOption + , withConnectName + , withEcho + , withAuthToken + , withAuthTokenHandler + , withUserPass + , withUserPassHandler + , withNKey + , withNKeyHandler + , withJWT + , withJWTHandlers + , withTLS + , withTLSCert + , withTLSRootCA + , withTLSServerName + , withTLSInsecure + , withMinimumLogLevel + , withLogAction + , withConnectionAttempts + , withConnectTimeoutMicros + , withCallbackConcurrency + , withMessageLimit + , withPendingDeliveryLimits + , withErrorHandler + , withBufferLimit + , withExitAction + , LogLevel (..) + , LogEntry (..) + , renderLogEntry + , AuthTokenData + , AuthTokenHandler + , UserPassData + , UserPassHandler + , NKeyData + , NKeyPublicKey + , JWTTokenData + , JWTHandler + , SignatureHandler + , AuthError (..) + , TLSPublicKey + , TLSPrivateKey + , TLSCertData + , TLSConfig (..) + , ClientExitReason (..) + , ServerError + , serverErrorReason + , ConnectError (..) + , ConnectAttemptError (..) + , ConnectFailure (..) + ) where + +import Auth.Config (authMethods, mergeAuth) +import qualified Auth.Jwt as AuthJwt +import qualified Auth.NKey as AuthNKey +import qualified Auth.None as AuthNone +import qualified Auth.Token as AuthToken +import Auth.Types + ( Auth + , AuthError (..) + , AuthTokenData + , AuthTokenHandler + , JWTHandler + , JWTTokenData + , NKeyData + , NKeyPublicKey + , SignatureHandler + , UserPassData + , UserPassHandler + ) +import qualified Auth.UserPass as AuthUserPass +import Client.API + ( Client (..) + , CloseConfig (..) + , FlushConfig (..) + , Message (..) + , NatsError (..) + , PingConfig (..) + , RequestConfig (..) + , ResetConfig (..) + , Subscription (..) + , UnsubscribeConfig (..) + ) +import Control.Concurrent (forkIO) +import Control.Concurrent.STM +import Control.Exception + ( SomeException + , displayException + , mask + , onException + ) +import Control.Monad (forM_, void, when) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Char (isAsciiUpper) +import Data.Maybe (fromMaybe) +import Data.Time.Clock (NominalDiffTime) +import Data.Version (showVersion) +import Engine (closeClient, resetClient, runEngine) +import Lib.CallOption (CallOption, applyCallOptions) +import Lib.Logger + ( LogEntry (..) + , LogLevel (..) + , LoggerConfig (..) + , MonadLogger (..) + , defaultLogger + , newLogContext + , renderLogEntry + ) +import Network.Connection (connectionApi) +import Network.ConnectionAPI (newConn) +import Parser.Attoparsec (parserApiWithMessageLimit) +import qualified Paths_natskell as Package +import Pipeline.Broadcasting (broadcastingApi) +import Pipeline.Streaming (streamingApi) +import Publish (defaultPublishConfig) +import Publish.Config (PublishConfig (..)) +import Queue.API (QueueItem (QueueItem)) +import Queue.TransactionalQueue (newQueue) +import State.Store + ( ClientState + , config + , enqueue + , newClientState + , nextInbox + , nextSid + , pushPingAction + , readServerInfo + , readStatus + , runClient + , setConnectName + , waitForClosed + , waitForInitialConnection + , waitForNotRunning + ) +import State.Types + ( ClientConfig (..) + , ClientExitReason (..) + , ClientStatus (..) + , ConnectAttemptError (..) + , ConnectError (..) + , ConnectFailure (..) + , ServerError + , TLSCertData + , TLSConfig (..) + , TLSPrivateKey + , TLSPublicKey + , defaultTLSConfig + , serverErrorReason + ) +import Subscription.Store + ( SubscriptionStore + , awaitNoTrackedExpiries + , hasTrackedExpiries + , newSubscriptionStore + , register + , startExpiryWorker + , startWorkers + , unregister + ) +import Subscription.Types + ( PendingLimits (..) + , SubscribeConfig (..) + , SubscriptionMeta (SubscriptionMeta) + , defaultPendingLimits + ) +import qualified Types.Connect as Connect +import qualified Types.Info as Info +import qualified Types.Msg as Msg +import Types.Ping (Ping (..)) +import qualified Types.Pub as Pub +import qualified Types.Sub as Sub +import qualified Types.Unsub as Unsub +import Validators.Validators (validate) + +data ClientOptions = ClientOptions + { optionConnectConfig :: Connect.Connect + , optionAuth :: Auth + , optionTlsConfig :: Maybe TLSConfig + , optionLoggerConfig :: LoggerConfig + , optionConnectionAttempts :: Int + , optionConnectTimeoutMicros :: Int + , optionCallbackConcurrency :: Int + , optionMessageLimit :: Int + , optionPendingLimits :: PendingLimits + , optionErrorHandler :: NatsError -> IO () + , optionExitAction :: ClientExitReason -> IO () + , optionConnectOptions :: [(String, Int)] + } + +-- | An opaque NATS server endpoint. +-- +-- Keeping this representation private allows endpoint schemes and transports +-- to be added without changing the connection API. +data Server = Server String Int + deriving (Eq, Show) + +data ServerConfigError = EmptyServerHost + | InvalidServerPort Int + deriving (Eq, Show) + +-- | Construct a TCP NATS server endpoint. +server :: String -> Int -> Either ServerConfigError Server +server host port + | null host = Left EmptyServerHost + | port < 1 || port > 65535 = Left (InvalidServerPort port) + | otherwise = Right (Server host port) + +-- | Construct a server using the standard NATS port, 4222. +serverWithDefaultPort :: String -> Either ServerConfigError Server +serverWithDefaultPort host = server host 4222 + +serverHost :: Server -> String +serverHost (Server host _) = host + +serverPort :: Server -> Int +serverPort (Server _ port) = port + +-- | Connect to one of the configured NATS servers. +connect :: [Server] -> [ConfigOption] -> IO (Either ConnectError Client) +connect servers = + newClient [(serverHost endpoint, serverPort endpoint) | endpoint <- servers] + +-- | Compatibility connection entry point using raw @(host, port)@ tuples. +newClient :: [(String, Int)] -> [ConfigOption] -> IO (Either ConnectError Client) +newClient servers configOptions = do + loggerConfig' <- defaultLogger + ctx <- newLogContext + let defaultOptions = applyCallOptions configOptions ClientOptions + { optionConnectConfig = defaultConnect + , optionAuth = AuthNone.auth + , optionTlsConfig = Nothing + , optionLoggerConfig = loggerConfig' + , optionConnectionAttempts = 5 + , optionConnectTimeoutMicros = 2 * 1000000 + , optionCallbackConcurrency = 1 + , optionMessageLimit = 1024 * 1024 + , optionPendingLimits = defaultPendingLimits + , optionErrorHandler = const (pure ()) + , optionExitAction = const (pure ()) + , optionConnectOptions = servers + } + clientConfig = + ClientConfig + { connectionAttempts = optionConnectionAttempts defaultOptions + , connectTimeoutMicros = optionConnectTimeoutMicros defaultOptions + , callbackConcurrency = optionCallbackConcurrency defaultOptions + , messageLimit = optionMessageLimit defaultOptions + , connectConfig = optionConnectConfig defaultOptions + , loggerConfig = optionLoggerConfig defaultOptions + , tlsConfig = optionTlsConfig defaultOptions + , exitAction = optionExitAction defaultOptions + , connectOptions = optionConnectOptions defaultOptions + } + configuredAuth = optionAuth defaultOptions + + queue <- newQueue + conn <- newConn connectionApi + clientState <- newClientState clientConfig queue conn ctx + store <- + newSubscriptionStore + (optionPendingLimits defaultOptions) + (handleSlowConsumer clientState (optionErrorHandler defaultOptions)) + + setConnectName clientState (Connect.name (optionConnectConfig defaultOptions)) + logStaticConfiguration clientState defaultOptions + + startWorkers + (callbackConcurrency clientConfig) + store + (do + waitForClosed clientState + awaitNoTrackedExpiries store) + (handleCallbackError clientState) + + startExpiryWorker store $ + shouldStopExpiryWorker clientState store + + void . forkIO $ + runEngine + connectionApi + streamingApi + broadcastingApi + (parserApiWithMessageLimit (messageLimit clientConfig)) + clientState + store + configuredAuth + + let client = Client + { publish = \subject payload publishOptions -> do + let cfg = applyCallOptions publishOptions defaultPublishConfig + publishClient clientState subject payload cfg + , subscribe = \subject subscribeOptions callback -> do + let cfg = applyCallOptions subscribeOptions defaultSubscribeConfig + subscribeClient clientState store False subject cfg callback + , subscribeOnce = \subject subscribeOptions callback -> do + let cfg = applyCallOptions subscribeOptions defaultSubscribeConfig + subscribeClient clientState store True subject cfg callback + , request = \subject payload requestOptions -> do + let cfg = applyCallOptions requestOptions defaultRequestConfig + requestClient clientState store subject payload cfg + , unsubscribe = \subscription options -> + case applyCallOptions options UnsubscribeConfig of + UnsubscribeConfig -> unsubscribeClient clientState store subscription + , newInbox = nextInbox clientState + , ping = \options -> + case applyCallOptions options PingConfig of + PingConfig -> flushClient clientState + , flush = \options -> + case applyCallOptions options FlushConfig of + FlushConfig -> flushClient clientState + , reset = \options -> + case applyCallOptions options ResetConfig of + ResetConfig -> resetClient connectionApi clientState store + , close = \options -> + case applyCallOptions options CloseConfig of + CloseConfig -> closeClient connectionApi clientState store + } + + initialResult <- atomically (waitForInitialConnection clientState) + pure (client <$ initialResult) + +type ConfigOption = CallOption ClientOptions + +withConnectName :: BS.ByteString -> ConfigOption +withConnectName name config = + config + { optionConnectConfig = + (optionConnectConfig config) { Connect.name = Just name } + } + +withEcho :: Bool -> ConfigOption +withEcho enabled config = + config + { optionConnectConfig = + (optionConnectConfig config) { Connect.echo = Just enabled } + } + +withAuthToken :: AuthTokenData -> ConfigOption +withAuthToken token = addAuth (AuthToken.auth token) + +-- | Fetch a token for every connection and reconnection attempt. +withAuthTokenHandler :: AuthTokenHandler -> ConfigOption +withAuthTokenHandler handler = addAuth (AuthToken.authHandler handler) + +withUserPass :: UserPassData -> ConfigOption +withUserPass userPass = addAuth (AuthUserPass.auth userPass) + +-- | Fetch a username and password for every connection and reconnection attempt. +withUserPassHandler :: UserPassHandler -> ConfigOption +withUserPassHandler handler = addAuth (AuthUserPass.authHandler handler) + +withNKey :: NKeyData -> ConfigOption +withNKey seed = addAuth (AuthNKey.auth seed) + +-- | Authenticate with a public NKey and a handler that signs the server nonce. +-- The handler returns the raw 64-byte Ed25519 signature; natskell performs the +-- protocol's base64url encoding. +withNKeyHandler :: NKeyPublicKey -> SignatureHandler -> ConfigOption +withNKeyHandler publicKey handler = addAuth (AuthNKey.authHandler publicKey handler) + +withJWT :: JWTTokenData -> ConfigOption +withJWT creds = addAuth (AuthJwt.auth creds) + +-- | Fetch a user JWT and sign the server nonce for every connection attempt. +withJWTHandlers :: JWTHandler -> SignatureHandler -> ConfigOption +withJWTHandlers jwtHandler signatureHandler = + addAuth (AuthJwt.authHandlers jwtHandler signatureHandler) + +addAuth :: Auth -> ConfigOption +addAuth auth config = + config { optionAuth = mergeAuth (optionAuth config) auth } + +-- | Require TLS using the operating system trust store. +withTLS :: ConfigOption +withTLS = modifyTLSConfig id + +withTLSCert :: TLSCertData -> ConfigOption +withTLSCert cert = + modifyTLSConfig $ \tls -> tls { tlsClientCertificate = Just cert } + +-- | Trust a PEM-encoded root certificate for this client. Once configured, +-- these roots replace the operating-system trust store for the connection. +withTLSRootCA :: BS.ByteString -> ConfigOption +withTLSRootCA root = + modifyTLSConfig $ \tls -> + tls { tlsRootCertificates = tlsRootCertificates tls ++ [root] } + +-- | Override the host name used for certificate verification and SNI. +withTLSServerName :: String -> ConfigOption +withTLSServerName serverName = + modifyTLSConfig $ \tls -> tls { tlsServerName = Just serverName } + +-- | Disable server certificate verification. This is unsafe and should only +-- be used when the peer is trusted by some mechanism outside TLS. +withTLSInsecure :: ConfigOption +withTLSInsecure = + modifyTLSConfig $ \tls -> tls { tlsInsecure = True } + +modifyTLSConfig :: (TLSConfig -> TLSConfig) -> ConfigOption +modifyTLSConfig update config = + config + { optionTlsConfig = + Just (update (fromMaybe defaultTLSConfig (optionTlsConfig config))) + } + +withMinimumLogLevel :: LogLevel -> ConfigOption +withMinimumLogLevel minimumLogLevel config = + config + { optionLoggerConfig = + (optionLoggerConfig config) { minLogLevel = minimumLogLevel } + } + +withLogAction :: (LogEntry -> IO ()) -> ConfigOption +withLogAction logAction config = + config + { optionLoggerConfig = + (optionLoggerConfig config) { logFn = logAction } + } + +withConnectionAttempts :: Int -> ConfigOption +withConnectionAttempts attempts config = + config { optionConnectionAttempts = max 1 attempts } + +-- | Set the maximum time for INFO, TLS, CONNECT, and PONG negotiation on each +-- server attempt. Values below one microsecond are clamped to one. +withConnectTimeoutMicros :: Int -> ConfigOption +withConnectTimeoutMicros timeoutMicros config = + config { optionConnectTimeoutMicros = max 1 timeoutMicros } + +withCallbackConcurrency :: Int -> ConfigOption +withCallbackConcurrency concurrency config = + config { optionCallbackConcurrency = concurrency } + +-- | Set the largest encoded message body this client accepts. For messages +-- with headers, the encoded header block and payload both count toward the +-- limit. Outbound messages are also constrained by the server's max_payload. +withMessageLimit :: Int -> ConfigOption +withMessageLimit limit config = + config { optionMessageLimit = max 1 limit } + +-- | Bound callback deliveries pending across the entire client. The first +-- limit is a message count and the second is encoded message bytes. +withPendingDeliveryLimits :: Int -> Int -> ConfigOption +withPendingDeliveryLimits maximumMessages maximumBytes config = + config + { optionPendingLimits = + PendingLimits + { pendingMessageLimit = max 1 maximumMessages + , pendingByteLimit = max 1 maximumBytes + } + } + +-- | Receive asynchronous client errors such as slow-consumer notifications. +-- The handler runs on the callback worker pool, never the socket reader. +withErrorHandler :: (NatsError -> IO ()) -> ConfigOption +withErrorHandler handler config = + config { optionErrorHandler = handler } + +-- | Compatibility alias for 'withMessageLimit'. +withBufferLimit :: Int -> ConfigOption +withBufferLimit = withMessageLimit + +{-# DEPRECATED withBufferLimit "Use withMessageLimit instead." #-} + +withExitAction :: (ClientExitReason -> IO ()) -> ConfigOption +withExitAction action config = config { optionExitAction = action } + +defaultConnect :: Connect.Connect +defaultConnect = + Connect.Connect + { Connect.verbose = False + , Connect.pedantic = True + , Connect.tls_required = False + , Connect.auth_token = Nothing + , Connect.user = Nothing + , Connect.pass = Nothing + , Connect.name = Nothing + , Connect.lang = "haskell" + , Connect.version = BC.pack (showVersion Package.version) + , Connect.protocol = Nothing + , Connect.echo = Just True + , Connect.sig = Nothing + , Connect.jwt = Nothing + , Connect.nkey = Nothing + , Connect.no_responders = Just True + , Connect.headers = Just True + } + +defaultSubscribeConfig :: SubscribeConfig +defaultSubscribeConfig = SubscribeConfig Nothing Nothing + +defaultRequestConfig :: RequestConfig +defaultRequestConfig = RequestConfig 2 Nothing + +logStaticConfiguration :: ClientState -> ClientOptions -> IO () +logStaticConfiguration client options = + runClient client $ do + case authMethods (optionAuth options) of + [] -> logMessage Info "no authentication method provided" + methods -> forM_ methods $ \method -> + logMessage Info ("using " ++ method) + case optionTlsConfig options of + Nothing -> + pure () + Just tls -> do + logMessage Info "using tls" + when (tlsInsecure tls) $ + logMessage Warn "tls certificate verification is disabled" + +handleCallbackError :: ClientState -> SomeException -> IO () +handleCallbackError client err = + runClient client $ + logMessage Error ("callback failed: " ++ displayException err) + +handleSlowConsumer :: ClientState -> (NatsError -> IO ()) -> IO () +handleSlowConsumer client handler = do + runClient client $ + logMessage Error "slow consumer: global pending delivery limit reached" + handler NatsSlowConsumer + +shouldStopExpiryWorker :: ClientState -> SubscriptionStore -> IO Bool +shouldStopExpiryWorker client store = do + status <- readStatus client + tracked <- hasTrackedExpiries store + pure $ + case status of + Closed _ -> not tracked + _ -> False + +toMessage :: Msg.Msg -> Message +toMessage msg = + Message + { subject = Msg.subject msg + , sid = Msg.sid msg + , replyTo = Msg.replyTo msg + , payload = fromMaybe BS.empty (Msg.payload msg) + , headers = Msg.headers msg + } + +publishClient + :: ClientState + -> Msg.Subject + -> Msg.Payload + -> PublishConfig + -> IO (Either NatsError ()) +publishClient client subject messagePayload cfg = do + runClient client $ + logMessage Debug ("publishing to subject: " ++ show subject) + let publishMessage = + Pub.Pub + { Pub.subject = subject + , Pub.payload = + if BS.null messagePayload then Nothing else Just messagePayload + , Pub.replyTo = publishReplyTo cfg + , Pub.headers = publishHeaders cfg + } + validation <- canPublish client publishMessage + case validation of + Left err -> pure (Left err) + Right () -> do + enqueue client (QueueItem publishMessage) + pure (Right ()) + +canPublish :: ClientState -> Pub.Pub -> IO (Either NatsError ()) +canPublish client publishMessage = + case validate publishMessage of + Left reason -> do + runClient client $ + logMessage Error ("rejecting invalid publish: " ++ BC.unpack reason) + pure (Left (NatsValidationError reason)) + Right () -> do + statusResult <- runningResult client + case statusResult of + Left err -> pure (Left err) + Right () -> do + serverInfo <- readServerInfo client + let actual = Pub.messageSize publishMessage + clientMaximum = messageLimit (config client) + maximumSize = + maybe + clientMaximum + (min clientMaximum . Info.max_payload) + serverInfo + if actual > maximumSize + then do + runClient client $ + logMessage Error + ( "rejecting publish: message size " + ++ show actual + ++ " exceeds effective limit " + ++ show maximumSize + ) + pure (Left (NatsPayloadTooLarge actual maximumSize)) + else + pure (Right ()) + +subscribeClient + :: ClientState + -> SubscriptionStore + -> Bool + -> Msg.Subject + -> SubscribeConfig + -> (Message -> IO ()) + -> IO (Either NatsError Subscription) +subscribeClient client store isReply subject cfg callback = + subscribeRawClient client store isReply subject cfg (maybe (pure ()) (callback . toMessage)) + +subscribeRawClient + :: ClientState + -> SubscriptionStore + -> Bool + -> Msg.Subject + -> SubscribeConfig + -> (Maybe Msg.Msg -> IO ()) + -> IO (Either NatsError Subscription) +subscribeRawClient client store isReply subject cfg callback = do + subscribeRawClientWithOverflow + client + store + isReply + subject + cfg + callback + (pure ()) + +subscribeRawClientWithOverflow + :: ClientState + -> SubscriptionStore + -> Bool + -> Msg.Subject + -> SubscribeConfig + -> (Maybe Msg.Msg -> IO ()) + -> IO () + -> IO (Either NatsError Subscription) +subscribeRawClientWithOverflow client store isReply subject cfg callback onDropped = do + runClient client $ + logMessage Debug ("subscribing to subject: " ++ show subject) + statusResult <- runningResult client + let queueGroup = subscribeQueueGroup cfg + sid <- nextSid client + let + subscriptionMessage = + Sub.Sub + { Sub.subject = subject + , Sub.queueGroup = queueGroup + , Sub.sid = sid + } + case validate subscriptionMessage of + Left reason -> do + runClient client $ + logMessage Error ("rejecting invalid subscription: " ++ BC.unpack reason) + pure (Left (NatsValidationError reason)) + Right () -> case statusResult of + Left err -> pure (Left err) + Right () -> do + let meta = + SubscriptionMeta subject queueGroup isReply + register store sid meta cfg callback onDropped + enqueue client $ + QueueItem + subscriptionMessage + when isReply $ + enqueue client + (QueueItem + Unsub.Unsub + { Unsub.sid = sid + , Unsub.maxMsg = Just 1 + }) + pure (Right (Subscription sid)) + +requestClient + :: ClientState + -> SubscriptionStore + -> Msg.Subject + -> Msg.Payload + -> RequestConfig + -> IO (Either NatsError Message) +requestClient client store requestSubject requestPayload cfg = + mask $ \restore -> do + response <- newEmptyTMVarIO + deadline <- registerDelay (durationMicros (requestTimeout cfg)) + inbox <- nextInbox client + let subscriptionConfig = SubscribeConfig Nothing Nothing + deliver Nothing = atomically (void (tryPutTMVar response (Left NatsRequestTimedOut))) + deliver (Just msg) = + atomically . void . tryPutTMVar response $ + let message = toMessage msg + in if isNoResponders message + then Left NatsNoResponders + else Right message + rejectSlowConsumer = + atomically (void (tryPutTMVar response (Left NatsSlowConsumer))) + subscriptionResult <- + subscribeRawClientWithOverflow + client + store + True + inbox + subscriptionConfig + deliver + rejectSlowConsumer + case subscriptionResult of + Left err -> pure (Left err) + Right subscription -> do + let publishConfig = PublishConfig (requestHeaders cfg) (Just inbox) + cleanup = void (unsubscribeClient client store subscription) + publishResult <- publishClient client requestSubject requestPayload publishConfig + case publishResult of + Left err -> cleanup >> pure (Left err) + Right () -> do + result <- restore (awaitRequest client deadline response) `onException` cleanup + cleanup + pure result + +awaitRequest + :: ClientState + -> TVar Bool + -> TMVar (Either NatsError Message) + -> IO (Either NatsError Message) +awaitRequest client deadline response = do + outcome <- atomically $ + (Just <$> readTMVar response) + `orElse` (do + expired <- readTVar deadline + check expired + pure (Just (Left NatsRequestTimedOut))) + `orElse` (Nothing <$ waitForNotRunning client) + case outcome of + Just result -> pure result + Nothing -> do + status <- readStatus client + pure (Left (closedError status)) + +durationMicros :: NominalDiffTime -> Int +durationMicros duration = + fromInteger (min (toInteger (maxBound :: Int)) micros) + where + micros = max 0 (floor (realToFrac duration * (1000000 :: Double))) + +isNoResponders :: Message -> Bool +isNoResponders message = + case headers message of + Nothing -> False + Just messageHeaders -> + any isNoRespondersHeader messageHeaders + where + isNoRespondersHeader (name, value) = + BC.map toAsciiLower name == "status" && value == "503" + toAsciiLower byte + | isAsciiUpper byte = toEnum (fromEnum byte + 32) + | otherwise = byte + +unsubscribeClient + :: ClientState + -> SubscriptionStore + -> Subscription + -> IO (Either NatsError ()) +unsubscribeClient client store (Subscription sid) = do + runClient client $ + logMessage Debug ("unsubscribing SID: " ++ show sid) + unregister store sid + statusResult <- runningResult client + case statusResult of + Left err -> pure (Left err) + Right () -> do + enqueue client $ + QueueItem + Unsub.Unsub + { Unsub.sid = sid + , Unsub.maxMsg = Nothing + } + pure (Right ()) + +pingClient :: ClientState -> IO () -> IO () +pingClient client action = do + runClient client $ + logMessage Debug "sending ping to server" + pushPingAction client action + enqueue client (QueueItem Ping) + +flushClient :: ClientState -> IO (Either NatsError ()) +flushClient client = do + statusResult <- runningResult client + case statusResult of + Left err -> pure (Left err) + Right () -> do + ponged <- newEmptyTMVarIO + pingClient client (atomically (void (tryPutTMVar ponged ()))) + outcome <- atomically $ + (True <$ readTMVar ponged) + `orElse` (False <$ waitForNotRunning client) + if outcome + then pure (Right ()) + else do + status <- readStatus client + pure (Left (closedError status)) + +runningResult :: ClientState -> IO (Either NatsError ()) +runningResult client = do + status <- readStatus client + pure $ + case status of + Running -> Right () + Closing reason -> Left (NatsConnectionClosed reason) + Closed reason -> Left (NatsConnectionClosed reason) + +closedError :: ClientStatus -> NatsError +closedError status = + case status of + Closing reason -> NatsConnectionClosed reason + Closed reason -> NatsConnectionClosed reason + Running -> NatsConnectionClosed ExitResetRequested diff --git a/client/JetStream/API.hs b/client/JetStream/API.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/API.hs @@ -0,0 +1,46 @@ +-- | Capability record for the JetStream client surface. +module JetStream.API + ( JetStream + , streams + , consumers + , publisher + , messages + , management + , JetStreamRequestOption + , withRequestTimeout + , JetStreamApiError + , apiErrorCode + , apiErrorCodeDetail + , apiErrorDescription + , JetStreamError (..) + , module JetStream.API.Consumer + , module JetStream.API.Management + , module JetStream.API.Message + , module JetStream.API.Publish + , module JetStream.API.Stream + ) where + +import JetStream.API.Consumer +import JetStream.API.Management +import JetStream.API.Message +import JetStream.API.Publish +import JetStream.API.Stream +import JetStream.Error + ( JetStreamApiError + , JetStreamError (..) + , apiErrorCode + , apiErrorCodeDetail + , apiErrorDescription + ) +import JetStream.Options + ( JetStream + , consumers + , management + , messages + , publisher + , streams + ) +import JetStream.Types + ( JetStreamRequestOption + , withRequestTimeout + ) diff --git a/client/JetStream/API/Consumer.hs b/client/JetStream/API/Consumer.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/API/Consumer.hs @@ -0,0 +1,4 @@ +-- | JetStream consumer management contract. +module JetStream.API.Consumer (module JetStream.Consumer.API) where + +import JetStream.Consumer.API diff --git a/client/JetStream/API/Management.hs b/client/JetStream/API/Management.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/API/Management.hs @@ -0,0 +1,40 @@ +-- | JetStream account and tier management contract. +module JetStream.API.Management + ( ManagementAPI + , accountInfo + , AccountAPIStats + , accountAPILevel + , accountAPITotal + , accountAPIErrors + , accountAPIInflight + , AccountInfo + , accountInfoTier + , accountInfoDomain + , accountInfoAPI + , accountInfoTiers + , AccountLimits + , accountLimitsMaxMemory + , accountLimitsMaxStorage + , accountLimitsMaxStreams + , accountLimitsMaxConsumers + , accountLimitsMaxAckPending + , accountLimitsMemoryMaxStreamBytes + , accountLimitsStorageMaxStreamBytes + , accountLimitsMaxBytesRequired + , AccountTier + , accountTierMemory + , accountTierStorage + , accountTierReservedMemory + , accountTierReservedStorage + , accountTierStreams + , accountTierConsumers + , accountTierLimits + ) where + +import JetStream.Management.API (ManagementAPI, accountInfo) +import JetStream.Types + ( AccountAPIStats (..) + , AccountInfo (..) + , AccountLimits (..) + , AccountTier (..) + ) diff --git a/client/JetStream/API/Message.hs b/client/JetStream/API/Message.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/API/Message.hs @@ -0,0 +1,4 @@ +-- | JetStream message consumption and acknowledgement contract. +module JetStream.API.Message (module JetStream.Message.API) where + +import JetStream.Message.API diff --git a/client/JetStream/API/Publish.hs b/client/JetStream/API/Publish.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/API/Publish.hs @@ -0,0 +1,4 @@ +-- | JetStream publishing contract. +module JetStream.API.Publish (module JetStream.Publish.API) where + +import JetStream.Publish.API diff --git a/client/JetStream/API/Stream.hs b/client/JetStream/API/Stream.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/API/Stream.hs @@ -0,0 +1,4 @@ +-- | JetStream stream management contract. +module JetStream.API.Stream (module JetStream.Stream.API) where + +import JetStream.Stream.API diff --git a/client/JetStream/Client.hs b/client/JetStream/Client.hs new file mode 100644 --- /dev/null +++ b/client/JetStream/Client.hs @@ -0,0 +1,42 @@ +-- | High-level client implementation for JetStream. +module JetStream.Client + ( newJetStream + , JetStream + , JetStreamOption + , JetStreamRequestOption + , JetStreamConfigError (..) + , withDomain + , withRequestTimeout + , withRequestTimeoutMicros + ) where + +import qualified API as Nats +import JetStream.API (JetStream) +import qualified JetStream.Consumer as Consumer +import qualified JetStream.Management as Management +import qualified JetStream.Message as Message +import JetStream.Options + ( JetStream (..) + , JetStreamConfigError (..) + , JetStreamOption + , tryNewJetStreamContext + , withDomain + , withRequestTimeout + , withRequestTimeoutMicros + ) +import qualified JetStream.Publish as Publish +import qualified JetStream.Stream as Stream +import JetStream.Types (JetStreamRequestOption) + +-- | Build JetStream capabilities from an existing NATS client. +newJetStream :: Nats.Client -> [JetStreamOption] -> Either JetStreamConfigError JetStream +newJetStream client options = do + ctx <- tryNewJetStreamContext client options + let consumerAPI = Consumer.consumerAPI ctx + pure JetStream + { streams = Stream.streamAPI ctx + , consumers = consumerAPI + , publisher = Publish.publishAPI ctx + , messages = Message.messageAPI ctx consumerAPI + , management = Management.managementAPI ctx + } diff --git a/internal/Client/API.hs b/internal/Client/API.hs new file mode 100644 --- /dev/null +++ b/internal/Client/API.hs @@ -0,0 +1,157 @@ +-- | Internal home for the core NATS client capability surface. +-- +-- The public @API@ module deliberately hides the constructors defined here. +-- Keeping construction private lets the capability grow without changing the +-- public representation. +module Client.API + ( Client (..) + , Message (..) + , MsgView + , Subscription (..) + , subscriptionSid + , NatsError (..) + , Subject + , Payload + , Headers + , PublishOption + , SubscribeOption + , RequestOption + , UnsubscribeOption + , PingOption + , FlushOption + , ResetOption + , CloseOption + , RequestConfig (..) + , UnsubscribeConfig (..) + , PingConfig (..) + , FlushConfig (..) + , ResetConfig (..) + , CloseConfig (..) + , withSubscriptionExpiry + , withQueueGroup + , withReplyTo + , withHeaders + , withRequestTimeout + , withRequestHeaders + ) where + +import qualified Data.ByteString as BS +import Data.Time.Clock (NominalDiffTime) +import Lib.CallOption (CallOption) +import Publish.Config (PublishConfig (..)) +import State.Types (ClientExitReason) +import Subscription.Types (SubscribeConfig (..)) +import Types.Msg (Headers, Payload, SID, Subject) + +-- | Client capabilities for publishing, subscribing, request-reply, and +-- lifecycle control. +-- +-- This constructor is internal. Public callers receive a 'Client' from +-- 'Client.newClient' and use the named operations re-exported by @API@. +data Client = Client + { publish :: Subject -> Payload -> [PublishOption] -> IO (Either NatsError ()) + -- ^ Publish a message. + , subscribe :: Subject -> [SubscribeOption] -> (Message -> IO ()) -> IO (Either NatsError Subscription) + -- ^ Subscribe to a subject until explicitly unsubscribed. + , subscribeOnce :: Subject -> [SubscribeOption] -> (Message -> IO ()) -> IO (Either NatsError Subscription) + -- ^ Subscribe until the first message is delivered. + , request :: Subject -> Payload -> [RequestOption] -> IO (Either NatsError Message) + -- ^ Publish a request and wait for one response. + , unsubscribe :: Subscription -> [UnsubscribeOption] -> IO (Either NatsError ()) + -- ^ Unsubscribe a subscription created by this client. + , newInbox :: IO Subject + -- ^ Create a unique inbox subject for replies. + , ping :: [PingOption] -> IO (Either NatsError ()) + -- ^ Round trip a PING through the server. + , flush :: [FlushOption] -> IO (Either NatsError ()) + -- ^ Wait until all previously buffered writes reach the server. + , reset :: [ResetOption] -> IO () + -- ^ Reset the client connection state. + , close :: [CloseOption] -> IO () + -- ^ Close the client connection and release resources. + } + +-- | A message delivered by NATS. +-- +-- NATS always has a payload, which may be empty. This is intentionally +-- represented by a strict 'BS.ByteString', not @Maybe ByteString@. +data Message = Message + { subject :: Subject + , sid :: SID + , replyTo :: Maybe Subject + , payload :: Payload + , headers :: Maybe Headers + } + deriving (Eq, Show) + +-- | Compatibility name retained for code written against pre-0.4 releases. +type MsgView = Message + +-- | An opaque handle to an active subscription. +newtype Subscription = Subscription SID + deriving (Eq, Show) + +-- | The protocol identifier used by the legacy @API@ module. New code should +-- treat 'Subscription' as an opaque handle. +subscriptionSid :: Subscription -> SID +subscriptionSid (Subscription value) = value + +-- | Failures reported by core NATS operations. +data NatsError = NatsValidationError BS.ByteString + | NatsPayloadTooLarge Int Int + | NatsSlowConsumer + | NatsConnectionClosed ClientExitReason + | NatsRequestTimedOut + | NatsNoResponders + deriving (Eq, Show) + +type PublishOption = CallOption PublishConfig + +type SubscribeOption = CallOption SubscribeConfig + +data RequestConfig = RequestConfig + { requestTimeout :: NominalDiffTime + , requestHeaders :: Maybe Headers + } + +type RequestOption = CallOption RequestConfig + +data UnsubscribeConfig = UnsubscribeConfig +type UnsubscribeOption = CallOption UnsubscribeConfig + +data PingConfig = PingConfig +type PingOption = CallOption PingConfig + +data FlushConfig = FlushConfig +type FlushOption = CallOption FlushConfig + +data ResetConfig = ResetConfig +type ResetOption = CallOption ResetConfig + +data CloseConfig = CloseConfig +type CloseOption = CallOption CloseConfig + +-- | Stop a one-shot subscription if it has not received a message within the +-- given interval. +withSubscriptionExpiry :: NominalDiffTime -> SubscribeOption +withSubscriptionExpiry expirySeconds cfg = cfg { expiry = Just expirySeconds } + +-- | Deliver messages to one member of the named queue group. +withQueueGroup :: Subject -> SubscribeOption +withQueueGroup queueGroup cfg = cfg { subscribeQueueGroup = Just queueGroup } + +-- | Set the reply subject on an outgoing publish. +withReplyTo :: Subject -> PublishOption +withReplyTo replySubject cfg = cfg { publishReplyTo = Just replySubject } + +-- | Set headers on an outgoing publish. +withHeaders :: Headers -> PublishOption +withHeaders messageHeaders cfg = cfg { publishHeaders = Just messageHeaders } + +-- | Set the amount of time to wait for a request reply. +withRequestTimeout :: NominalDiffTime -> RequestOption +withRequestTimeout timeout cfg = cfg { requestTimeout = max 0 timeout } + +-- | Set headers on an outgoing request. +withRequestHeaders :: Headers -> RequestOption +withRequestHeaders messageHeaders cfg = cfg { requestHeaders = Just messageHeaders } diff --git a/internal/Lib/CallOption.hs b/internal/Lib/CallOption.hs --- a/internal/Lib/CallOption.hs +++ b/internal/Lib/CallOption.hs @@ -3,4 +3,4 @@ type CallOption a = (a -> a) applyCallOptions :: [CallOption a] -> a -> a -applyCallOptions options value = foldr ($) value options +applyCallOptions options value = foldl (flip ($)) value options diff --git a/internal/Plumbing/Network/Connection/Tls.hs b/internal/Plumbing/Network/Connection/Tls.hs --- a/internal/Plumbing/Network/Connection/Tls.hs +++ b/internal/Plumbing/Network/Connection/Tls.hs @@ -7,9 +7,11 @@ import Control.Exception import Control.Monad -import Data.ByteString (ByteString) -import qualified Data.ByteString.Char8 as BC -import qualified Data.ByteString.Lazy as LBS +import qualified Data.ByteString.Char8 as BC +import qualified Data.ByteString.Lazy as LBS +import Data.Maybe (fromMaybe) +import Data.X509.CertificateStore (makeCertificateStore) +import Data.X509.Memory (readSignedObjectFromMemory) import Network.Connection.Core ( bufferRead , currentTransport @@ -21,9 +23,11 @@ , Transport (..) , TransportOption (..) ) -import qualified Network.Socket as NS -import qualified Network.Socket.ByteString as NSB -import qualified Network.TLS as TLS +import qualified Network.Socket as NS +import qualified Network.Socket.ByteString as NSB +import qualified Network.TLS as TLS +import System.X509 (getSystemCertificateStore) +import Types.TLS (TLSConfig (..)) tlsTransport :: TLS.Context -> Transport tlsTransport ctx = @@ -61,7 +65,7 @@ Nothing -> return $ Left "Transport not initialized" Just currentTransport' -> case transportUpgrade currentTransport' of - Nothing -> return $ Right () + Nothing -> return $ Left "Transport does not support TLS" Just upgrade -> do result <- upgrade params case result of @@ -70,7 +74,7 @@ pointTransport conn newTransport return $ Right () -upgradeToTLSWithConfig :: Conn -> String -> Maybe (ByteString, ByteString) -> IO (Either String ()) +upgradeToTLSWithConfig :: Conn -> String -> TLSConfig -> IO (Either String ()) upgradeToTLSWithConfig conn host tlsConfig = do paramsResult <- buildTlsParams host tlsConfig case paramsResult of @@ -82,29 +86,65 @@ let useTls = transportTlsRequested transportOption || transportTlsRequired transportOption if useTls then do - result <- upgradeToTLSWithConfig conn (transportHost transportOption) (transportTlsCert transportOption) - case result of - Left err -> return (Left err) - Right () -> do - enableReadWorker conn - return (Right ()) + case transportTlsConfig transportOption of + Nothing -> return (Left "TLS transport requested without TLS configuration") + Just tlsConfig -> do + result <- upgradeToTLSWithConfig conn (transportHost transportOption) tlsConfig + case result of + Left err -> return (Left err) + Right () -> do + enableReadWorker conn + return (Right ()) else do bufferRead conn (transportInitialBytes transportOption) enableReadWorker conn return (Right ()) -buildTlsParams :: String -> Maybe (ByteString, ByteString) -> IO (Either String TLS.ClientParams) +buildTlsParams :: String -> TLSConfig -> IO (Either String TLS.ClientParams) buildTlsParams host tlsConfig = do - let base = TLS.defaultParamsClient host (BC.pack host) - hooks = TLS.clientHooks base - shared = TLS.clientShared base - acceptAny = hooks { TLS.onServerCertificate = \_ _ _ _ -> return [] } - case tlsConfig of - Just (certPem, keyPem) -> - case TLS.credentialLoadX509FromMemory certPem keyPem of - Left err -> return (Left err) - Right cred -> - let hooks' = acceptAny { TLS.onCertificateRequest = \_ -> return (Just cred) } - shared' = shared { TLS.sharedCredentials = TLS.Credentials [cred] } - in return (Right base { TLS.clientHooks = hooks', TLS.clientShared = shared' }) - Nothing -> return (Right base { TLS.clientHooks = acceptAny }) + result <- try @SomeException $ do + systemStore <- getSystemCertificateStore + let verificationHost = fromMaybe host (tlsServerName tlsConfig) + base = TLS.defaultParamsClient verificationHost (BC.pack verificationHost) + hooks = TLS.clientHooks base + shared = TLS.clientShared base + parsedRoots = map readSignedObjectFromMemory (tlsRootCertificates tlsConfig) + invalidRoot = any null parsedRoots + caStore = + if null parsedRoots + then systemStore + else makeCertificateStore (concat parsedRoots) + verificationHooks = + if tlsInsecure tlsConfig + then hooks { TLS.onServerCertificate = \_ _ _ _ -> return [] } + else hooks + sharedWithRoots = shared { TLS.sharedCAStore = caStore } + if invalidRoot + then fail "TLS root certificate PEM is invalid" + else + case tlsClientCertificate tlsConfig of + Nothing -> + pure base + { TLS.clientHooks = verificationHooks + , TLS.clientShared = sharedWithRoots + } + Just (certPem, keyPem) -> + case TLS.credentialLoadX509FromMemory certPem keyPem of + Left err -> fail err + Right cred -> + pure base + { TLS.clientHooks = + verificationHooks + { TLS.onCertificateRequest = \_ -> return (Just cred) + } + , TLS.clientShared = + sharedWithRoots + { TLS.sharedCredentials = TLS.Credentials [cred] + } + } + case result of + Left err -> + case fromException err :: Maybe SomeAsyncException of + Just _ -> throwIO err + Nothing -> pure (Left (displayException err)) + Right params -> pure (Right params) diff --git a/internal/Plumbing/Network/Connection/Types.hs b/internal/Plumbing/Network/Connection/Types.hs --- a/internal/Plumbing/Network/Connection/Types.hs +++ b/internal/Plumbing/Network/Connection/Types.hs @@ -10,6 +10,7 @@ import Data.ByteString (ByteString) import qualified Data.ByteString.Lazy as LBS import qualified Network.TLS as TLS +import Types.TLS (TLSConfig) type ReadError = String type WriteError = String @@ -24,10 +25,10 @@ } data TransportOption = TransportOption - { transportHost :: String - , transportTlsRequired :: Bool + { transportHost :: String + , transportTlsRequired :: Bool , transportTlsRequested :: Bool - , transportTlsCert :: Maybe (ByteString, ByteString) + , transportTlsConfig :: Maybe TLSConfig , transportInitialBytes :: ByteString } diff --git a/internal/Plumbing/Parser/API.hs b/internal/Plumbing/Parser/API.hs --- a/internal/Plumbing/Parser/API.hs +++ b/internal/Plumbing/Parser/API.hs @@ -24,6 +24,7 @@ | ParsedErr Err | ParsedInfo Info | ParsedMsg Msg + | ParsedMessageTooLarge Int Int deriving (Eq, Show) newtype ParserAPI a = ParserAPI { parse :: ByteString -> ParseStep a } diff --git a/internal/Plumbing/Parser/Attoparsec.hs b/internal/Plumbing/Parser/Attoparsec.hs --- a/internal/Plumbing/Parser/Attoparsec.hs +++ b/internal/Plumbing/Parser/Attoparsec.hs @@ -2,6 +2,7 @@ module Parser.Attoparsec ( parserApi + , parserApiWithMessageLimit ) where import Control.Applicative ((<|>)) @@ -11,14 +12,15 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as B8 import qualified Data.ByteString.Lazy as BSL +import Data.Char (toLower) import Data.Word (Word8) import Parser.API ( ParseStep (DropPrefix, Emit, NeedMore) - , ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong) + , ParsedMessage (ParsedErr, ParsedInfo, ParsedMessageTooLarge, ParsedMsg, ParsedOk, ParsedPing, ParsedPong) , ParserAPI (ParserAPI) ) import Types.Err - ( Err (ErrAuthTimeout, ErrAuthViolation, ErrErr, ErrInvalidProtocol, ErrInvalidSubject, ErrMaxConnsEx, ErrMaxControlLineEx, ErrMaxPayload, ErrPermViolation, ErrRoutePortConn, ErrSlowConsumer, ErrStaleConn, ErrTlsRequired, ErrUnknownOp) + ( Err (ErrAccountAuthExpired, ErrAuthExpired, ErrAuthRevoked, ErrAuthTimeout, ErrAuthViolation, ErrErr, ErrInvalidProtocol, ErrInvalidSubject, ErrMaxConnsEx, ErrMaxControlLineEx, ErrMaxPayload, ErrPermViolation, ErrRoutePortConn, ErrSlowConsumer, ErrStaleConn, ErrTlsRequired, ErrUnknownOp) ) import Types.Msg (Msg (Msg)) import Types.Ok (Ok (Ok)) @@ -26,11 +28,15 @@ import Types.Pong (Pong (Pong)) parserApi :: ParserAPI ParsedMessage -parserApi = ParserAPI parseStep +parserApi = parserApiWithMessageLimit maxBound -parseStep :: BS.ByteString -> ParseStep ParsedMessage -parseStep bytes = - case A.parse parsedMessageParser bytes of +parserApiWithMessageLimit :: Int -> ParserAPI ParsedMessage +parserApiWithMessageLimit maximumMessageSize = + ParserAPI (parseStep (max 1 maximumMessageSize)) + +parseStep :: Int -> BS.ByteString -> ParseStep ParsedMessage +parseStep maximumMessageSize bytes = + case A.parse (parsedMessageParser maximumMessageSize) bytes of A.Done rest parsedMessage -> Emit parsedMessage rest A.Partial _ -> @@ -38,8 +44,8 @@ A.Fail rest _ reason -> DropPrefix (dropBytes bytes rest) reason -parsedMessageParser :: A.Parser ParsedMessage -parsedMessageParser = do +parsedMessageParser :: Int -> A.Parser ParsedMessage +parsedMessageParser maximumMessageSize = do prefix <- A.peekWord8' case prefix of 43 -> @@ -47,11 +53,11 @@ 45 -> errParser 72 -> - hmsgParser + hmsgParser maximumMessageSize 73 -> infoParser 77 -> - msgParser + msgParser maximumMessageSize 80 -> pingOrPongParser _ -> @@ -87,29 +93,29 @@ Right info -> pure (ParsedInfo info) -msgParser :: A.Parser ParsedMessage -msgParser = do +msgParser :: Int -> A.Parser ParsedMessage +msgParser maximumMessageSize = do _ <- A.string "MSG" skipHorizontalSpace1 controlLine <- lineParser case B8.words controlLine of [subjectName, sidValue, payloadSizeBytes] -> - parseMsgPayload subjectName sidValue Nothing payloadSizeBytes + parseMsgPayload maximumMessageSize subjectName sidValue Nothing payloadSizeBytes [subjectName, sidValue, replySubject, payloadSizeBytes] -> - parseMsgPayload subjectName sidValue (Just replySubject) payloadSizeBytes + parseMsgPayload maximumMessageSize subjectName sidValue (Just replySubject) payloadSizeBytes _ -> fail ("invalid MSG control line: " ++ B8.unpack controlLine) -hmsgParser :: A.Parser ParsedMessage -hmsgParser = do +hmsgParser :: Int -> A.Parser ParsedMessage +hmsgParser maximumMessageSize = do _ <- A.string "HMSG" skipHorizontalSpace1 controlLine <- lineParser case B8.words controlLine of [subjectName, sidValue, headerSizeBytes, totalSizeBytes] -> - parseHMsgPayload subjectName sidValue Nothing headerSizeBytes totalSizeBytes + parseHMsgPayload maximumMessageSize subjectName sidValue Nothing headerSizeBytes totalSizeBytes [subjectName, sidValue, replySubject, headerSizeBytes, totalSizeBytes] -> - parseHMsgPayload subjectName sidValue (Just replySubject) headerSizeBytes totalSizeBytes + parseHMsgPayload maximumMessageSize subjectName sidValue (Just replySubject) headerSizeBytes totalSizeBytes _ -> fail ("invalid HMSG control line: " ++ B8.unpack controlLine) @@ -127,27 +133,30 @@ _ <- A.string "PONG\r\n" pure (ParsedPong Pong) -parseMsgPayload :: BS.ByteString -> BS.ByteString -> Maybe BS.ByteString -> BS.ByteString -> A.Parser ParsedMessage -parseMsgPayload subjectName sidValue replySubject payloadSizeBytes = +parseMsgPayload :: Int -> BS.ByteString -> BS.ByteString -> Maybe BS.ByteString -> BS.ByteString -> A.Parser ParsedMessage +parseMsgPayload maximumMessageSize subjectName sidValue replySubject payloadSizeBytes = case integerBytes "MSG payload size" payloadSizeBytes of Left parseReason -> fail parseReason - Right payloadSize -> do - payloadBytes <- A.take payloadSize - _ <- A.string "\r\n" - pure - ( ParsedMsg - ( Msg - subjectName - sidValue - replySubject - (nonEmpty payloadBytes) - Nothing + Right payloadSize + | payloadSize > maximumMessageSize -> + pure (ParsedMessageTooLarge payloadSize maximumMessageSize) + | otherwise -> do + payloadBytes <- A.take payloadSize + _ <- A.string "\r\n" + pure + ( ParsedMsg + ( Msg + subjectName + sidValue + replySubject + (nonEmpty payloadBytes) + Nothing + ) ) - ) -parseHMsgPayload :: BS.ByteString -> BS.ByteString -> Maybe BS.ByteString -> BS.ByteString -> BS.ByteString -> A.Parser ParsedMessage -parseHMsgPayload subjectName sidValue replySubject headerSizeBytes totalSizeBytes = +parseHMsgPayload :: Int -> BS.ByteString -> BS.ByteString -> Maybe BS.ByteString -> BS.ByteString -> BS.ByteString -> A.Parser ParsedMessage +parseHMsgPayload maximumMessageSize subjectName sidValue replySubject headerSizeBytes totalSizeBytes = case (integerBytes "HMSG header size" headerSizeBytes, integerBytes "HMSG total size" totalSizeBytes) of (Left parseReason, _) -> fail parseReason @@ -161,6 +170,8 @@ ++ " is smaller than header size " ++ show headerSize ) + | totalSize > maximumMessageSize -> + pure (ParsedMessageTooLarge totalSize maximumMessageSize) | otherwise -> do headerBytes <- A.take headerSize payloadBytes <- A.take (totalSize - headerSize) @@ -247,10 +258,16 @@ Right (ErrUnknownOp reason) | reason == "Attempted To Connect To Route Port" = Right (ErrRoutePortConn reason) - | reason == "Authorization Violation" = + | "authorization violation" `BS.isPrefixOf` normalized = Right (ErrAuthViolation reason) - | reason == "Authorization Timeout" = + | "authorization timeout" `BS.isPrefixOf` normalized = Right (ErrAuthTimeout reason) + | "user authentication expired" `BS.isPrefixOf` normalized = + Right (ErrAuthExpired reason) + | "user authentication revoked" `BS.isPrefixOf` normalized = + Right (ErrAuthRevoked reason) + | "account authentication expired" `BS.isPrefixOf` normalized = + Right (ErrAccountAuthExpired reason) | reason == "Invalid Client Protocol" = Right (ErrInvalidProtocol reason) | reason == "Maximum Control Line Exceeded" = @@ -272,7 +289,9 @@ | "Permissions Violation" `BS.isPrefixOf` reason = Right (ErrPermViolation reason) | otherwise = - Left ("unknown -ERR reason: " ++ B8.unpack reason) + Right (ErrErr reason) + where + normalized = B8.map toLower reason lineParser :: A.Parser BS.ByteString lineParser = do diff --git a/internal/Plumbing/Pipeline/Broadcasting.hs b/internal/Plumbing/Pipeline/Broadcasting.hs --- a/internal/Plumbing/Pipeline/Broadcasting.hs +++ b/internal/Plumbing/Pipeline/Broadcasting.hs @@ -12,9 +12,9 @@ import Queue.API (Queue) runBroadcasting :: (MonadLogger m , MonadIO m) - => Int -> Queue -> WriterAPI writer -> writer -> m () -runBroadcasting bufferLimit q writerApi writer = do - runConduit $ source q .| transformer bufferLimit .| sink writerApi writer + => Queue -> WriterAPI writer -> writer -> m () +runBroadcasting q writerApi writer = do + runConduit $ source q .| transformer .| sink writerApi writer broadcastingApi :: BroadcastingAPI broadcastingApi = BroadcastingAPI diff --git a/internal/Plumbing/Pipeline/Broadcasting/API.hs b/internal/Plumbing/Pipeline/Broadcasting/API.hs --- a/internal/Plumbing/Pipeline/Broadcasting/API.hs +++ b/internal/Plumbing/Pipeline/Broadcasting/API.hs @@ -9,4 +9,4 @@ import Network.ConnectionAPI (WriterAPI) import Queue.API (Queue) -newtype BroadcastingAPI = BroadcastingAPI { run :: forall m writer. (MonadLogger m, MonadIO m) => Int -> Queue -> WriterAPI writer -> writer -> m () } +newtype BroadcastingAPI = BroadcastingAPI { run :: forall m writer. (MonadLogger m, MonadIO m) => Queue -> WriterAPI writer -> writer -> m () } diff --git a/internal/Plumbing/Pipeline/Broadcasting/Transformer.hs b/internal/Plumbing/Pipeline/Broadcasting/Transformer.hs --- a/internal/Plumbing/Pipeline/Broadcasting/Transformer.hs +++ b/internal/Plumbing/Pipeline/Broadcasting/Transformer.hs @@ -2,17 +2,9 @@ import Conduit import qualified Data.ByteString.Lazy as LBS -import Lib.Logger.Types (LogLevel (..), MonadLogger (..)) +import Lib.Logger.Types (MonadLogger) import Transformers.Transformers transformer :: (MonadLogger m, MonadIO m, Transformer t) - => Int - -> ConduitT t LBS.ByteString m () -transformer bufferLimit = awaitForever $ \t -> do - let bytes = transform t - byteCount = fromIntegral (LBS.length bytes) - if byteCount > bufferLimit - then - lift . logMessage Error $ - "dropping outbound message: " ++ show byteCount ++ " bytes exceeds limit " ++ show bufferLimit - else yield bytes + => ConduitT t LBS.ByteString m () +transformer = awaitForever (yield . transform) diff --git a/internal/Plumbing/Pipeline/Streaming/Parser.hs b/internal/Plumbing/Pipeline/Streaming/Parser.hs --- a/internal/Plumbing/Pipeline/Streaming/Parser.hs +++ b/internal/Plumbing/Pipeline/Streaming/Parser.hs @@ -31,12 +31,11 @@ if bsLen > bufferLimit then do lift . logMessage Error $ - "overloaded buffer: " + "closing connection: incomplete protocol frame buffered " ++ show bsLen ++ " bytes exceeds limit " ++ show bufferLimit - lift . logMessage Debug $ ("invalid prefix: " ++ show (take bufferLimit bs)) - handleChunk (drop bufferLimit bs) + return () else do lift . logMessage Debug $ "message spans frame, waiting for more data" loop bs @@ -49,15 +48,5 @@ lift . logMessage Error $ ("parser rejected inbound data: " ++ reason) Emit message rest -> do lift . logMessage Debug $ "parsed message" - let consumedBytes = length bs - length rest - if consumedBytes > bufferLimit - then do - lift . logMessage Error $ - "dropping inbound message: " - ++ show consumedBytes - ++ " bytes exceeds limit " - ++ show bufferLimit - handleChunk rest - else do - yield message - handleChunk rest + yield message + handleChunk rest diff --git a/internal/Policy/Auth/Config.hs b/internal/Policy/Auth/Config.hs new file mode 100644 --- /dev/null +++ b/internal/Policy/Auth/Config.hs @@ -0,0 +1,35 @@ +module Auth.Config + ( authMethods + , mergeAuth + ) where + +import Auth.Types + +-- | Merge auth configuration while preserving independent token and user/pass +-- methods. Challenge auth (NKey or JWT) remains mutually exclusive and the +-- right-hand configuration wins. +mergeAuth :: Auth -> Auth -> Auth +mergeAuth left right = + Auth + { authTokenHandler = preferRight (authTokenHandler left) (authTokenHandler right) + , authUserPassHandler = preferRight (authUserPassHandler left) (authUserPassHandler right) + , authChallenge = preferRight (authChallenge left) (authChallenge right) + } + +preferRight :: Maybe a -> Maybe a -> Maybe a +preferRight left Nothing = left +preferRight _ right@(Just _) = right + +authMethods :: Auth -> [String] +authMethods auth = + tokenMethod ++ userPassMethod ++ challengeMethod + where + tokenMethod = ["auth token" | Just _ <- [authTokenHandler auth]] + userPassMethod = ["user/pass" | Just _ <- [authUserPassHandler auth]] + challengeMethod = + case authChallenge auth of + Nothing -> [] + Just (AuthNKeySeed _) -> ["nkey"] + Just (AuthNKeyHandler _ _) -> ["nkey"] + Just (AuthJWTBundle _) -> ["jwt"] + Just (AuthJWTHandlers _ _) -> ["jwt"] diff --git a/internal/Policy/Auth/Credentials.hs b/internal/Policy/Auth/Credentials.hs new file mode 100644 --- /dev/null +++ b/internal/Policy/Auth/Credentials.hs @@ -0,0 +1,55 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Auth.Credentials + ( JwtBundle (..) + , parseJwtBundle + ) where + +import qualified Data.ByteString as BS +import Data.Word (Word8) + +data JwtBundle = JwtBundle + { jwtToken :: BS.ByteString + , jwtSeed :: BS.ByteString + } + deriving (Eq) + +instance Show JwtBundle where + show _ = "JwtBundle {jwtToken = , jwtSeed = }" + +parseJwtBundle :: BS.ByteString -> Maybe JwtBundle +parseJwtBundle input = do + jwt <- extractBlock jwtStart jwtEnd input + seed <- extractBlock seedStart seedEnd input + pure (JwtBundle jwt seed) + +extractBlock :: BS.ByteString -> BS.ByteString -> BS.ByteString -> Maybe BS.ByteString +extractBlock startMarker endMarker input = do + let (_, rest) = BS.breakSubstring startMarker input + if BS.null rest + then Nothing + else do + let afterStart = BS.drop (BS.length startMarker) rest + (block, end) = BS.breakSubstring endMarker afterStart + trimmed = trimAscii block + if BS.null end || BS.null trimmed + then Nothing + else Just trimmed + +jwtStart, jwtEnd, seedStart, seedEnd :: BS.ByteString +jwtStart = "-----BEGIN NATS USER JWT-----" +jwtEnd = "------END NATS USER JWT------" +seedStart = "-----BEGIN USER NKEY SEED-----" +seedEnd = "------END USER NKEY SEED------" + +trimAscii :: BS.ByteString -> BS.ByteString +trimAscii = + dropWhileEndAscii isSpaceAscii . BS.dropWhile isSpaceAscii + +dropWhileEndAscii :: (Word8 -> Bool) -> BS.ByteString -> BS.ByteString +dropWhileEndAscii predicate = + BS.reverse . BS.dropWhile predicate . BS.reverse + +isSpaceAscii :: Word8 -> Bool +isSpaceAscii w = + w == 9 || w == 10 || w == 13 || w == 32 diff --git a/internal/Policy/Auth/Jwt.hs b/internal/Policy/Auth/Jwt.hs --- a/internal/Policy/Auth/Jwt.hs +++ b/internal/Policy/Auth/Jwt.hs @@ -1,10 +1,16 @@ module Auth.Jwt ( JwtBundle (..) , auth + , authHandlers , parseJwtBundle ) where +import Auth.Credentials (JwtBundle (..), parseJwtBundle) import Auth.Types auth :: JWTTokenData -> Auth -auth = AuthJWT +auth creds = emptyAuth { authChallenge = Just (AuthJWTBundle creds) } + +authHandlers :: JWTHandler -> SignatureHandler -> Auth +authHandlers jwtHandler signatureHandler = + emptyAuth { authChallenge = Just (AuthJWTHandlers jwtHandler signatureHandler) } diff --git a/internal/Policy/Auth/NKey.hs b/internal/Policy/Auth/NKey.hs --- a/internal/Policy/Auth/NKey.hs +++ b/internal/Policy/Auth/NKey.hs @@ -1,9 +1,20 @@ module Auth.NKey ( auth + , authHandler , signNonceWithSeed ) where +import qualified Auth.NKey.Codec as Codec import Auth.Types auth :: NKeyData -> Auth -auth = AuthNKey +auth seed = emptyAuth { authChallenge = Just (AuthNKeySeed seed) } + +authHandler :: NKeyPublicKey -> SignatureHandler -> Auth +authHandler publicKey handler = + emptyAuth { authChallenge = Just (AuthNKeyHandler publicKey handler) } + +signNonceWithSeed :: NKeyData -> Nonce -> Either String (NKeyPublicKey, Signature) +signNonceWithSeed seed nonce = do + (publicKey, signature) <- Codec.signNonceWithSeedRaw seed nonce + pure (publicKey, Codec.encodeSignature signature) diff --git a/internal/Policy/Auth/NKey/Codec.hs b/internal/Policy/Auth/NKey/Codec.hs new file mode 100644 --- /dev/null +++ b/internal/Policy/Auth/NKey/Codec.hs @@ -0,0 +1,139 @@ +module Auth.NKey.Codec + ( encodeSignature + , signNonceWithSeedRaw + , validateUserPublicKey + ) where + +import qualified Crypto.Error as Crypto +import qualified Crypto.PubKey.Ed25519 as Ed25519 +import qualified Data.Bits as Bits +import qualified Data.ByteArray as ByteArray +import qualified Data.ByteArray.Encoding as Encoding +import qualified Data.ByteString as BS +import qualified Data.List as List +import Data.Word (Word16, Word8) + +signNonceWithSeedRaw + :: BS.ByteString + -> BS.ByteString + -> Either String (BS.ByteString, BS.ByteString) +signNonceWithSeedRaw seed nonce = do + seedBytes <- decodeSeed seed + secretKey <- toSecretKey seedBytes + let publicKey = Ed25519.toPublic secretKey + signature = Ed25519.sign secretKey publicKey nonce + signatureBytes = ByteArray.convert signature + publicBytes = ByteArray.convert publicKey + pure (encodePublicKey publicBytes, signatureBytes) + +encodeSignature :: BS.ByteString -> BS.ByteString +encodeSignature = Encoding.convertToBase Encoding.Base64URLUnpadded + +validateUserPublicKey :: BS.ByteString -> Either String () +validateUserPublicKey encoded = do + raw <- decodeBase32Canonical encoded + payload <- verifyChecksum "public key" raw + if BS.length payload /= 33 + then Left "public key length is invalid" + else + if BS.head payload /= prefixByteUser + then Left "public key is not a user nkey" + else Right () + +decodeSeed :: BS.ByteString -> Either String BS.ByteString +decodeSeed encoded = do + raw <- decodeBase32Canonical (trimAscii encoded) + payload <- verifyChecksum "seed" raw + if BS.length payload /= 34 + then Left "seed length is invalid" + else do + let prefix0 = BS.index payload 0 + prefix1 = BS.index payload 1 + seedPrefix = prefix0 Bits..&. 0xF8 + publicPrefix = + ((prefix0 Bits..&. 0x07) `Bits.shiftL` 5) + Bits..|. (prefix1 `Bits.shiftR` 3) + if seedPrefix /= prefixByteSeed + then Left "seed prefix mismatch" + else + if publicPrefix /= prefixByteUser + then Left "seed is not a user nkey" + else Right (BS.drop 2 payload) + +decodeBase32Canonical :: BS.ByteString -> Either String BS.ByteString +decodeBase32Canonical encoded = do + let paddingLength = (8 - BS.length encoded `mod` 8) `mod` 8 + padded = encoded <> BS.replicate paddingLength 61 + decoded <- + case Encoding.convertFromBase Encoding.Base32 padded of + Left err -> Left ("invalid base32 encoding: " ++ err) + Right bytes -> Right bytes + let canonical = BS.takeWhile (/= 61) (Encoding.convertToBase Encoding.Base32 decoded) + if canonical == encoded + then Right decoded + else Left "base32 encoding is not canonical" + +verifyChecksum :: String -> BS.ByteString -> Either String BS.ByteString +verifyChecksum label raw + | BS.length raw < 3 = Left (label ++ " is too short") + | otherwise = do + let (payload, checksumBytes) = BS.splitAt (BS.length raw - 2) raw + expected <- decodeChecksum checksumBytes + if crc16 payload == expected + then Right payload + else Left (label ++ " checksum mismatch") + +toSecretKey :: BS.ByteString -> Either String Ed25519.SecretKey +toSecretKey raw = + case Ed25519.secretKey raw of + Crypto.CryptoPassed key -> Right key + Crypto.CryptoFailed _ -> Left "seed could not be parsed as a secret key" + +encodePublicKey :: BS.ByteString -> BS.ByteString +encodePublicKey raw = + let payload = BS.cons prefixByteUser raw + checksum = crc16 payload + in Encoding.convertToBase Encoding.Base32 (payload <> encodeChecksum checksum) + +encodeChecksum :: Word16 -> BS.ByteString +encodeChecksum value = + BS.pack + [ fromIntegral (value Bits..&. 0xFF) + , fromIntegral (value `Bits.shiftR` 8) + ] + +decodeChecksum :: BS.ByteString -> Either String Word16 +decodeChecksum bytes + | BS.length bytes /= 2 = Left "checksum length is invalid" + | otherwise = + let b0 = fromIntegral (BS.index bytes 0) + b1 = fromIntegral (BS.index bytes 1) + in Right (b0 Bits..|. (b1 `Bits.shiftL` 8)) + +crc16 :: BS.ByteString -> Word16 +crc16 = + List.foldl' update 0 . BS.unpack + where + update crc byte = + List.foldl' step (crc `Bits.xor` (fromIntegral byte `Bits.shiftL` 8)) [0 .. 7] + + step crc _ = + if Bits.testBit crc 15 + then (crc `Bits.shiftL` 1) `Bits.xor` 0x1021 + else crc `Bits.shiftL` 1 + +trimAscii :: BS.ByteString -> BS.ByteString +trimAscii = + dropWhileEndAscii isSpaceAscii . BS.dropWhile isSpaceAscii + +dropWhileEndAscii :: (Word8 -> Bool) -> BS.ByteString -> BS.ByteString +dropWhileEndAscii predicate = + BS.reverse . BS.dropWhile predicate . BS.reverse + +isSpaceAscii :: Word8 -> Bool +isSpaceAscii w = + w == 9 || w == 10 || w == 13 || w == 32 + +prefixByteSeed, prefixByteUser :: Word8 +prefixByteSeed = 18 `Bits.shiftL` 3 +prefixByteUser = 20 `Bits.shiftL` 3 diff --git a/internal/Policy/Auth/None.hs b/internal/Policy/Auth/None.hs --- a/internal/Policy/Auth/None.hs +++ b/internal/Policy/Auth/None.hs @@ -5,4 +5,4 @@ import Auth.Types auth :: Auth -auth = AuthNone +auth = emptyAuth diff --git a/internal/Policy/Auth/Resolver.hs b/internal/Policy/Auth/Resolver.hs new file mode 100644 --- /dev/null +++ b/internal/Policy/Auth/Resolver.hs @@ -0,0 +1,157 @@ +{-# LANGUAGE TypeApplications #-} + +module Auth.Resolver + ( applyAuthPatch + , buildAuthPatch + ) where + +import Auth.Credentials (JwtBundle (..), parseJwtBundle) +import qualified Auth.NKey.Codec as NKeyCodec +import Auth.Types +import Control.Exception + ( SomeAsyncException + , SomeException + , displayException + , fromException + , throwIO + , try + ) +import qualified Data.ByteString as BS +import qualified Types.Connect as Connect + +buildAuthPatch :: Auth -> AuthContext -> IO (Either AuthError AuthPatch) +buildAuthPatch auth ctx = do + tokenResult <- resolveOptional "auth token" validateToken (authTokenHandler auth) + userPassResult <- resolveOptional "user/pass" validateUserPass (authUserPassHandler auth) + challengeResult <- buildChallengePatch (authChallenge auth) ctx + pure $ do + token <- tokenResult + userPass <- userPassResult + challengePatch <- challengeResult + pure challengePatch + { patchAuthToken = token + , patchUser = fst <$> userPass + , patchPass = snd <$> userPass + } + +resolveOptional + :: String + -> (a -> Either AuthError ()) + -> Maybe (IO (Either AuthError a)) + -> IO (Either AuthError (Maybe a)) +resolveOptional _ _ Nothing = pure (Right Nothing) +resolveOptional label validate (Just handler) = do + result <- runAuthHandler label handler + pure $ do + value <- result + validate value + pure (Just value) + +buildChallengePatch :: Maybe ChallengeAuth -> AuthContext -> IO (Either AuthError AuthPatch) +buildChallengePatch Nothing _ = pure (Right emptyAuthPatch) +buildChallengePatch (Just challenge) ctx = + case requireNonce ctx of + Left err -> pure (Left err) + Right nonce -> + case challenge of + AuthNKeySeed seed -> + pure $ do + (publicKey, signature) <- signNonceWithSeedRaw seed nonce + pure emptyAuthPatch + { patchNKey = Just publicKey + , patchSig = Just (NKeyCodec.encodeSignature signature) + } + AuthNKeyHandler publicKey signatureHandler -> + case validatePublicKey publicKey of + Left err -> pure (Left err) + Right () -> do + signatureResult <- resolveSignature signatureHandler nonce + pure $ do + signature <- signatureResult + pure emptyAuthPatch + { patchNKey = Just publicKey + , patchSig = Just signature + } + AuthJWTBundle creds -> + pure $ do + bundle <- + maybe (Left (AuthError "jwt credentials bundle is invalid")) Right (parseJwtBundle creds) + (_, signature) <- signNonceWithSeedRaw (jwtSeed bundle) nonce + pure emptyAuthPatch + { patchJwt = Just (jwtToken bundle) + , patchSig = Just (NKeyCodec.encodeSignature signature) + } + AuthJWTHandlers jwtHandler signatureHandler -> do + jwtResult <- runAuthHandler "jwt" jwtHandler + case jwtResult >>= validateRawJwtValue of + Left err -> pure (Left err) + Right jwt -> do + signatureResult <- resolveSignature signatureHandler nonce + pure $ do + signature <- signatureResult + pure emptyAuthPatch + { patchJwt = Just jwt + , patchSig = Just signature + } + +signNonceWithSeedRaw + :: BS.ByteString + -> BS.ByteString + -> Either AuthError (BS.ByteString, BS.ByteString) +signNonceWithSeedRaw seed nonce = + either (Left . AuthError) Right (NKeyCodec.signNonceWithSeedRaw seed nonce) + +resolveSignature :: SignatureHandler -> Nonce -> IO (Either AuthError BS.ByteString) +resolveSignature handler nonce = do + result <- runAuthHandler "nonce signature" (handler nonce) + pure $ do + signature <- result + if BS.length signature == 64 + then Right (NKeyCodec.encodeSignature signature) + else Left (AuthError "nonce signature must contain 64 raw bytes") + +runAuthHandler :: String -> IO (Either AuthError a) -> IO (Either AuthError a) +runAuthHandler label action = do + result <- try @SomeException action + case result of + Right value -> pure value + Left err -> + case fromException err :: Maybe SomeAsyncException of + Just _ -> throwIO err + Nothing -> pure (Left (AuthError (label ++ " handler failed: " ++ displayException err))) + +applyAuthPatch :: AuthPatch -> Connect.Connect -> Connect.Connect +applyAuthPatch patch connect = + connect + { Connect.auth_token = patchAuthToken patch + , Connect.user = patchUser patch + , Connect.pass = patchPass patch + , Connect.jwt = patchJwt patch + , Connect.nkey = patchNKey patch + , Connect.sig = patchSig patch + } + +validateToken :: AuthTokenData -> Either AuthError () +validateToken token + | BS.null token = Left (AuthError "auth token must not be empty") + | otherwise = Right () + +validateUserPass :: UserPassData -> Either AuthError () +validateUserPass (user, pass) + | BS.null user = Left (AuthError "user must not be empty") + | BS.null pass = Left (AuthError "pass must not be empty") + | otherwise = Right () + +validatePublicKey :: NKeyPublicKey -> Either AuthError () +validatePublicKey = + either (Left . AuthError . ("nkey public key is invalid: " ++)) Right + . NKeyCodec.validateUserPublicKey + +validateRawJwtValue :: JWTTokenData -> Either AuthError JWTTokenData +validateRawJwtValue token + | BS.null token = Left (AuthError "jwt must not be empty") + | otherwise = Right token + +requireNonce :: AuthContext -> Either AuthError BS.ByteString +requireNonce (AuthContext maybeNonce) = + maybe (Left (AuthError "auth method requires a server nonce")) Right maybeNonce diff --git a/internal/Policy/Auth/Token.hs b/internal/Policy/Auth/Token.hs --- a/internal/Policy/Auth/Token.hs +++ b/internal/Policy/Auth/Token.hs @@ -1,8 +1,12 @@ module Auth.Token ( auth + , authHandler ) where import Auth.Types auth :: AuthTokenData -> Auth -auth = AuthToken +auth = authHandler . pure . Right + +authHandler :: AuthTokenHandler -> Auth +authHandler handler = emptyAuth { authTokenHandler = Just handler } diff --git a/internal/Policy/Auth/Types.hs b/internal/Policy/Auth/Types.hs --- a/internal/Policy/Auth/Types.hs +++ b/internal/Policy/Auth/Types.hs @@ -1,51 +1,47 @@ -{-# LANGUAGE OverloadedStrings #-} - module Auth.Types ( User , Pass , UserPassData , NKeyData + , NKeyPublicKey , AuthTokenData , JWTTokenData - , JwtBundle (..) + , Nonce + , Signature + , AuthTokenHandler + , UserPassHandler + , JWTHandler + , SignatureHandler , AuthError (..) , AuthContext (..) , AuthPatch (..) , Auth (..) + , ChallengeAuth (..) + , emptyAuth , emptyAuthPatch - , validateAuth - , buildAuthPatch - , applyAuthPatch - , parseJwtBundle - , signNonceWithSeed ) where -import qualified Crypto.Error as Crypto -import qualified Crypto.PubKey.Ed25519 as Ed25519 -import qualified Data.Bits as Bits -import qualified Data.ByteArray as ByteArray -import qualified Data.ByteString as BS -import Data.Char (ord) -import qualified Data.List as List -import Data.Word (Word16, Word32, Word8) -import qualified Types.Connect as Connect +import qualified Data.ByteString as BS type User = BS.ByteString type Pass = BS.ByteString type UserPassData = (User, Pass) type NKeyData = BS.ByteString +type NKeyPublicKey = BS.ByteString type AuthTokenData = BS.ByteString type JWTTokenData = BS.ByteString -data JwtBundle = JwtBundle - { jwtToken :: BS.ByteString - , jwtSeed :: BS.ByteString - } - deriving (Eq, Show) +type Nonce = BS.ByteString +type Signature = BS.ByteString +type AuthTokenHandler = IO (Either AuthError AuthTokenData) +type UserPassHandler = IO (Either AuthError UserPassData) +type JWTHandler = IO (Either AuthError JWTTokenData) +type SignatureHandler = Nonce -> IO (Either AuthError Signature) + newtype AuthError = AuthError String deriving (Eq, Show) @@ -60,15 +56,41 @@ , patchNKey :: Maybe BS.ByteString , patchSig :: Maybe BS.ByteString } - deriving (Eq, Show) + deriving (Eq) -data Auth = AuthNone - | AuthToken AuthTokenData - | AuthUserPass UserPassData - | AuthNKey NKeyData - | AuthJWT JWTTokenData - deriving (Eq, Show) +instance Show AuthPatch where + show patch = + "AuthPatch {patchAuthToken = " + ++ present (patchAuthToken patch) + ++ ", patchUser = " + ++ present (patchUser patch) + ++ ", patchPass = " + ++ present (patchPass patch) + ++ ", patchJwt = " + ++ present (patchJwt patch) + ++ ", patchNKey = " + ++ present (patchNKey patch) + ++ ", patchSig = " + ++ present (patchSig patch) + ++ "}" + where + present Nothing = "Nothing" + present (Just _) = "" +data ChallengeAuth = AuthNKeySeed NKeyData + | AuthNKeyHandler NKeyPublicKey SignatureHandler + | AuthJWTBundle JWTTokenData + | AuthJWTHandlers JWTHandler SignatureHandler + +data Auth = Auth + { authTokenHandler :: Maybe AuthTokenHandler + , authUserPassHandler :: Maybe UserPassHandler + , authChallenge :: Maybe ChallengeAuth + } + +emptyAuth :: Auth +emptyAuth = Auth Nothing Nothing Nothing + emptyAuthPatch :: AuthPatch emptyAuthPatch = AuthPatch @@ -79,317 +101,3 @@ , patchNKey = Nothing , patchSig = Nothing } - -validateAuth :: Auth -> Either AuthError () -validateAuth auth = - case auth of - AuthNone -> - Right () - AuthToken token -> - validateToken token - AuthUserPass userPass -> - validateUserPass userPass - AuthNKey seed -> - validateNKey seed - AuthJWT creds -> - validateJwt creds - -buildAuthPatch :: Auth -> AuthContext -> Either AuthError AuthPatch -buildAuthPatch auth ctx = - case auth of - AuthNone -> - Right emptyAuthPatch - AuthToken token -> do - validateToken token - Right emptyAuthPatch - { patchAuthToken = Just token - } - AuthUserPass (user, pass) -> do - validateUserPass (user, pass) - Right emptyAuthPatch - { patchUser = Just user - , patchPass = Just pass - } - AuthNKey seed -> do - validateNKey seed - nonce <- requireNonce ctx - (publicKey, signature) <- either (Left . AuthError) Right (signNonceWithSeed seed nonce) - Right emptyAuthPatch - { patchNKey = Just publicKey - , patchSig = Just signature - } - AuthJWT creds -> do - validateJwt creds - bundle <- - case parseJwtBundle creds of - Nothing -> Left (AuthError "jwt credentials bundle is invalid") - Just parsed -> Right parsed - nonce <- requireNonce ctx - (_, signature) <- either (Left . AuthError) Right (signNonceWithSeed (jwtSeed bundle) nonce) - Right emptyAuthPatch - { patchJwt = Just (jwtToken bundle) - , patchSig = Just signature - } - -applyAuthPatch :: AuthPatch -> Connect.Connect -> Connect.Connect -applyAuthPatch patch connect = - connect - { Connect.auth_token = patchAuthToken patch - , Connect.user = patchUser patch - , Connect.pass = patchPass patch - , Connect.jwt = patchJwt patch - , Connect.nkey = patchNKey patch - , Connect.sig = patchSig patch - } - -validateToken :: AuthTokenData -> Either AuthError () -validateToken token - | BS.null token = Left (AuthError "auth token must not be empty") - | otherwise = Right () - -validateUserPass :: UserPassData -> Either AuthError () -validateUserPass (user, pass) - | BS.null user = Left (AuthError "user must not be empty") - | BS.null pass = Left (AuthError "pass must not be empty") - | otherwise = Right () - -validateNKey :: NKeyData -> Either AuthError () -validateNKey seed - | BS.null seed = Left (AuthError "nkey seed must not be empty") - | otherwise = Right () - -validateJwt :: JWTTokenData -> Either AuthError () -validateJwt creds = - case parseJwtBundle creds of - Nothing -> Left (AuthError "jwt credentials bundle is invalid") - Just _ -> Right () - -requireNonce :: AuthContext -> Either AuthError BS.ByteString -requireNonce (AuthContext maybeNonce) = - maybe (Left (AuthError "auth method requires a server nonce")) Right maybeNonce - -parseJwtBundle :: BS.ByteString -> Maybe JwtBundle -parseJwtBundle input = - fmap (uncurry JwtBundle) (parseCreds input) - -parseCreds :: BS.ByteString -> Maybe (BS.ByteString, BS.ByteString) -parseCreds input = do - jwt <- extractBlock jwtStart jwtEnd input - seed <- extractBlock seedStart seedEnd input - pure (jwt, seed) - -extractBlock :: BS.ByteString -> BS.ByteString -> BS.ByteString -> Maybe BS.ByteString -extractBlock startMarker endMarker input = do - let (_, rest) = BS.breakSubstring startMarker input - if BS.null rest - then Nothing - else do - let afterStart = BS.drop (BS.length startMarker) rest - (block, _) = BS.breakSubstring endMarker afterStart - if BS.null block - then Nothing - else Just (trimAscii block) - -jwtStart, jwtEnd, seedStart, seedEnd :: BS.ByteString -jwtStart = "-----BEGIN NATS USER JWT-----" -jwtEnd = "------END NATS USER JWT------" -seedStart = "-----BEGIN USER NKEY SEED-----" -seedEnd = "------END USER NKEY SEED------" - -signNonceWithSeed :: BS.ByteString -> BS.ByteString -> Either String (BS.ByteString, BS.ByteString) -signNonceWithSeed seed nonce = do - seedBytes <- decodeSeed seed - secretKey <- toSecretKey seedBytes - let publicKey = Ed25519.toPublic secretKey - signature = Ed25519.sign secretKey publicKey nonce - sigEncoded = encodeBase64Url (ByteArray.convert signature :: BS.ByteString) - publicEncoded = encodePublicKey (ByteArray.convert publicKey :: BS.ByteString) - pure (publicEncoded, sigEncoded) - -decodeSeed :: BS.ByteString -> Either String BS.ByteString -decodeSeed encoded = do - raw <- decodeBase32 (trimAscii encoded) - if BS.length raw < 4 - then Left "seed is too short" - else do - let (payload, checksumBytes) = BS.splitAt (BS.length raw - 2) raw - expected <- decodeChecksum checksumBytes - let actual = crc16 payload - if actual /= expected - then Left "seed checksum mismatch" - else do - let prefix0 = BS.index payload 0 - prefix1 = BS.index payload 1 - seedPrefix = prefix0 Bits..&. 0xF8 - if seedPrefix /= prefixByteSeed - then Left "seed prefix mismatch" - else do - let typ = ((prefix0 Bits..&. 0x07) `Bits.shiftL` 5) Bits..|. (prefix1 `Bits.shiftR` 3) - if typ /= prefixByteUser - then Left "seed is not a user nkey" - else do - let seedBytes = BS.drop 2 payload - if BS.length seedBytes /= 32 - then Left "seed length is invalid" - else pure seedBytes - -toSecretKey :: BS.ByteString -> Either String Ed25519.SecretKey -toSecretKey raw = - case Ed25519.secretKey raw of - Crypto.CryptoPassed key -> Right key - Crypto.CryptoFailed _ -> Left "seed could not be parsed as a secret key" - -trimAscii :: BS.ByteString -> BS.ByteString -trimAscii = - dropWhileEndAscii isSpaceAscii . BS.dropWhile isSpaceAscii - -dropWhileEndAscii :: (Word8 -> Bool) -> BS.ByteString -> BS.ByteString -dropWhileEndAscii predicate = - BS.reverse . BS.dropWhile predicate . BS.reverse - -isSpaceAscii :: Word8 -> Bool -isSpaceAscii w = - w == 9 || w == 10 || w == 13 || w == 32 - -encodePublicKey :: BS.ByteString -> BS.ByteString -encodePublicKey raw = - let payload = BS.cons prefixByteUser raw - checksum = crc16 payload - in encodeBase32 (payload <> encodeChecksum checksum) - -decodeBase32 :: BS.ByteString -> Either String BS.ByteString -decodeBase32 input = - let cleaned = BS.map toUpperAscii input - in go 0 0 [] (BS.unpack cleaned) - where - go :: Word32 -> Int -> [Word8] -> [Word8] -> Either String BS.ByteString - go _ _ acc [] = - Right (BS.pack (reverse acc)) - go buffer bits acc (x : xs) = - case base32Value x of - Nothing -> Left "invalid base32 character" - Just value -> do - let buffer' = (buffer `Bits.shiftL` 5) Bits..|. fromIntegral value - bits' = bits + 5 - (acc', buffer'', bits'') = flush buffer' bits' acc - go buffer'' bits'' acc' xs - - flush :: Word32 -> Int -> [Word8] -> ([Word8], Word32, Int) - flush buffer bits acc - | bits >= 8 = - let bits' = bits - 8 - byte = fromIntegral ((buffer `Bits.shiftR` bits') Bits..&. 0xFF) - in flush buffer bits' (byte : acc) - | otherwise = - (acc, buffer, bits) - -encodeBase32 :: BS.ByteString -> BS.ByteString -encodeBase32 input = - let (acc, buffer, bits) = List.foldl' step ([], 0, 0) (BS.unpack input) - acc' = if bits == 0 then acc else encodeRemaining acc buffer bits - in BS.pack (reverse acc') - where - step :: ([Word8], Word32, Int) -> Word8 -> ([Word8], Word32, Int) - step (out, buffer, bits) byte = - let buffer' = (buffer `Bits.shiftL` 8) Bits..|. fromIntegral byte - bits' = bits + 8 - (out', buffer'', bits'') = emit out buffer' bits' - in (out', buffer'', bits'') - - emit :: [Word8] -> Word32 -> Int -> ([Word8], Word32, Int) - emit out buffer bits - | bits >= 5 = - let bits' = bits - 5 - idx = fromIntegral ((buffer `Bits.shiftR` bits') Bits..&. 0x1F) - char = base32Alphabet idx - in emit (char : out) buffer bits' - | otherwise = - (out, buffer, bits) - - encodeRemaining :: [Word8] -> Word32 -> Int -> [Word8] - encodeRemaining out buffer bits = - let idx = fromIntegral ((buffer `Bits.shiftL` (5 - bits)) Bits..&. 0x1F) - in base32Alphabet idx : out - -base32Alphabet :: Word8 -> Word8 -base32Alphabet idx = - if idx < 26 - then fromIntegral (ord 'A' + fromIntegral idx) - else fromIntegral (ord '2' + fromIntegral idx - 26) - -base32Value :: Word8 -> Maybe Word8 -base32Value w - | w >= 65 && w <= 90 = Just (w - 65) - | w >= 50 && w <= 55 = Just (w - 24) - | otherwise = Nothing - -toUpperAscii :: Word8 -> Word8 -toUpperAscii w - | w >= 97 && w <= 122 = w - 32 - | otherwise = w - -encodeChecksum :: Word16 -> BS.ByteString -encodeChecksum value = - BS.pack - [ fromIntegral (value Bits..&. 0xFF) - , fromIntegral (value `Bits.shiftR` 8) - ] - -decodeChecksum :: BS.ByteString -> Either String Word16 -decodeChecksum bytes - | BS.length bytes /= 2 = Left "checksum length is invalid" - | otherwise = - let b0 = fromIntegral (BS.index bytes 0) - b1 = fromIntegral (BS.index bytes 1) - in Right (b0 Bits..|. (b1 `Bits.shiftL` 8)) - -crc16 :: BS.ByteString -> Word16 -crc16 = - List.foldl' update 0 . BS.unpack - where - update crc byte = - List.foldl' step (crc `Bits.xor` (fromIntegral byte `Bits.shiftL` 8)) [0 .. 7] - - step crc _ = - if Bits.testBit crc 15 - then (crc `Bits.shiftL` 1) `Bits.xor` 0x1021 - else crc `Bits.shiftL` 1 - -encodeBase64Url :: BS.ByteString -> BS.ByteString -encodeBase64Url input = - let (out, buffer, bits) = List.foldl' step ([], 0 :: Word32, 0 :: Int) (BS.unpack input) - out' = flush out buffer bits - in BS.pack (reverse out') - where - step (acc, buffer, bits) byte = - let buffer' = (buffer `Bits.shiftL` 8) Bits..|. fromIntegral byte - bits' = bits + 8 - (acc', buffer'', bits'') = emit acc buffer' bits' - in (acc', buffer'', bits'') - - emit acc buffer bits - | bits >= 6 = - let bits' = bits - 6 - idx = fromIntegral ((buffer `Bits.shiftR` bits') Bits..&. 0x3F) - char = base64UrlAlphabet idx - in emit (char : acc) buffer bits' - | otherwise = - (acc, buffer, bits) - - flush acc buffer bits - | bits == 0 = acc - | otherwise = - let idx = fromIntegral ((buffer `Bits.shiftL` (6 - bits)) Bits..&. 0x3F) - in base64UrlAlphabet idx : acc - -base64UrlAlphabet :: Word8 -> Word8 -base64UrlAlphabet idx - | idx < 26 = fromIntegral (ord 'A' + fromIntegral idx) - | idx < 52 = fromIntegral (ord 'a' + fromIntegral idx - 26) - | idx < 62 = fromIntegral (ord '0' + fromIntegral idx - 52) - | idx == 62 = 45 - | otherwise = 95 - -prefixByteSeed, prefixByteUser :: Word8 -prefixByteSeed = 18 `Bits.shiftL` 3 -prefixByteUser = 20 `Bits.shiftL` 3 diff --git a/internal/Policy/Auth/UserPass.hs b/internal/Policy/Auth/UserPass.hs --- a/internal/Policy/Auth/UserPass.hs +++ b/internal/Policy/Auth/UserPass.hs @@ -1,8 +1,12 @@ module Auth.UserPass ( auth + , authHandler ) where import Auth.Types auth :: UserPassData -> Auth -auth = AuthUserPass +auth = authHandler . pure . Right + +authHandler :: UserPassHandler -> Auth +authHandler handler = emptyAuth { authUserPassHandler = Just handler } diff --git a/internal/Policy/Engine.hs b/internal/Policy/Engine.hs --- a/internal/Policy/Engine.hs +++ b/internal/Policy/Engine.hs @@ -9,7 +9,11 @@ import Control.Concurrent.STM import Control.Monad (forM_, unless, void, when) import Data.Foldable (for_) -import Handshake.Nats (performHandshake) +import Data.Maybe (listToMaybe) +import Handshake.Nats + ( HandshakeError (..) + , performHandshake + ) import Lib.Logger (LogLevel (..), MonadLogger (..)) import Network.ConnectionAPI ( Conn @@ -33,6 +37,7 @@ , config , connection , enqueue + , failInitialConnection , incrementAttemptIndex , markClosed , markConnected @@ -52,6 +57,9 @@ ( ClientConfig (..) , ClientExitReason (..) , ClientStatus (..) + , ConnectAttemptError (..) + , ConnectError (..) + , ConnectFailure (..) ) import Subscription.Store ( SubscriptionStore @@ -75,15 +83,24 @@ -> Auth -> IO () runEngine connectionApi streamingApi broadcastingApi parserApi state store auth = - loop (connectionAttempts (config state)) Nothing + case connectOptions (config state) of + [] -> do + failInitialConnection state ConnectNoServers + finalize (ExitRetriesExhausted (Just "No servers provided")) + _ -> + loop (connectionAttempts (config state)) [] where StreamingAPI runStreaming = streamingApi BroadcastingAPI runBroadcasting = broadcastingApi - loop remaining lastErr + loop remaining attemptErrors | remaining <= 0 = do runClient state $ logMessage Info "retries exhausted; exiting" + let orderedErrors = reverse attemptErrors + connectError = ConnectAttemptsExhausted orderedErrors + lastErr = show <$> listToMaybe orderedErrors + failInitialConnection state connectError finalize (ExitRetriesExhausted lastErr) | otherwise = do status <- readStatus state @@ -97,7 +114,7 @@ runClient state $ logMessage Info "retrying client connection" incrementAttemptIndex state - loop (remaining - 1) (Just attemptErr) + loop (remaining - 1) (attemptErr : attemptErrors) Closing reason -> finalize reason Closed _ -> @@ -110,21 +127,27 @@ runAttempt = do openQueue state open connectionApi (connection state) - transportResult <- acquireTransport + let cfg = config state + endpoints = connectOptions cfg + attemptIndex <- readAttemptIndex state + let endpoint@(host, _port) = + endpoints !! (attemptIndex `mod` length endpoints) + setEndpoint state endpoint + transportResult <- acquireTransport endpoint case transportResult of Left err -> do runClient state $ logMessage Error ("connection attempt failed: " ++ show err) close connectionApi (connection state) - pure err - Right (conn, (host, _port)) -> do + pure (ConnectAttemptError endpoint (ConnectTransportFailure err)) + Right conn -> do handshakeResult <- performHandshake connectionApi parserApi state auth conn host case handshakeResult of Left err -> do runClient state $ logMessage Error ("connection initialization failed: " ++ show err) close connectionApi conn - pure (show err) + pure (ConnectAttemptError endpoint (handshakeFailure err)) Right () -> do resubscribeIfNeeded markConnectionReady state @@ -134,27 +157,24 @@ ConnectionDisconnected -> do runClient state $ logMessage Info "connection disconnected" - pure "connection disconnected" + pure (ConnectAttemptError endpoint (ConnectTransportFailure "connection disconnected")) ConnectionExit reason -> do setClosing state reason - pure (show reason) + pure (ConnectAttemptError endpoint (ConnectProtocolFailure (show reason))) - acquireTransport = do - let cfg = config state - conn = connection state - case connectOptions cfg of - [] -> - pure (Left "No servers provided") - endpoints -> do - attemptIndex <- readAttemptIndex state - let endpoint@(host, port) = - endpoints !! (attemptIndex `mod` length endpoints) - setEndpoint state endpoint - result <- connectTcp connectionApi conn host port - case result of - Left err -> pure (Left err) - Right () -> pure (Right (conn, endpoint)) + acquireTransport (host, port) = do + let conn = connection state + result <- connectTcp connectionApi conn host port + case result of + Left err -> pure (Left err) + Right () -> pure (Right conn) + handshakeFailure (HandshakeTransportError err) = ConnectTransportFailure err + handshakeFailure (HandshakeTLSError err) = ConnectTLSFailure err + handshakeFailure (HandshakeProtocolError err) = ConnectProtocolFailure err + handshakeFailure (HandshakeAuthError err) = ConnectAuthenticationFailure (show err) + handshakeFailure HandshakeTimeout = ConnectHandshakeTimeout + resubscribeIfNeeded = do alreadyConnected <- markConnected state when alreadyConnected $ do @@ -193,7 +213,6 @@ runClient state $ do logMessage Debug "starting broadcasting thread" runBroadcasting - (bufferLimit (config state)) (queue state) (writer connectionApi) conn @@ -206,7 +225,7 @@ runClient state $ do logMessage Debug "starting streaming thread" runStreaming - (bufferLimit (config state)) + (streamBufferLimit (messageLimit (config state))) (reader connectionApi) conn parserApi @@ -239,6 +258,13 @@ isResumable :: (a, SubscriptionMeta) -> Bool isResumable (_, SubscriptionMeta _ _ isReply) = not isReply + + streamBufferLimit :: Int -> Int + streamBufferLimit maximumMessageSize + | maximumMessageSize > maxBound - controlLineAllowance = maxBound + | otherwise = maximumMessageSize + controlLineAllowance + + controlLineAllowance = 64 * 1024 closeClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO () closeClient connectionApi state store = diff --git a/internal/Policy/Handshake/Nats.hs b/internal/Policy/Handshake/Nats.hs --- a/internal/Policy/Handshake/Nats.hs +++ b/internal/Policy/Handshake/Nats.hs @@ -3,8 +3,8 @@ , performHandshake ) where +import Auth.Resolver (applyAuthPatch, buildAuthPatch) import Auth.Types -import Control.Monad (unless) import qualified Data.ByteString as BS import Data.Maybe (fromMaybe, isJust) import Network.ConnectionAPI @@ -30,74 +30,93 @@ , updateLogContextFromInfo ) import State.Types (ClientConfig (..)) +import System.Timeout (timeout) import Transformers.Transformers (Transformer (transform)) import qualified Types.Connect as Connect +import qualified Types.Err as Err import qualified Types.Info as I +import Types.Ping (Ping (Ping)) +import Types.Pong (Pong (Pong)) +import Types.TLS (defaultTLSConfig) data HandshakeError = HandshakeTransportError String + | HandshakeTLSError String | HandshakeProtocolError String | HandshakeAuthError AuthError + | HandshakeTimeout deriving (Eq, Show) performHandshake :: ConnectionAPI -> ParserAPI ParsedMessage -> ClientState -> Auth -> Conn -> String -> IO (Either HandshakeError ()) performHandshake connectionApi parserApi state auth conn host = do - infoResult <- readInitialInfo - case infoResult of - Left err -> - pure (Left (HandshakeProtocolError err)) - Right (info, rest) -> do - let cfg = config state - tlsRequested = isJust (tlsCert cfg) || Connect.tls_required (connectConfig cfg) - tlsRequired = fromMaybe False (I.tls_required info) - transportOption = - TransportOption - { transportHost = host - , transportTlsRequired = tlsRequired - , transportTlsRequested = tlsRequested - , transportTlsCert = tlsCert cfg - , transportInitialBytes = rest - } - transportResult <- configure connectionApi conn transportOption - case transportResult of + result <- timeout (max 1 (connectTimeoutMicros (config state))) handshake + pure (fromMaybe (Left HandshakeTimeout) result) + where + handshake = do + infoResult <- readInitialInfo + case infoResult of Left err -> - pure (Left (HandshakeTransportError err)) - Right () -> do - setServerInfo state info - updateLogContextFromInfo state info - let authContext = AuthContext { authNonce = I.nonce info } - connectPayload = - (connectConfig cfg) - { Connect.tls_required = tlsRequired || tlsRequested + pure (Left err) + Right (info, rest) -> do + let cfg = config state + tlsRequested = isJust (tlsConfig cfg) || Connect.tls_required (connectConfig cfg) + tlsRequired = fromMaybe False (I.tls_required info) + tlsAvailable = fromMaybe False (I.tls_available info) + useTls = tlsRequired || (tlsRequested && tlsAvailable) + transportOption = + TransportOption + { transportHost = host + , transportTlsRequired = tlsRequired + , transportTlsRequested = tlsRequested && tlsAvailable + , transportTlsConfig = + if useTls + then Just (fromMaybe defaultTLSConfig (tlsConfig cfg)) + else Nothing + , transportInitialBytes = rest } - case validateAuth auth of - Left err -> - pure (Left (HandshakeAuthError err)) - Right () -> - case buildAuthPatch auth authContext of + if tlsRequested && not (tlsRequired || tlsAvailable) + then pure (Left (HandshakeTLSError "server does not offer TLS")) + else do + transportResult <- configure connectionApi conn transportOption + case transportResult of Left err -> - pure (Left (HandshakeAuthError err)) - Right patch -> do - writeResult <- - writeDataLazy - (writer connectionApi) - conn - (transform (applyAuthPatch patch connectPayload)) - case writeResult of + pure (Left (HandshakeTLSError err)) + Right () -> do + setServerInfo state info + updateLogContextFromInfo state info + let authContext = AuthContext { authNonce = I.nonce info } + connectPayload = + (connectConfig cfg) + { Connect.tls_required = useTls + } + authResult <- buildAuthPatch auth authContext + case authResult of Left err -> - pure (Left (HandshakeTransportError err)) - Right () -> - pure (Right ()) - where + pure (Left (HandshakeAuthError err)) + Right patch -> do + writeResult <- + writeDataLazy + (writer connectionApi) + conn + ( transform (applyAuthPatch patch connectPayload) + <> transform Ping + ) + case writeResult of + Left err -> + pure (Left (HandshakeTransportError err)) + Right () -> + awaitPong mempty + + readInitialInfo :: IO (Either HandshakeError (I.Info, BS.ByteString)) readInitialInfo = go mempty where go acc = do result <- readData (reader connectionApi) conn 4096 case result of Left err -> - pure (Left err) + pure (Left (HandshakeTransportError err)) Right chunk -> if BS.null chunk - then pure (Left "read returned empty chunk before INFO") + then pure (Left (HandshakeTransportError "read returned empty chunk before INFO")) else do let bytes = acc <> chunk case parse parserApi bytes of @@ -106,11 +125,66 @@ DropPrefix n _ -> go (BS.drop n bytes) Reject reason -> - pure (Left reason) + pure (Left (HandshakeProtocolError reason)) Emit (ParsedInfo info) rest -> pure (Right (info, rest)) - Emit (ParsedErr err) _ -> - pure (Left ("server error before INFO: " ++ show err)) - Emit _ rest -> do - unless (BS.null rest) (pure ()) + Emit (ParsedErr err) _ + | isAuthenticationError err -> + pure (Left (HandshakeAuthError (AuthError (show err)))) + | otherwise -> + pure (Left (HandshakeProtocolError ("server error before INFO: " ++ show err))) + Emit _ rest -> go rest + + awaitPong :: BS.ByteString -> IO (Either HandshakeError ()) + awaitPong acc = do + result <- readData (reader connectionApi) conn 4096 + case result of + Left err -> + pure (Left (HandshakeTransportError err)) + Right chunk + | BS.null chunk -> + pure (Left (HandshakeTransportError "read returned empty chunk before PONG")) + | otherwise -> + consumePongFrames (acc <> chunk) + + consumePongFrames :: BS.ByteString -> IO (Either HandshakeError ()) + consumePongFrames bytes = + case parse parserApi bytes of + NeedMore -> + awaitPong bytes + DropPrefix _ reason -> + pure (Left (HandshakeProtocolError reason)) + Reject reason -> + pure (Left (HandshakeProtocolError reason)) + Emit (ParsedPong _) _ -> + pure (Right ()) + Emit (ParsedOk _) rest -> + continueWith rest + Emit (ParsedPing _) rest -> do + pongResult <- writeDataLazy (writer connectionApi) conn (transform Pong) + case pongResult of + Left err -> pure (Left (HandshakeTransportError err)) + Right () -> continueWith rest + Emit (ParsedInfo info) rest -> do + setServerInfo state info + updateLogContextFromInfo state info + continueWith rest + Emit (ParsedErr err) _ + | isAuthenticationError err -> + pure (Left (HandshakeAuthError (AuthError (show err)))) + | otherwise -> + pure (Left (HandshakeProtocolError (show err))) + Emit message _ -> + pure (Left (HandshakeProtocolError ("unexpected frame before PONG: " ++ show message))) + + continueWith rest + | BS.null rest = awaitPong mempty + | otherwise = consumePongFrames rest + + isAuthenticationError (Err.ErrAuthViolation _) = True + isAuthenticationError (Err.ErrAuthTimeout _) = True + isAuthenticationError (Err.ErrAuthExpired _) = True + isAuthenticationError (Err.ErrAuthRevoked _) = True + isAuthenticationError (Err.ErrAccountAuthExpired _) = True + isAuthenticationError _ = False diff --git a/internal/Policy/Publish.hs b/internal/Policy/Publish.hs --- a/internal/Policy/Publish.hs +++ b/internal/Policy/Publish.hs @@ -2,7 +2,7 @@ ( defaultPublishConfig ) where -import Publish.Config (PublishConfig) +import Publish.Config (PublishConfig (..)) defaultPublishConfig :: PublishConfig -defaultPublishConfig = (Nothing, Nothing, Nothing, Nothing) +defaultPublishConfig = PublishConfig Nothing Nothing diff --git a/internal/Policy/Publish/Config.hs b/internal/Policy/Publish/Config.hs --- a/internal/Policy/Publish/Config.hs +++ b/internal/Policy/Publish/Config.hs @@ -1,10 +1,11 @@ module Publish.Config - ( Payload - , Headers - , PublishConfig + ( Headers + , PublishConfig (..) ) where -import qualified Types.Msg as M -import Types.Msg (Headers, Payload, Subject) +import Types.Msg (Headers, Subject) -type PublishConfig = (Maybe Payload, Maybe (Maybe M.Msg -> IO ()), Maybe Headers, Maybe Subject) +data PublishConfig = PublishConfig + { publishHeaders :: Maybe Headers + , publishReplyTo :: Maybe Subject + } diff --git a/internal/Policy/Router/Nats.hs b/internal/Policy/Router/Nats.hs --- a/internal/Policy/Router/Nats.hs +++ b/internal/Policy/Router/Nats.hs @@ -6,7 +6,7 @@ import Control.Monad.IO.Class (liftIO) import Lib.Logger (LogLevel (..), MonadLogger (..)) import Parser.API - ( ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong) + ( ParsedMessage (ParsedErr, ParsedInfo, ParsedMessageTooLarge, ParsedMsg, ParsedOk, ParsedPing, ParsedPong) ) import Queue.API (QueueItem (QueueItem)) import State.Store @@ -17,8 +17,15 @@ , setServerInfo , updateLogContextFromInfo ) -import State.Types (ClientExitReason (ExitServerError)) -import Subscription.Store (SubscriptionStore, dispatchMessage) +import State.Types + ( ClientExitReason (ExitInboundMessageTooLarge, ExitServerError) + , serverErrorFromProtocol + ) +import Subscription.Store + ( DispatchResult (DispatchDropped, DispatchMissing, DispatchQueued) + , SubscriptionStore + , dispatchMessage + ) import qualified Types.Err as Err import qualified Types.Msg as Msg import Types.Pong (Pong (..)) @@ -33,12 +40,28 @@ case parsed of ParsedMsg msg -> do logMessage Debug ("routing MSG: " ++ show msg) - handled <- liftIO $ dispatchMessage store msg - if handled - then pure RouteContinue - else do + result <- liftIO $ dispatchMessage store msg + case result of + DispatchQueued -> + pure RouteContinue + DispatchDropped reportSlowConsumer -> do + if reportSlowConsumer + then + logMessage Error "slow consumer: global pending delivery limit reached" + else + logMessage Debug "dropping delivery while client remains a slow consumer" + pure RouteContinue + DispatchMissing -> do logMessage Error ("callback missing for SID: " ++ show (Msg.sid msg)) pure RouteContinue + ParsedMessageTooLarge actual maximumSize -> do + logMessage Error + ( "inbound message size " + ++ show actual + ++ " exceeds client limit " + ++ show maximumSize + ) + pure (RouteExit (ExitInboundMessageTooLarge actual maximumSize)) ParsedInfo info -> do logMessage Debug ("routing INFO: " ++ show info) liftIO $ setServerInfo state info @@ -60,7 +83,7 @@ if Err.isFatal err then do logMessage Error ("fatal server error: " ++ show err) - pure (RouteExit (ExitServerError err)) + pure (RouteExit (ExitServerError (serverErrorFromProtocol err))) else do logMessage Warn ("server error: " ++ show err) pure RouteContinue diff --git a/internal/Policy/State/Store.hs b/internal/Policy/State/Store.hs --- a/internal/Policy/State/Store.hs +++ b/internal/Policy/State/Store.hs @@ -28,6 +28,8 @@ , waitForServerInfo , markConnected , markConnectionReady + , waitForInitialConnection + , failInitialConnection , readConnectionGeneration , waitForConnectionGenerationAfter , readAttemptIndex @@ -35,7 +37,7 @@ ) where import Control.Concurrent.STM -import Control.Monad (unless) +import Control.Monad (unless, void) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BC import Data.Maybe (fromMaybe) @@ -55,18 +57,19 @@ import Types.Msg (SID, Subject) data ClientState = ClientState - { clientConfig :: ClientConfig - , clientQueue :: Queue - , clientPings :: TQueue (IO ()) - , clientConnectedOnce :: TVar Bool + { clientConfig :: ClientConfig + , clientQueue :: Queue + , clientPings :: TQueue (IO ()) + , clientConnectedOnce :: TVar Bool , clientConnectionGeneration :: TVar Int - , clientSidCounter :: TVar SIDCounter - , clientInboxNuid :: TVar Nuid - , clientServerInfo :: TVar (Maybe Info) - , clientAttemptIndex :: TVar Int - , clientStatus :: TVar ClientStatus - , clientConnection :: Conn - , clientLogContext :: TVar LogContext + , clientInitialConnection :: TMVar (Either ConnectError ()) + , clientSidCounter :: TVar SIDCounter + , clientInboxNuid :: TVar Nuid + , clientServerInfo :: TVar (Maybe Info) + , clientAttemptIndex :: TVar Int + , clientStatus :: TVar ClientStatus + , clientConnection :: Conn + , clientLogContext :: TVar LogContext } newClientState :: ClientConfig -> Queue -> Conn -> TVar LogContext -> IO ClientState @@ -74,6 +77,7 @@ pings <- newTQueueIO connectedOnce <- newTVarIO False connectionGeneration <- newTVarIO 0 + initialConnection <- newEmptyTMVarIO sidCounter <- newTVarIO initialSIDCounter inboxNuid <- newTVarIO =<< newNuidIO serverInfo <- newTVarIO Nothing @@ -85,6 +89,7 @@ , clientPings = pings , clientConnectedOnce = connectedOnce , clientConnectionGeneration = connectionGeneration + , clientInitialConnection = initialConnection , clientSidCounter = sidCounter , clientInboxNuid = inboxNuid , clientServerInfo = serverInfo @@ -225,8 +230,16 @@ markConnectionReady :: ClientState -> IO () markConnectionReady client = - atomically $ + atomically $ do modifyTVar' (clientConnectionGeneration client) (+ 1) + void (tryPutTMVar (clientInitialConnection client) (Right ())) + +waitForInitialConnection :: ClientState -> STM (Either ConnectError ()) +waitForInitialConnection = readTMVar . clientInitialConnection + +failInitialConnection :: ClientState -> ConnectError -> IO () +failInitialConnection client err = + atomically $ void (tryPutTMVar (clientInitialConnection client) (Left err)) readConnectionGeneration :: ClientState -> IO Int readConnectionGeneration = diff --git a/internal/Policy/State/Types.hs b/internal/Policy/State/Types.hs --- a/internal/Policy/State/Types.hs +++ b/internal/Policy/State/Types.hs @@ -4,34 +4,74 @@ ( TLSPublicKey , TLSPrivateKey , TLSCertData + , TLSConfig (..) + , defaultTLSConfig , ClientConfig (..) + , ConnectError (..) + , ConnectAttemptError (..) + , ConnectFailure (..) + , ServerError + , serverErrorReason + , serverErrorFromProtocol , ClientExitReason (..) , ClientStatus (..) ) where -import qualified Data.ByteString as BS +import Data.ByteString (ByteString) import Lib.Logger.Types (LoggerConfig) import Types.Connect (Connect) import qualified Types.Err as Err - -type TLSPublicKey = BS.ByteString -type TLSPrivateKey = BS.ByteString -type TLSCertData = (TLSPublicKey, TLSPrivateKey) +import Types.TLS data ClientConfig = ClientConfig - { connectionAttempts :: Int - , callbackConcurrency :: Int - , bufferLimit :: Int - , connectConfig :: Connect - , loggerConfig :: LoggerConfig - , tlsCert :: Maybe TLSCertData - , exitAction :: ClientExitReason -> IO () - , connectOptions :: [(String, Int)] + { connectionAttempts :: Int + , connectTimeoutMicros :: Int + , callbackConcurrency :: Int + , messageLimit :: Int + , connectConfig :: Connect + , loggerConfig :: LoggerConfig + , tlsConfig :: Maybe TLSConfig + , exitAction :: ClientExitReason -> IO () + , connectOptions :: [(String, Int)] } +-- | Why an initial connection could not be established. +data ConnectError = ConnectNoServers + | ConnectAttemptsExhausted [ConnectAttemptError] + deriving (Eq, Show) + +-- | Failure from one server in the configured pool. +data ConnectAttemptError = ConnectAttemptError + { attemptedEndpoint :: (String, Int) + , attemptFailure :: ConnectFailure + } + deriving (Eq, Show) + +-- | The stage at which a connection attempt failed. +data ConnectFailure = ConnectTransportFailure String + | ConnectTLSFailure String + | ConnectProtocolFailure String + | ConnectAuthenticationFailure String + | ConnectHandshakeTimeout + deriving (Eq, Show) + +-- | A server-reported protocol error. +-- +-- Its wire representation is intentionally opaque so that new server error +-- categories do not force changes to the public type. +newtype ServerError = ServerError ByteString + deriving (Eq, Show) + +serverErrorReason :: ServerError -> ByteString +serverErrorReason (ServerError reason) = reason + +serverErrorFromProtocol :: Err.Err -> ServerError +serverErrorFromProtocol = ServerError . Err.errReason + data ClientExitReason = ExitClosedByUser | ExitRetriesExhausted (Maybe String) - | ExitServerError Err.Err + | ExitServerError ServerError + | ExitInboundMessageTooLarge Int Int | ExitResetRequested deriving (Eq, Show) diff --git a/internal/Policy/Subscription/Store.hs b/internal/Policy/Subscription/Store.hs --- a/internal/Policy/Subscription/Store.hs +++ b/internal/Policy/Subscription/Store.hs @@ -1,5 +1,6 @@ module Subscription.Store ( SubscriptionStore + , DispatchResult (..) , newSubscriptionStore , register , unregister @@ -15,7 +16,7 @@ import Control.Concurrent (forkIO, threadDelay) import Control.Concurrent.STM import Control.Exception (SomeException, finally) -import Control.Monad (unless, void) +import Control.Monad (unless, void, when) import qualified Data.Heap as Heap import qualified Data.Map as Map import Data.Time.Clock @@ -25,38 +26,62 @@ import Types.Msg (SID) data SubscriptionState = SubscriptionState - { subscriptionCallbacks :: Map.Map SID (Maybe M.Msg -> IO ()) + { subscriptionCallbacks :: Map.Map SID SubscriptionCallback , subscriptionExpiryHeap :: Heap.MinHeap (UTCTime, SID) , subscriptionTrackedExpiries :: Map.Map SID UTCTime , subscriptionMeta :: Map.Map SID SubscriptionMeta } +data SubscriptionCallback = SubscriptionCallback + { deliverMessage :: Maybe M.Msg -> IO () + , dropMessage :: IO () + } + +data DispatchResult = DispatchMissing + | DispatchQueued + | DispatchDropped Bool + deriving (Eq, Show) + data SubscriptionStore = SubscriptionStore { storeState :: TVar SubscriptionState , storeCallbackQueue :: TQueue (IO ()) , storeCallbackPending :: TVar Int + , storeDeliveryPending :: TVar Int + , storeDeliveryBytes :: TVar Int + , storeSlowConsumer :: TVar Bool + , storePendingLimits :: PendingLimits + , storeSlowAction :: IO () } -newSubscriptionStore :: IO SubscriptionStore -newSubscriptionStore = +newSubscriptionStore :: PendingLimits -> IO () -> IO SubscriptionStore +newSubscriptionStore limits slowAction = SubscriptionStore <$> newTVarIO emptySubscriptionState <*> newTQueueIO <*> newTVarIO 0 + <*> newTVarIO 0 + <*> newTVarIO 0 + <*> newTVarIO False + <*> pure (normalizePendingLimits limits) + <*> pure slowAction emptySubscriptionState :: SubscriptionState emptySubscriptionState = SubscriptionState Map.empty Heap.empty Map.empty Map.empty -register :: SubscriptionStore -> SID -> SubscriptionMeta -> SubscribeConfig -> (Maybe M.Msg -> IO ()) -> IO () -register store sid meta cfg callback = do +register :: SubscriptionStore -> SID -> SubscriptionMeta -> SubscribeConfig -> (Maybe M.Msg -> IO ()) -> IO () -> IO () +register store sid meta cfg callback onDropped = do expiryAt <- case expiry cfg of Nothing -> pure Nothing Just ttl -> Just . addUTCTime ttl <$> getCurrentTime atomically . modifyTVar' (storeState store) $ \state -> let stateWithCallback = state - { subscriptionCallbacks = Map.insert sid callback (subscriptionCallbacks state) + { subscriptionCallbacks = + Map.insert + sid + (SubscriptionCallback callback onDropped) + (subscriptionCallbacks state) , subscriptionMeta = Map.insert sid meta (subscriptionMeta state) } in case expiryAt of @@ -70,14 +95,14 @@ atomically $ modifyTVar' (storeState store) (removeSubscriptionLocal sid) -dispatchMessage :: SubscriptionStore -> M.Msg -> IO Bool -dispatchMessage store msg = - atomically $ do +dispatchMessage :: SubscriptionStore -> M.Msg -> IO DispatchResult +dispatchMessage store msg = do + (result, onDropped) <- atomically $ do state <- readTVar (storeState store) let sid = M.sid msg case Map.lookup sid (subscriptionCallbacks state) of Nothing -> - pure False + pure (DispatchMissing, pure ()) Just callback -> do let shouldRemove = maybe False isReply (Map.lookup sid (subscriptionMeta state)) @@ -85,9 +110,21 @@ if shouldRemove then removeSubscriptionLocal sid state else state + messageBytes = M.messageContentSize (M.headers msg) (M.payload msg) + canQueue <- hasDeliveryCapacity store messageBytes writeTVar (storeState store) nextState - enqueueCallbackSTM store (callback (Just msg)) - pure True + if canQueue + then do + enqueueDeliverySTM store messageBytes (deliverMessage callback (Just msg)) + pure (DispatchQueued, pure ()) + else do + alreadySlow <- readTVar (storeSlowConsumer store) + unless alreadySlow $ do + writeTVar (storeSlowConsumer store) True + enqueueControlCallbackSTM store (storeSlowAction store) + pure (DispatchDropped (not alreadySlow), dropMessage callback) + onDropped + pure result active :: SubscriptionStore -> IO [(SID, SubscriptionMeta)] active store = @@ -123,20 +160,52 @@ hasTrackedExpiries store = not . Map.null . subscriptionTrackedExpiries <$> readTVarIO (storeState store) -enqueueCallbackSTM :: SubscriptionStore -> IO () -> STM () -enqueueCallbackSTM store action = do +enqueueControlCallbackSTM :: SubscriptionStore -> IO () -> STM () +enqueueControlCallbackSTM store action = do modifyTVar' (storeCallbackPending store) (+1) let wrapped = action `finally` atomically (modifyTVar' (storeCallbackPending store) (subtract 1)) writeTQueue (storeCallbackQueue store) wrapped +enqueueDeliverySTM :: SubscriptionStore -> Int -> IO () -> STM () +enqueueDeliverySTM store messageBytes action = do + modifyTVar' (storeCallbackPending store) (+1) + modifyTVar' (storeDeliveryPending store) (+1) + modifyTVar' (storeDeliveryBytes store) (+ messageBytes) + let wrapped = + action `finally` atomically (releaseDeliverySTM store messageBytes) + writeTQueue (storeCallbackQueue store) wrapped + +releaseDeliverySTM :: SubscriptionStore -> Int -> STM () +releaseDeliverySTM store messageBytes = do + modifyTVar' (storeCallbackPending store) (subtract 1) + modifyTVar' (storeDeliveryPending store) (subtract 1) + modifyTVar' (storeDeliveryBytes store) (subtract messageBytes) + pendingMessages <- readTVar (storeDeliveryPending store) + pendingBytes <- readTVar (storeDeliveryBytes store) + let limits = storePendingLimits store + belowLowWater = + pendingMessages <= pendingMessageLimit limits `div` 2 + && pendingBytes <= pendingByteLimit limits `div` 2 + when belowLowWater $ + writeTVar (storeSlowConsumer store) False + +hasDeliveryCapacity :: SubscriptionStore -> Int -> STM Bool +hasDeliveryCapacity store messageBytes = do + pendingMessages <- readTVar (storeDeliveryPending store) + pendingBytes <- readTVar (storeDeliveryBytes store) + let limits = storePendingLimits store + pure $ + pendingMessages < pendingMessageLimit limits + && messageBytes <= pendingByteLimit limits - pendingBytes + expireReadySubscriptions :: SubscriptionStore -> UTCTime -> IO Bool expireReadySubscriptions store now = atomically $ do state <- readTVar (storeState store) let (actions, nextState) = collectExpiredCallbacks now state writeTVar (storeState store) nextState - mapM_ (enqueueCallbackSTM store) actions + mapM_ (enqueueControlCallbackSTM store) actions pure (not (null actions)) trackSubscriptionExpiry :: SID -> UTCTime -> SubscriptionState -> SubscriptionState @@ -173,7 +242,11 @@ | otherwise -> let callback = Map.lookup sid (subscriptionCallbacks state) state' = removeSubscriptionLocal sid (dropHeapHead state) - actions' = maybe actions ((: actions) . ($ Nothing)) callback + actions' = + maybe + actions + ((: actions) . ($ Nothing) . deliverMessage) + callback in go actions' state' dropHeapHead :: SubscriptionState -> SubscriptionState @@ -181,3 +254,10 @@ case Heap.view (subscriptionExpiryHeap state) of Nothing -> state Just (_, heapTail) -> state { subscriptionExpiryHeap = heapTail } + +normalizePendingLimits :: PendingLimits -> PendingLimits +normalizePendingLimits limits = + limits + { pendingMessageLimit = max 1 (pendingMessageLimit limits) + , pendingByteLimit = max 1 (pendingByteLimit limits) + } diff --git a/internal/Policy/Subscription/Types.hs b/internal/Policy/Subscription/Types.hs --- a/internal/Policy/Subscription/Types.hs +++ b/internal/Policy/Subscription/Types.hs @@ -1,6 +1,8 @@ module Subscription.Types ( SubscribeConfig (..) , SubscriptionMeta (..) + , PendingLimits (..) + , defaultPendingLimits ) where import Data.Time.Clock (NominalDiffTime) @@ -17,3 +19,15 @@ , isReply :: Bool } deriving (Eq, Show) + +data PendingLimits = PendingLimits + { pendingMessageLimit :: Int + , pendingByteLimit :: Int + } + deriving (Eq, Show) + +defaultPendingLimits :: PendingLimits +defaultPendingLimits = PendingLimits + { pendingMessageLimit = 65536 + , pendingByteLimit = 64 * 1024 * 1024 + } diff --git a/internal/Types/Err.hs b/internal/Types/Err.hs --- a/internal/Types/Err.hs +++ b/internal/Types/Err.hs @@ -1,6 +1,7 @@ module Types.Err ( Reason , Err (..) + , errReason , isFatal ) where @@ -12,6 +13,9 @@ | ErrRoutePortConn Reason | ErrAuthViolation Reason | ErrAuthTimeout Reason + | ErrAuthExpired Reason + | ErrAuthRevoked Reason + | ErrAccountAuthExpired Reason | ErrInvalidProtocol Reason | ErrMaxControlLineEx Reason | ErrErr Reason @@ -28,3 +32,24 @@ isFatal (ErrInvalidSubject _) = False isFatal (ErrPermViolation _) = False isFatal _ = True + +errReason :: Err -> Reason +errReason err = + case err of + ErrUnknownOp reason -> reason + ErrRoutePortConn reason -> reason + ErrAuthViolation reason -> reason + ErrAuthTimeout reason -> reason + ErrAuthExpired reason -> reason + ErrAuthRevoked reason -> reason + ErrAccountAuthExpired reason -> reason + ErrInvalidProtocol reason -> reason + ErrMaxControlLineEx reason -> reason + ErrErr reason -> reason + ErrTlsRequired reason -> reason + ErrStaleConn reason -> reason + ErrMaxConnsEx reason -> reason + ErrSlowConsumer reason -> reason + ErrMaxPayload reason -> reason + ErrInvalidSubject reason -> reason + ErrPermViolation reason -> reason diff --git a/internal/Types/Info.hs b/internal/Types/Info.hs --- a/internal/Types/Info.hs +++ b/internal/Types/Info.hs @@ -21,6 +21,7 @@ , nonce :: Maybe BS.ByteString , auth_required :: Maybe Bool , tls_required :: Maybe Bool + , tls_available :: Maybe Bool , connect_urls :: Maybe [BS.ByteString] , ldm :: Maybe Bool , headers :: Maybe Bool @@ -53,6 +54,7 @@ , infoJSON_nonce :: Maybe Utf8ByteString , infoJSON_auth_required :: Maybe Bool , infoJSON_tls_required :: Maybe Bool + , infoJSON_tls_available :: Maybe Bool , infoJSON_connect_urls :: Maybe [Utf8ByteString] , infoJSON_ldm :: Maybe Bool , infoJSON_headers :: Maybe Bool @@ -73,6 +75,7 @@ , nonce = unUtf8ByteString <$> infoJSON_nonce infoJson , auth_required = infoJSON_auth_required infoJson , tls_required = infoJSON_tls_required infoJson + , tls_available = infoJSON_tls_available infoJson , connect_urls = fmap (map unUtf8ByteString) (infoJSON_connect_urls infoJson) , ldm = infoJSON_ldm infoJson , headers = infoJSON_headers infoJson diff --git a/internal/Types/Msg.hs b/internal/Types/Msg.hs --- a/internal/Types/Msg.hs +++ b/internal/Types/Msg.hs @@ -4,9 +4,11 @@ , Payload , Headers , Msg (..) + , messageContentSize ) where import Data.ByteString (ByteString) +import qualified Data.ByteString as BS type Subject = ByteString type SID = ByteString @@ -21,3 +23,18 @@ , headers :: Maybe Headers } deriving (Eq, Show) + +-- | Bytes retained for a NATS message body: encoded headers plus payload. +messageContentSize :: Maybe Headers -> Maybe Payload -> Int +messageContentSize maybeHeaders maybePayload = + fromInteger (min (toInteger (maxBound :: Int)) totalBytes) + where + totalBytes = headerBytes + maybe 0 (toInteger . BS.length) maybePayload + headerBytes :: Integer + headerBytes = + case maybeHeaders of + Nothing -> 0 + Just headerPairs -> + 12 + sum (map headerPairSize headerPairs) + headerPairSize (key, value) = + toInteger (BS.length key) + toInteger (BS.length value) + 3 diff --git a/internal/Types/Pub.hs b/internal/Types/Pub.hs --- a/internal/Types/Pub.hs +++ b/internal/Types/Pub.hs @@ -8,7 +8,12 @@ import Data.Maybe import Prelude hiding (concat, length, null) import Transformers.Transformers (Transformer (..)) -import Types.Msg (Headers, Payload, Subject) +import Types.Msg + ( Headers + , Payload + , Subject + , messageContentSize + ) import Validators.Validators data Pub = Pub @@ -73,11 +78,7 @@ messageSize :: Pub -> Int messageSize pubMsg = - case headers pubMsg of - Nothing -> - length (payloadChunk (payload pubMsg)) - Just headerList -> - length (headerString headerList) + length (payloadChunk (payload pubMsg)) + messageContentSize (headers pubMsg) (payload pubMsg) headerString :: Headers -> ByteString headerString hs = diff --git a/internal/Types/TLS.hs b/internal/Types/TLS.hs new file mode 100644 --- /dev/null +++ b/internal/Types/TLS.hs @@ -0,0 +1,47 @@ +module Types.TLS + ( TLSPublicKey + , TLSPrivateKey + , TLSCertData + , TLSConfig (..) + , defaultTLSConfig + ) where + +import qualified Data.ByteString as BS +import Data.Maybe (isJust) + +type TLSPublicKey = BS.ByteString +type TLSPrivateKey = BS.ByteString +type TLSCertData = (TLSPublicKey, TLSPrivateKey) + +-- | TLS trust and client-identity configuration. +data TLSConfig = TLSConfig + { tlsClientCertificate :: Maybe TLSCertData + , tlsRootCertificates :: [BS.ByteString] + , tlsServerName :: Maybe String + , tlsInsecure :: Bool + } + deriving (Eq) + +instance Show TLSConfig where + show config = + "TLSConfig {tlsClientCertificate = " + ++ configured (isJust (tlsClientCertificate config)) + ++ ", tlsRootCertificates = " + ++ show (length (tlsRootCertificates config)) + ++ " configured, tlsServerName = " + ++ show (tlsServerName config) + ++ ", tlsInsecure = " + ++ show (tlsInsecure config) + ++ "}" + where + configured True = "" + configured False = "Nothing" + +defaultTLSConfig :: TLSConfig +defaultTLSConfig = + TLSConfig + { tlsClientCertificate = Nothing + , tlsRootCertificates = [] + , tlsServerName = Nothing + , tlsInsecure = False + } diff --git a/jetstream-client/JetStream/API.hs b/jetstream-client/JetStream/API.hs deleted file mode 100644 --- a/jetstream-client/JetStream/API.hs +++ /dev/null @@ -1,53 +0,0 @@ --- | Capability record for the JetStream client surface. -module JetStream.API - ( - JetStream (..) - , AccountAPIStats (..) - , AccountInfo (..) - , AccountLimits (..) - , AccountTier (..) - , JetStreamApiError (..) - , JetStreamError (..) - , StreamName - , ConsumerName - , Subject - , Payload - , module JetStream.Consumer.API - , module JetStream.Message.API - , module JetStream.Publish.API - , module JetStream.Stream.API - ) where - -import JetStream.Consumer.API -import JetStream.Error - ( JetStreamApiError (..) - , JetStreamError (..) - ) -import JetStream.Message.API -import JetStream.Publish.API -import JetStream.Stream.API -import JetStream.Types - ( AccountAPIStats (..) - , AccountInfo (..) - , AccountLimits (..) - , AccountTier (..) - , ConsumerName - , Payload - , StreamName - , Subject - ) - --- | JetStream capabilities for streams, consumers, publishing, and pull --- message workflows. -data JetStream = JetStream - { streams :: StreamAPI - -- ^ Manage streams and inspect stream state. - , consumers :: ConsumerAPI - -- ^ Manage stream consumers. - , publisher :: PublishAPI - -- ^ Publish messages to JetStream subjects. - , messages :: MessageAPI - -- ^ Fetch and acknowledge pull-consumer messages. - , accountInfo :: IO (Either JetStreamError AccountInfo) - -- ^ Fetch account-level JetStream usage and limits. - } diff --git a/jetstream-client/JetStream/Client.hs b/jetstream-client/JetStream/Client.hs deleted file mode 100644 --- a/jetstream-client/JetStream/Client.hs +++ /dev/null @@ -1,37 +0,0 @@ --- | High-level client implementation for JetStream. -module JetStream.Client - ( newJetStream - , JetStream (..) - , JetStreamOption - , withDomain - , withRequestTimeoutMicros - ) where - -import qualified API as Nats -import JetStream.API (JetStream (..)) -import qualified JetStream.Consumer as Consumer -import qualified JetStream.Message as Message -import JetStream.Options - ( JetStreamOption - , newJetStreamContext - , withDomain - , withRequestTimeoutMicros - ) -import qualified JetStream.Protocol.Request as Request -import qualified JetStream.Protocol.Subject as Subject -import qualified JetStream.Publish as Publish -import qualified JetStream.Stream as Stream - --- | Build a JetStream capability record from an existing NATS client. -newJetStream :: Nats.Client -> [JetStreamOption] -> JetStream -newJetStream client options = - let ctx = newJetStreamContext client options - consumerAPI = Consumer.consumerAPI ctx - in JetStream - { streams = Stream.streamAPI ctx - , consumers = consumerAPI - , publisher = Publish.publishAPI ctx - , messages = Message.messageAPI ctx consumerAPI - , accountInfo = - Request.requestJSON ctx (Subject.accountInfoSubject ctx) Nothing - } diff --git a/jetstream/JetStream/Consumer.hs b/jetstream/JetStream/Consumer.hs --- a/jetstream/JetStream/Consumer.hs +++ b/jetstream/JetStream/Consumer.hs @@ -7,13 +7,10 @@ import Data.Aeson (Value, object, toJSON, (.=)) import Data.Maybe (catMaybes) -import JetStream.Consumer.API (ConsumerAPI (..)) +import JetStream.Consumer.API import JetStream.Consumer.Types - ( ConsumerAction (..) - , ConsumerConfigOption + ( ConsumerAPI (..) , ConsumerConfigRequest - , ConsumerKind - , ConsumerTarget (..) , applyConsumerKind , applyConsumerTarget , consumerActionValue @@ -38,7 +35,7 @@ consumerAPI :: JetStreamContext -> ConsumerAPI consumerAPI context = ConsumerAPI - { putConsumer = \stream action target kind options -> + { putConsumer = \stream action target kind options requestOptions -> case consumerRequest context stream action target kind options of Left err -> pure (Left err) @@ -46,34 +43,42 @@ Request.requestJSON context subject (Just (createConsumerRequest stream actionValue config)) - , consumerInfo = \stream consumer -> + requestOptions + , consumerInfo = \stream consumer requestOptions -> Request.requestJSON context (Subject.consumerInfoSubject context stream consumer) Nothing - , pauseConsumer = \stream consumer pauseUntil -> + requestOptions + , pauseConsumer = \stream consumer pauseUntil requestOptions -> Request.requestJSON context (Subject.consumerPauseSubject context stream consumer) (Just (toJSON (consumerPauseRequest (Just pauseUntil)))) - , resumeConsumer = \stream consumer -> + requestOptions + , resumeConsumer = \stream consumer requestOptions -> Request.requestJSON context (Subject.consumerPauseSubject context stream consumer) (Just (toJSON (consumerPauseRequest Nothing))) - , resetConsumer = \stream consumer options -> + requestOptions + , resetConsumer = \stream consumer options requestOptions -> Request.requestJSON context (Subject.consumerResetSubject context stream consumer) (Just (toJSON (consumerResetRequest options))) - , deleteConsumer = \stream consumer -> + requestOptions + , deleteConsumer = \stream consumer requestOptions -> Request.requestJSON context (Subject.consumerDeleteSubject context stream consumer) Nothing - , listConsumers = \stream options -> + requestOptions + , listConsumers = \stream options requestOptions -> Request.requestJSON context (Subject.consumerListSubject context stream) (Just (toJSON (consumerListRequest options))) - , consumerNames = \stream options -> + requestOptions + , consumerNames = \stream options requestOptions -> Request.requestJSON context (Subject.consumerNamesSubject context stream) (Just (toJSON (consumerNamesRequest options))) + requestOptions } consumerRequest diff --git a/jetstream/JetStream/Consumer/API.hs b/jetstream/JetStream/Consumer/API.hs --- a/jetstream/JetStream/Consumer/API.hs +++ b/jetstream/JetStream/Consumer/API.hs @@ -1,21 +1,81 @@ module JetStream.Consumer.API - ( ConsumerAPI (..) + ( ConsumerAPI + , putConsumer + , consumerInfo + , pauseConsumer + , resumeConsumer + , resetConsumer + , deleteConsumer + , listConsumers + , consumerNames + , Consumer + , consumerStreamName + , consumerName + , lookupConsumer + , putConsumerHandle + , getConsumerInfo , AckPolicy (..) , ConsumerAction (..) - , ConsumerConfig (..) + , ConsumerConfig + , consumerConfigDurableName + , consumerConfigName + , consumerConfigDescription + , consumerConfigDeliverSubject + , consumerConfigDeliverGroup + , consumerConfigDeliverPolicy + , consumerConfigAckPolicy + , consumerConfigReplayPolicy + , consumerConfigFilterSubject + , consumerConfigFilterSubjects + , consumerConfigAckWait + , consumerConfigMaxDeliver + , consumerConfigMaxWaiting + , consumerConfigMaxAckPending + , consumerConfigInactiveThreshold + , consumerConfigIdleHeartbeat + , consumerConfigHeadersOnly + , consumerConfigReplicas + , consumerConfigMemoryStorage , ConsumerConfigOption , ConsumerFilter (..) , ConsumerKind (..) - , ConsumerInfo (..) + , ConsumerInfo + , consumerInfoStreamName + , consumerInfoName + , consumerInfoCreated + , consumerInfoConfig + , consumerInfoDelivered + , consumerInfoAckFloor + , consumerInfoNumAckPending + , consumerInfoNumRedelivered + , consumerInfoNumWaiting + , consumerInfoNumPending , ConsumerListOption - , ConsumerListResponse (..) - , ConsumerNamesResponse (..) - , ConsumerPauseResponse (..) + , ConsumerListResponse + , consumerListTotal + , consumerListOffset + , consumerListLimit + , consumerListConsumers + , ConsumerNamesResponse + , consumerNamesTotal + , consumerNamesOffset + , consumerNamesLimit + , consumerNamesConsumers + , ConsumerPauseResponse + , consumerPausePaused + , consumerPauseUntilTime + , consumerPauseRemaining , ConsumerResetOption - , ConsumerResetResponse (..) - , ConsumerSequenceInfo (..) + , ConsumerResetResponse + , consumerResetInfo + , consumerResetResponseSequence + , ConsumerSequenceInfo + , consumerSequenceConsumer + , consumerSequenceStream + , consumerSequenceLast , ConsumerTarget (..) - , DeleteConsumerResponse (..) + , DeleteConsumerResponse + , deleteConsumerSuccess , DeliverPolicy (..) , ReplayPolicy (..) , withConsumerAckPolicy @@ -37,27 +97,43 @@ , withConsumerResetSequence ) where -import Data.Time.Clock (UTCTime) import JetStream.Consumer.Types import JetStream.Error (JetStreamError) -import JetStream.Types (ConsumerName, StreamName) +import JetStream.Types + ( ConsumerName + , JetStreamRequestOption + , StreamName + ) --- | Consumer management operations. -data ConsumerAPI = ConsumerAPI - { putConsumer :: StreamName -> ConsumerAction -> ConsumerTarget -> ConsumerKind -> [ConsumerConfigOption] -> IO (Either JetStreamError ConsumerInfo) - -- ^ Create, update, or create-or-update a stream consumer. - , consumerInfo :: StreamName -> ConsumerName -> IO (Either JetStreamError ConsumerInfo) - -- ^ Fetch consumer configuration and state. - , pauseConsumer :: StreamName -> ConsumerName -> UTCTime -> IO (Either JetStreamError ConsumerPauseResponse) - -- ^ Pause a consumer until the given time. - , resumeConsumer :: StreamName -> ConsumerName -> IO (Either JetStreamError ConsumerPauseResponse) - -- ^ Resume a paused consumer. - , resetConsumer :: StreamName -> ConsumerName -> [ConsumerResetOption] -> IO (Either JetStreamError ConsumerResetResponse) - -- ^ Reset a consumer, optionally to a stream sequence. - , deleteConsumer :: StreamName -> ConsumerName -> IO (Either JetStreamError DeleteConsumerResponse) - -- ^ Delete a consumer from a stream. - , listConsumers :: StreamName -> [ConsumerListOption] -> IO (Either JetStreamError ConsumerListResponse) - -- ^ List consumers for a stream. - , consumerNames :: StreamName -> [ConsumerListOption] -> IO (Either JetStreamError ConsumerNamesResponse) - -- ^ List consumer names for a stream. - } +lookupConsumer + :: ConsumerAPI + -> StreamName + -> ConsumerName + -> [JetStreamRequestOption] + -> IO (Either JetStreamError Consumer) +lookupConsumer api stream name requestOptions = + fmap (fmap (const (Consumer stream name))) + (consumerInfo api stream name requestOptions) + +putConsumerHandle + :: ConsumerAPI + -> StreamName + -> ConsumerAction + -> ConsumerTarget + -> ConsumerKind + -> [ConsumerConfigOption] + -> [JetStreamRequestOption] + -> IO (Either JetStreamError Consumer) +putConsumerHandle api stream action target kind configOptions requestOptions = + fmap (fmap toHandle) + (putConsumer api stream action target kind configOptions requestOptions) + where + toHandle result = Consumer (consumerInfoStreamName result) (consumerInfoName result) + +getConsumerInfo + :: ConsumerAPI + -> Consumer + -> [JetStreamRequestOption] + -> IO (Either JetStreamError ConsumerInfo) +getConsumerInfo api consumer = + consumerInfo api (consumerStreamName consumer) (consumerName consumer) diff --git a/jetstream/JetStream/Consumer/Types.hs b/jetstream/JetStream/Consumer/Types.hs --- a/jetstream/JetStream/Consumer/Types.hs +++ b/jetstream/JetStream/Consumer/Types.hs @@ -1,7 +1,9 @@ {-# LANGUAGE OverloadedStrings #-} module JetStream.Consumer.Types - ( AckPolicy (..) + ( ConsumerAPI (..) + , Consumer (..) + , AckPolicy (..) , ConsumerConfig (..) , ConsumerAction (..) , ConsumerConfigOption @@ -56,11 +58,14 @@ import Data.Maybe (catMaybes) import qualified Data.Text as T import Data.Time.Clock (NominalDiffTime, UTCTime) +import Data.Word (Word64) +import JetStream.Error (JetStreamError) import JetStream.Types ( AckPolicy (..) , CallOption , ConsumerName , DeliverPolicy (..) + , JetStreamRequestOption , ReplayPolicy (..) , StreamName , Subject @@ -70,6 +75,26 @@ , parseByteString ) +-- | Consumer management operations. The public API hides the constructor so +-- new operations can be added without breaking downstream code. +data ConsumerAPI = ConsumerAPI + { putConsumer :: StreamName -> ConsumerAction -> ConsumerTarget -> ConsumerKind -> [ConsumerConfigOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerInfo) + , consumerInfo :: StreamName -> ConsumerName -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerInfo) + , pauseConsumer :: StreamName -> ConsumerName -> UTCTime -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerPauseResponse) + , resumeConsumer :: StreamName -> ConsumerName -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerPauseResponse) + , resetConsumer :: StreamName -> ConsumerName -> [ConsumerResetOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerResetResponse) + , deleteConsumer :: StreamName -> ConsumerName -> [JetStreamRequestOption] -> IO (Either JetStreamError DeleteConsumerResponse) + , listConsumers :: StreamName -> [ConsumerListOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerListResponse) + , consumerNames :: StreamName -> [ConsumerListOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerNamesResponse) + } + +-- | Stable identity of a consumer within a stream. +data Consumer = Consumer + { consumerStreamName :: StreamName + , consumerName :: ConsumerName + } + deriving (Eq, Ord, Show) + data ConsumerConfigRequest = ConsumerConfigRequest { consumerConfigRequestDurableName :: Maybe ConsumerName , consumerConfigRequestName :: Maybe ConsumerName @@ -278,8 +303,8 @@ deriving (Eq, Show) data ConsumerSequenceInfo = ConsumerSequenceInfo - { consumerSequenceConsumer :: Integer - , consumerSequenceStream :: Integer + { consumerSequenceConsumer :: Word64 + , consumerSequenceStream :: Word64 , consumerSequenceLast :: Maybe UTCTime } deriving (Eq, Show) @@ -301,7 +326,7 @@ } deriving (Eq, Show) -newtype ConsumerResetRequest = ConsumerResetRequest { consumerResetRequestSequence :: Maybe Integer } +newtype ConsumerResetRequest = ConsumerResetRequest { consumerResetRequestSequence :: Maybe Word64 } deriving (Eq, Show) type ConsumerResetOption = CallOption ConsumerResetRequest @@ -313,13 +338,13 @@ { consumerResetRequestSequence = Nothing } -withConsumerResetSequence :: Integer -> ConsumerResetOption +withConsumerResetSequence :: Word64 -> ConsumerResetOption withConsumerResetSequence sequenceNumber request = request { consumerResetRequestSequence = Just sequenceNumber } data ConsumerResetResponse = ConsumerResetResponse { consumerResetInfo :: ConsumerInfo - , consumerResetResponseSequence :: Integer + , consumerResetResponseSequence :: Word64 } deriving (Eq, Show) @@ -592,7 +617,7 @@ "last_per_subject" -> pure DeliverLastPerSubject value -> - fail ("unknown deliver policy: " ++ T.unpack value) + pure (DeliverPolicyUnknown value) parseByteStringField :: Object -> Key -> Parser BS.ByteString parseByteStringField obj key = diff --git a/jetstream/JetStream/Error.hs b/jetstream/JetStream/Error.hs --- a/jetstream/JetStream/Error.hs +++ b/jetstream/JetStream/Error.hs @@ -5,6 +5,7 @@ , JetStreamError (..) ) where +import qualified Client.API as Nats import Data.Aeson import qualified Data.ByteString as BS import qualified Data.Text.Encoding as E @@ -21,6 +22,7 @@ | JetStreamDecodeError String | JetStreamTimeout | JetStreamNoReply + | JetStreamNatsError Nats.NatsError deriving (Eq, Show) instance FromJSON JetStreamApiError where diff --git a/jetstream/JetStream/Management.hs b/jetstream/JetStream/Management.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Management.hs @@ -0,0 +1,14 @@ +-- | Account-level JetStream management implementation. +module JetStream.Management (managementAPI) where + +import JetStream.Management.API (ManagementAPI (..)) +import JetStream.Options (JetStreamContext) +import qualified JetStream.Protocol.Request as Request +import qualified JetStream.Protocol.Subject as Subject + +managementAPI :: JetStreamContext -> ManagementAPI +managementAPI context = + ManagementAPI + { accountInfo = + Request.requestJSON context (Subject.accountInfoSubject context) Nothing + } diff --git a/jetstream/JetStream/Management/API.hs b/jetstream/JetStream/Management/API.hs new file mode 100644 --- /dev/null +++ b/jetstream/JetStream/Management/API.hs @@ -0,0 +1,10 @@ +-- | Account-level JetStream management operations. +module JetStream.Management.API + ( ManagementAPI (..) + ) where + +import JetStream.Error (JetStreamError) +import JetStream.Types (AccountInfo, JetStreamRequestOption) + +-- | Account and tier management capabilities. +newtype ManagementAPI = ManagementAPI { accountInfo :: [JetStreamRequestOption] -> IO (Either JetStreamError AccountInfo) } diff --git a/jetstream/JetStream/Message.hs b/jetstream/JetStream/Message.hs --- a/jetstream/JetStream/Message.hs +++ b/jetstream/JetStream/Message.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE LambdaCase #-} + module JetStream.Message ( messageAPI , fetchMessages @@ -7,7 +9,7 @@ , termMessage ) where -import qualified API as Nats +import qualified Client.API as Nats import Control.Concurrent.STM import Control.Exception (bracket) import Control.Monad (unless, void) @@ -15,7 +17,13 @@ import qualified Data.ByteString.Char8 as BC import Data.Char (isAlphaNum) import Data.Maybe (catMaybes, fromMaybe) -import JetStream.Consumer.API (ConsumerAPI (..)) +import Data.Word (Word64) +import JetStream.Consumer.API + ( ConsumerAPI + , consumerInfo + , deleteConsumer + , putConsumer + ) import JetStream.Consumer.Types ( ConsumerAction (ConsumerCreateOrUpdate) , ConsumerConfigOption @@ -34,16 +42,19 @@ , withConsumerReplicas ) import JetStream.Error - ( JetStreamError (JetStreamDecodeError, JetStreamNoReply, JetStreamTimeout) + ( JetStreamError (JetStreamDecodeError, JetStreamNatsError, JetStreamNoReply, JetStreamTimeout) ) -import JetStream.Message.API (MessageAPI (..)) import JetStream.Message.Types -import JetStream.Options (JetStreamContext (..)) +import JetStream.Options + ( JetStreamContext (..) + , requestTimeoutMicros + ) import qualified JetStream.Protocol.Subject as Subject import JetStream.Types ( AckPolicy (AckNone) , ConsumerName , DeliverPolicy (DeliverByStartSequence) + , JetStreamRequestOption , StreamName , Subject ) @@ -53,16 +64,16 @@ messageAPI :: JetStreamContext -> ConsumerAPI -> MessageAPI messageAPI context consumerAPI = MessageAPI - { fetch = \stream consumer options -> - fetchMessages context stream consumer (pullRequest options) - , consumePush = \deliverSubject options handler -> + { fetch = \stream consumer options requestOptions -> + fetchMessages context stream consumer (pullRequest options) requestOptions + , consumePush = \deliverSubject options _ handler -> consumePushMessages context deliverSubject (pushConsumeConfig options) handler - , createOrderedConsumer = \stream options -> - createOrderedConsumerHandle context consumerAPI stream (orderedConsumerConfig options) - , ack = ackMessage (contextClient context) - , nak = nakMessage (contextClient context) - , inProgress = inProgressMessage (contextClient context) - , term = termMessage (contextClient context) + , createOrderedConsumer = \stream options requestOptions -> + createOrderedConsumerHandle context consumerAPI stream (orderedConsumerConfig options) requestOptions + , ack = \message _ -> ackMessage (contextClient context) message + , nak = \message _ -> nakMessage (contextClient context) message + , inProgress = \message _ -> inProgressMessage (contextClient context) message + , term = \message _ -> termMessage (contextClient context) message } consumePushMessages @@ -70,17 +81,16 @@ -> Subject -> PushConsumeConfig -> (Message -> IO ()) - -> IO PushSubscription + -> IO (Either JetStreamError PushSubscription) consumePushMessages context deliverSubject config handler = do - sid <- Nats.subscribe natsClient deliverSubject subscribeOptions $ \msgView -> - case fmap fromMsgView msgView of - Nothing -> - pure () - Just message -> - unless (isStatusMessage message) (handler message) - pure PushSubscription - { stopPushSubscription = Nats.unsubscribe natsClient sid - } + subscription <- Nats.subscribe natsClient deliverSubject subscribeOptions $ \msgView -> + let message = fromMsgView msgView + in unless (isStatusMessage message) (handler message) + pure $ case subscription of + Left err -> Left (JetStreamNatsError err) + Right handle -> Right PushSubscription + { stopPushSubscription = mapNatsResult <$> Nats.unsubscribe natsClient handle [] + } where natsClient = contextClient context subscribeOptions = @@ -93,7 +103,7 @@ , orderedStateConfig :: OrderedConsumerConfig , orderedStateNamePrefix :: ConsumerName , orderedStateSerial :: TVar Int - , orderedStateNextSequence :: TVar (Maybe Integer) + , orderedStateNextSequence :: TVar (Maybe Word64) , orderedStateCurrentName :: TVar (Maybe ConsumerName) , orderedStateStopped :: TVar Bool } @@ -103,15 +113,16 @@ -> ConsumerAPI -> StreamName -> OrderedConsumerConfig + -> [JetStreamRequestOption] -> IO (Either JetStreamError OrderedConsumer) -createOrderedConsumerHandle context consumerAPI stream config = do +createOrderedConsumerHandle context consumerAPI stream config requestOptions = do namePrefix <- orderedNamePrefix context config state <- OrderedState context consumerAPI stream config namePrefix <$> newTVarIO 0 <*> newTVarIO Nothing <*> newTVarIO Nothing <*> newTVarIO False - created <- resetOrderedConsumer state + created <- resetOrderedConsumer state requestOptions pure $ do void created Right OrderedConsumer @@ -120,8 +131,8 @@ , stopOrderedConsumer = stopOrdered state } -orderedInfo :: OrderedState -> IO (Either JetStreamError ConsumerInfo) -orderedInfo state = do +orderedInfo :: OrderedState -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerInfo) +orderedInfo state requestOptions = do currentName <- readTVarIO (orderedStateCurrentName state) case currentName of Nothing -> @@ -131,10 +142,11 @@ (orderedStateConsumers state) (orderedStateStream state) consumerName + requestOptions -fetchOrderedMessages :: OrderedState -> [FetchOption] -> IO (Either JetStreamError PullResponse) -fetchOrderedMessages state options = do - resetResult <- resetOrderedConsumer state +fetchOrderedMessages :: OrderedState -> [FetchOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError PullResponse) +fetchOrderedMessages state options requestOptions = do + resetResult <- resetOrderedConsumer state requestOptions case resetResult of Left err -> pure (Left err) @@ -144,6 +156,7 @@ (orderedStateStream state) (consumerInfoName info) (pullRequest options) + requestOptions case responseResult of Left err -> pure (Left err) @@ -158,8 +171,8 @@ writeTVar (orderedStateNextSequence state) (Just nextSequence) pure (Right response) -resetOrderedConsumer :: OrderedState -> IO (Either JetStreamError ConsumerInfo) -resetOrderedConsumer state = do +resetOrderedConsumer :: OrderedState -> [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerInfo) +resetOrderedConsumer state requestOptions = do stopped <- readTVarIO (orderedStateStopped state) if stopped then pure (Left JetStreamNoReply) @@ -174,6 +187,7 @@ (orderedStateConsumers state) (orderedStateStream state) consumerName + requestOptions (serial, nextSequence) <- atomically $ do serial <- succ <$> readTVar (orderedStateSerial state) writeTVar (orderedStateSerial state) serial @@ -188,6 +202,7 @@ (NamedConsumer consumerName) PullConsumer options + requestOptions case result of Left err -> pure (Left err) @@ -196,24 +211,25 @@ writeTVar (orderedStateCurrentName state) (Just consumerName) pure (Right info) -stopOrdered :: OrderedState -> IO () -stopOrdered state = do +stopOrdered :: OrderedState -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) +stopOrdered state requestOptions = do atomically $ writeTVar (orderedStateStopped state) True currentName <- readTVarIO (orderedStateCurrentName state) case currentName of Nothing -> - pure () + pure (Right ()) Just consumerName -> - void $ + void <$> deleteConsumer (orderedStateConsumers state) (orderedStateStream state) consumerName + requestOptions orderedConsumerOptions :: OrderedConsumerConfig - -> Maybe Integer + -> Maybe Word64 -> [ConsumerConfigOption] orderedConsumerOptions config nextSequence = catMaybes @@ -234,11 +250,11 @@ Nothing -> orderedConsumerDeliverPolicy config Just sequenceNumber -> - DeliverByStartSequence (fromInteger sequenceNumber) + DeliverByStartSequence sequenceNumber inactiveThreshold = fromMaybe 300 (orderedConsumerInactiveThreshold config) -orderedResponseNextSequence :: PullResponse -> Either JetStreamError (Maybe Integer) +orderedResponseNextSequence :: PullResponse -> Either JetStreamError (Maybe Word64) orderedResponseNextSequence response = case reverse (pullResponseMessages response) of [] -> @@ -272,24 +288,35 @@ | otherwise = '_' -- | Fetch messages for a pull consumer. -fetchMessages :: JetStreamContext -> StreamName -> ConsumerName -> PullRequest -> IO (Either JetStreamError PullResponse) -fetchMessages context stream consumer request +fetchMessages :: JetStreamContext -> StreamName -> ConsumerName -> PullRequest -> [JetStreamRequestOption] -> IO (Either JetStreamError PullResponse) +fetchMessages context stream consumer request requestOptions | pullRequestBatch request <= 0 = pure (Right (PullResponse [] Nothing)) | otherwise = do responseQueue <- newTQueueIO inbox <- Nats.newInbox natsClient bracket (Nats.subscribe natsClient inbox [] (atomically . writeTQueue responseQueue)) - (Nats.unsubscribe natsClient) - (\_ -> do - Nats.publish natsClient requestSubject - [ Nats.withPayload (pullRequestPayload (pullRequestBatch request) request) - , Nats.withReplyTo inbox - ] - collectResponses (responseTimeoutMicros request) responseQueue (pullRequestBatch request) []) + (\case + Left _ -> pure () + Right handle -> void (Nats.unsubscribe natsClient handle [])) + (\case + Left err -> + pure (Left (JetStreamNatsError err)) + Right _ -> do + published <- Nats.publish natsClient requestSubject + (pullRequestPayload (pullRequestBatch request) request) + [Nats.withReplyTo inbox] + case published of + Left err -> + pure (Left (JetStreamNatsError err)) + Right () -> + collectResponses waitMicros responseQueue (pullRequestBatch request) []) where natsClient = contextClient context requestSubject = Subject.consumerNextSubject context stream consumer + waitMicros = min + (responseTimeoutMicros request) + (requestTimeoutMicros context requestOptions) ackMessage :: Nats.Client -> Message -> IO (Either JetStreamError ()) ackMessage natsClient = publishAck natsClient Ack @@ -306,9 +333,8 @@ data PullResult = PullResultMessage Message | PullResultStatus PullStatus | PullResultTimeout - | PullResultClosed -collectResponses :: Int -> TQueue (Maybe Nats.MsgView) -> Int -> [Message] -> IO (Either JetStreamError PullResponse) +collectResponses :: Int -> TQueue Nats.MsgView -> Int -> [Message] -> IO (Either JetStreamError PullResponse) collectResponses _ _ 0 messages = pure (Right (PullResponse (reverse messages) Nothing)) collectResponses waitMicros responseQueue remaining messages = do @@ -320,13 +346,10 @@ pure (Right (PullResponse (reverse messages) (Just status))) PullResultTimeout -> pure (Left JetStreamTimeout) - PullResultClosed -> - pure (Left JetStreamNoReply) -classifyPullResult :: Maybe (Maybe Nats.MsgView) -> PullResult +classifyPullResult :: Maybe Nats.MsgView -> PullResult classifyPullResult Nothing = PullResultTimeout -classifyPullResult (Just Nothing) = PullResultClosed -classifyPullResult (Just (Just msgView)) = +classifyPullResult (Just msgView) = case messageStatus message of Nothing -> PullResultMessage message Just status -> PullResultStatus status @@ -349,7 +372,11 @@ Nothing -> pure (Left JetStreamNoReply) Just reply -> - Right <$> Nats.publish natsClient reply [Nats.withPayload (ackPayload verb)] + mapNatsResult <$> Nats.publish natsClient reply (ackPayload verb) [] + +mapNatsResult :: Either Nats.NatsError () -> Either JetStreamError () +mapNatsResult = + either (Left . JetStreamNatsError) Right responseTimeoutMicros :: PullRequest -> Int responseTimeoutMicros request = diff --git a/jetstream/JetStream/Message/API.hs b/jetstream/JetStream/Message/API.hs --- a/jetstream/JetStream/Message/API.hs +++ b/jetstream/JetStream/Message/API.hs @@ -1,14 +1,42 @@ module JetStream.Message.API - ( MessageAPI (..) + ( MessageAPI + , fetch + , consumePush + , createOrderedConsumer + , ack + , nak + , inProgress + , term , FetchOption , FetchWait (..) , Headers - , Message (..) - , OrderedConsumer (..) + , Message + , messageSubject + , messagePayload + , messageHeaders + , messageReplyTo + , messageStatus + , MessageMetadata + , messageMetadata + , messageMetadataStream + , messageMetadataConsumer + , messageMetadataStreamSequence + , messageMetadataConsumerSequence + , messageMetadataNumDelivered + , messageMetadataNumPending + , messageMetadataTimestamp + , messageMetadataDomain + , OrderedConsumer + , orderedConsumerInfo + , fetchOrdered + , stopOrderedConsumer , OrderedConsumerOption , PushConsumeOption - , PushSubscription (..) - , PullResponse (..) + , PushSubscription + , stopPushSubscription + , PullResponse + , pullResponseMessages + , pullResponseStatus , PullStatus (..) , withFetchBatch , withFetchWait @@ -21,18 +49,46 @@ , withPushQueueGroup ) where -import JetStream.Error (JetStreamError) import JetStream.Message.Types ( FetchOption , FetchWait (..) , Headers , Message (..) - , OrderedConsumer (..) + , MessageAPI + , MessageMetadata + , OrderedConsumer , OrderedConsumerOption - , PullResponse (..) + , PullResponse , PullStatus (..) , PushConsumeOption - , PushSubscription (..) + , PushSubscription + , ack + , consumePush + , createOrderedConsumer + , fetch + , fetchOrdered + , inProgress + , messageHeaders + , messageMetadata + , messageMetadataConsumer + , messageMetadataConsumerSequence + , messageMetadataDomain + , messageMetadataNumDelivered + , messageMetadataNumPending + , messageMetadataStream + , messageMetadataStreamSequence + , messageMetadataTimestamp + , messagePayload + , messageReplyTo + , messageStatus + , messageSubject + , nak + , orderedConsumerInfo + , pullResponseMessages + , pullResponseStatus + , stopOrderedConsumer + , stopPushSubscription + , term , withFetchBatch , withFetchWait , withOrderedConsumerDeliverPolicy @@ -43,22 +99,3 @@ , withOrderedConsumerReplayPolicy , withPushQueueGroup ) -import JetStream.Types (ConsumerName, StreamName, Subject) - --- | Pull-consumer message operations. -data MessageAPI = MessageAPI - { fetch :: StreamName -> ConsumerName -> [FetchOption] -> IO (Either JetStreamError PullResponse) - -- ^ Fetch messages for a pull consumer. - , consumePush :: Subject -> [PushConsumeOption] -> (Message -> IO ()) -> IO PushSubscription - -- ^ Subscribe to a push consumer deliver subject. - , createOrderedConsumer :: StreamName -> [OrderedConsumerOption] -> IO (Either JetStreamError OrderedConsumer) - -- ^ Create a client-managed ordered pull consumer. - , ack :: Message -> IO (Either JetStreamError ()) - -- ^ Acknowledge successful message processing. - , nak :: Message -> IO (Either JetStreamError ()) - -- ^ Negatively acknowledge a message. - , inProgress :: Message -> IO (Either JetStreamError ()) - -- ^ Tell the server that message processing is still in progress. - , term :: Message -> IO (Either JetStreamError ()) - -- ^ Terminate message redelivery. - } diff --git a/jetstream/JetStream/Message/Types.hs b/jetstream/JetStream/Message/Types.hs --- a/jetstream/JetStream/Message/Types.hs +++ b/jetstream/JetStream/Message/Types.hs @@ -1,7 +1,8 @@ {-# LANGUAGE OverloadedStrings #-} module JetStream.Message.Types - ( AckVerb (..) + ( MessageAPI (..) + , AckVerb (..) , FetchOption , FetchWait (..) , Headers @@ -48,6 +49,7 @@ import Data.Maybe (isJust) import Data.Time.Clock (NominalDiffTime, UTCTime) import qualified Data.Time.Clock.POSIX as Time +import Data.Word (Word64) import JetStream.Consumer.Types (ConsumerFilter, ConsumerInfo) import JetStream.Error (JetStreamError) import JetStream.Types @@ -55,6 +57,7 @@ , ConsumerName , DeliverPolicy (..) , Headers + , JetStreamRequestOption , Payload , ReplayPolicy , StreamName @@ -62,6 +65,18 @@ , applyCallOptions ) +-- | Message consumption and acknowledgement operations. The public module +-- hides the constructor so this capability can grow additively. +data MessageAPI = MessageAPI + { fetch :: StreamName -> ConsumerName -> [FetchOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError PullResponse) + , consumePush :: Subject -> [PushConsumeOption] -> [JetStreamRequestOption] -> (Message -> IO ()) -> IO (Either JetStreamError PushSubscription) + , createOrderedConsumer :: StreamName -> [OrderedConsumerOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError OrderedConsumer) + , ack :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) + , nak :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) + , inProgress :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) + , term :: Message -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) + } + data PullRequest = PullRequest { pullRequestBatch :: Int , pullRequestWait :: FetchWait @@ -101,7 +116,7 @@ data Message = Message { messageSubject :: Subject - , messagePayload :: Maybe Payload + , messagePayload :: Payload , messageHeaders :: Maybe Headers , messageReplyTo :: Maybe Subject , messageStatus :: Maybe PullStatus @@ -111,8 +126,8 @@ data MessageMetadata = MessageMetadata { messageMetadataStream :: StreamName , messageMetadataConsumer :: ConsumerName - , messageMetadataStreamSequence :: Integer - , messageMetadataConsumerSequence :: Integer + , messageMetadataStreamSequence :: Word64 + , messageMetadataConsumerSequence :: Word64 , messageMetadataNumDelivered :: Integer , messageMetadataNumPending :: Integer , messageMetadataTimestamp :: UTCTime @@ -120,7 +135,7 @@ } deriving (Eq, Show) -newtype PushSubscription = PushSubscription { stopPushSubscription :: IO () } +newtype PushSubscription = PushSubscription { stopPushSubscription :: IO (Either JetStreamError ()) } newtype PushConsumeConfig = PushConsumeConfig { pushConsumeQueueGroup :: Maybe Subject } @@ -138,9 +153,9 @@ config { pushConsumeQueueGroup = Just queueGroup } data OrderedConsumer = OrderedConsumer - { orderedConsumerInfo :: IO (Either JetStreamError ConsumerInfo) - , fetchOrdered :: [FetchOption] -> IO (Either JetStreamError PullResponse) - , stopOrderedConsumer :: IO () + { orderedConsumerInfo :: [JetStreamRequestOption] -> IO (Either JetStreamError ConsumerInfo) + , fetchOrdered :: [FetchOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError PullResponse) + , stopOrderedConsumer :: [JetStreamRequestOption] -> IO (Either JetStreamError ()) } data OrderedConsumerConfig = OrderedConsumerConfig @@ -283,8 +298,8 @@ stream <- tokenAt ackStreamTokenPos tokens consumer <- tokenAt ackConsumerTokenPos tokens delivered <- parseNum =<< tokenAt ackNumDeliveredTokenPos tokens - streamSeq <- parseNum =<< tokenAt ackStreamSeqTokenPos tokens - consumerSeq <- parseNum =<< tokenAt ackConsumerSeqTokenPos tokens + streamSeq <- parseSequence =<< tokenAt ackStreamSeqTokenPos tokens + consumerSeq <- parseSequence =<< tokenAt ackConsumerSeqTokenPos tokens timestamp <- nanosToTime <$> (parseNum =<< tokenAt ackTimestampTokenPos tokens) pending <- parseNum =<< tokenAt ackNumPendingTokenPos tokens let domain = nonEmpty =<< tokenAt ackDomainTokenPos tokens @@ -328,6 +343,13 @@ | BS.null bytes = Just 0 | BC.all isDigit bytes = Just (read (BC.unpack bytes)) | otherwise = Nothing + +parseSequence :: ByteString -> Maybe Word64 +parseSequence bytes = do + value <- parseNum bytes + if value < 0 || value > toInteger (maxBound :: Word64) + then Nothing + else Just (fromInteger value) nanosToTime :: Integer -> UTCTime nanosToTime nanoseconds = diff --git a/jetstream/JetStream/Options.hs b/jetstream/JetStream/Options.hs --- a/jetstream/JetStream/Options.hs +++ b/jetstream/JetStream/Options.hs @@ -1,16 +1,44 @@ +{-# LANGUAGE OverloadedStrings #-} + module JetStream.Options - ( JetStreamConfig (..) + ( JetStream (..) + , JetStreamConfig (..) + , JetStreamConfigError (..) , JetStreamContext (..) , JetStreamOption + , JetStreamRequestOption , defaultJetStreamConfig , newJetStreamContext + , tryNewJetStreamContext + , requestTimeoutMicros , withDomain + , withRequestTimeout , withRequestTimeoutMicros ) where -import qualified API as Nats -import qualified Data.ByteString as BS +import qualified Client.API as Nats +import qualified Data.ByteString as BS +import JetStream.Consumer.API (ConsumerAPI) +import JetStream.Management.API (ManagementAPI) +import JetStream.Message.API (MessageAPI) +import JetStream.Publish.API (PublishAPI) +import JetStream.Stream.API (StreamAPI) +import JetStream.Types + ( JetStreamRequestOption + , applyRequestOptions + , withRequestTimeout + ) +-- | JetStream capabilities. The constructor is kept in the internal package; +-- the public API exposes this type abstractly and exports its accessors. +data JetStream = JetStream + { streams :: StreamAPI + , consumers :: ConsumerAPI + , publisher :: PublishAPI + , messages :: MessageAPI + , management :: ManagementAPI + } + data JetStreamConfig = JetStreamConfig { configDomain :: Maybe BS.ByteString , configRequestTimeoutMicros :: Int @@ -23,8 +51,12 @@ , contextRequestTimeoutMicros :: Int } -type JetStreamOption = JetStreamConfig -> JetStreamConfig +newtype JetStreamOption = JetStreamOption (JetStreamConfig -> JetStreamConfig) +data JetStreamConfigError = EmptyJetStreamDomain + | InvalidJetStreamRequestTimeout Int + deriving (Eq, Show) + defaultJetStreamConfig :: JetStreamConfig defaultJetStreamConfig = JetStreamConfig @@ -34,17 +66,40 @@ newJetStreamContext :: Nats.Client -> [JetStreamOption] -> JetStreamContext newJetStreamContext client options = - let config = foldl (flip ($)) defaultJetStreamConfig options + let config = applyJetStreamOptions options defaultJetStreamConfig in JetStreamContext - { contextClient = client - , contextDomain = configDomain config - , contextRequestTimeoutMicros = configRequestTimeoutMicros config - } + { contextClient = client + , contextDomain = configDomain config + , contextRequestTimeoutMicros = max 1 (configRequestTimeoutMicros config) + } +tryNewJetStreamContext :: Nats.Client -> [JetStreamOption] -> Either JetStreamConfigError JetStreamContext +tryNewJetStreamContext client options = + let config = applyJetStreamOptions options defaultJetStreamConfig + in validateJetStreamConfig config >> pure (newJetStreamContext client options) + +applyJetStreamOptions :: [JetStreamOption] -> JetStreamConfig -> JetStreamConfig +applyJetStreamOptions options config = + foldl apply config options + where + apply value (JetStreamOption option) = option value + +validateJetStreamConfig :: JetStreamConfig -> Either JetStreamConfigError () +validateJetStreamConfig config + | configDomain config == Just "" = Left EmptyJetStreamDomain + | configRequestTimeoutMicros config <= 0 = + Left (InvalidJetStreamRequestTimeout (configRequestTimeoutMicros config)) + | otherwise = Right () + +requestTimeoutMicros :: JetStreamContext -> [JetStreamRequestOption] -> Int +requestTimeoutMicros context = + applyRequestOptions (contextRequestTimeoutMicros context) + withDomain :: BS.ByteString -> JetStreamOption -withDomain domain config = - config { configDomain = Just domain } +withDomain domain = + JetStreamOption $ \config -> config { configDomain = Just domain } withRequestTimeoutMicros :: Int -> JetStreamOption -withRequestTimeoutMicros timeoutMicros config = - config { configRequestTimeoutMicros = max 1 timeoutMicros } +withRequestTimeoutMicros timeoutMicros = + JetStreamOption $ \config -> + config { configRequestTimeoutMicros = timeoutMicros } diff --git a/jetstream/JetStream/Protocol/Headers.hs b/jetstream/JetStream/Protocol/Headers.hs --- a/jetstream/JetStream/Protocol/Headers.hs +++ b/jetstream/JetStream/Protocol/Headers.hs @@ -9,7 +9,7 @@ , statusError ) where -import qualified API as Nats +import qualified Client.API as Nats import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as BC import Data.Maybe (fromMaybe, listToMaybe) diff --git a/jetstream/JetStream/Protocol/Request.hs b/jetstream/JetStream/Protocol/Request.hs --- a/jetstream/JetStream/Protocol/Request.hs +++ b/jetstream/JetStream/Protocol/Request.hs @@ -8,8 +8,7 @@ , decodeJetStreamResponse ) where -import qualified API as Nats -import Control.Concurrent.STM +import qualified Client.API as Nats import Control.Monad (void) import Data.Aeson import qualified Data.Aeson as Aeson @@ -17,50 +16,50 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS import JetStream.Error -import JetStream.Options (JetStreamContext (..)) +import JetStream.Options + ( JetStreamContext (..) + , requestTimeoutMicros + ) import JetStream.Protocol.Headers (statusError) -import JetStream.Types (Headers, Payload, Subject) -import System.Timeout (timeout) +import JetStream.Types + ( Headers + , JetStreamRequestOption + , Payload + , Subject + ) -requestMsg :: JetStreamContext -> Subject -> Maybe Payload -> Headers -> IO (Either JetStreamError Nats.MsgView) -requestMsg ctx subject payload headers = do - replyBox <- newEmptyTMVarIO - let publishOptions = - maybe [] (\body -> [Nats.withPayload body]) payload - ++ [Nats.withReplyCallback (atomically . putTMVar replyBox)] - ++ [Nats.withHeaders headers | not (null headers)] - Nats.publish (contextClient ctx) subject publishOptions - result <- timeout (contextRequestTimeoutMicros ctx) (atomically (readTMVar replyBox)) - case result of - Nothing -> - pure (Left JetStreamTimeout) - Just Nothing -> - pure (Left JetStreamNoReply) - Just (Just msg) -> - pure (Right msg) +requestMsg :: JetStreamContext -> Subject -> Payload -> Headers -> [JetStreamRequestOption] -> IO (Either JetStreamError Nats.MsgView) +requestMsg ctx subject payload headers options = do + result <- Nats.request (contextClient ctx) subject payload requestOptions + pure $ case result of + Left Nats.NatsRequestTimedOut -> Left JetStreamTimeout + Left err -> Left (JetStreamNatsError err) + Right msg -> Right msg + where + timeoutSeconds = + fromRational (toRational (requestTimeoutMicros ctx options) / 1000000) + requestOptions = + Nats.withRequestTimeout timeoutSeconds + : [Nats.withRequestHeaders headers | not (null headers)] -requestJSON :: FromJSON a => JetStreamContext -> Subject -> Maybe Value -> IO (Either JetStreamError a) +requestJSON :: FromJSON a => JetStreamContext -> Subject -> Maybe Value -> [JetStreamRequestOption] -> IO (Either JetStreamError a) requestJSON ctx subject body = requestJSONWithHeaders ctx subject body [] -requestJSONWithHeaders :: FromJSON a => JetStreamContext -> Subject -> Maybe Value -> Headers -> IO (Either JetStreamError a) -requestJSONWithHeaders ctx subject body headers = do - msgResult <- requestMsg ctx subject (strictEncode <$> body) headers +requestJSONWithHeaders :: FromJSON a => JetStreamContext -> Subject -> Maybe Value -> Headers -> [JetStreamRequestOption] -> IO (Either JetStreamError a) +requestJSONWithHeaders ctx subject body headers options = do + msgResult <- requestMsg ctx subject (maybe BS.empty strictEncode body) headers options pure $ do msg <- msgResult case statusError msg of Just err -> Left err Nothing -> - case Nats.payload msg of - Nothing -> - Left JetStreamNoReply - Just payload -> - decodeJetStreamResponse payload + decodeJetStreamResponse (Nats.payload msg) -requestUnit :: JetStreamContext -> Subject -> Maybe Value -> IO (Either JetStreamError ()) -requestUnit ctx subject body = do - response <- requestJSON ctx subject body +requestUnit :: JetStreamContext -> Subject -> Maybe Value -> [JetStreamRequestOption] -> IO (Either JetStreamError ()) +requestUnit ctx subject body options = do + response <- requestJSON ctx subject body options pure (void (response :: Either JetStreamError SuccessResponse)) decodeJetStreamResponse :: FromJSON a => BS.ByteString -> Either JetStreamError a diff --git a/jetstream/JetStream/Publish.hs b/jetstream/JetStream/Publish.hs --- a/jetstream/JetStream/Publish.hs +++ b/jetstream/JetStream/Publish.hs @@ -3,7 +3,8 @@ , module JetStream.Publish.API ) where -import qualified API as Nats +import qualified Client.API as Nats +import qualified Data.ByteString as BS import JetStream.Error (JetStreamError (JetStreamNoReply)) import JetStream.Options (JetStreamContext) import JetStream.Protocol.Headers (statusError) @@ -12,8 +13,12 @@ , requestMsg ) import JetStream.Publish.API -import JetStream.Publish.Types (publishHeaders) -import JetStream.Types (Payload, Subject) +import JetStream.Publish.Types (PublishAPI (..), publishHeaders) +import JetStream.Types + ( JetStreamRequestOption + , Payload + , Subject + ) publishAPI :: JetStreamContext -> PublishAPI publishAPI context = @@ -26,17 +31,16 @@ -> Subject -> Payload -> [PublishOption] + -> [JetStreamRequestOption] -> IO (Either JetStreamError PublishAck) -publishMessage context subject payload options = do - msgResult <- requestMsg context subject (Just payload) (publishHeaders options) +publishMessage context subject payload options requestOptions = do + msgResult <- requestMsg context subject payload (publishHeaders options) requestOptions pure $ do msg <- msgResult case statusError msg of Just err -> Left err Nothing -> - case Nats.payload msg of - Nothing -> - Left JetStreamNoReply - Just ackPayload -> - decodeJetStreamResponse ackPayload + if BS.null (Nats.payload msg) + then Left JetStreamNoReply + else decodeJetStreamResponse (Nats.payload msg) diff --git a/jetstream/JetStream/Publish/API.hs b/jetstream/JetStream/Publish/API.hs --- a/jetstream/JetStream/Publish/API.hs +++ b/jetstream/JetStream/Publish/API.hs @@ -1,6 +1,11 @@ module JetStream.Publish.API - ( PublishAPI (..) - , PublishAck (..) + ( PublishAPI + , publish + , PublishAck + , publishAckStream + , publishAckSequence + , publishAckDuplicate + , publishAckDomain , PublishExpectation (..) , PublishOption , withMsgId @@ -9,20 +14,18 @@ , withHeaders ) where -import JetStream.Error (JetStreamError) import JetStream.Publish.Types - ( PublishAck (..) + ( PublishAPI + , PublishAck , PublishExpectation (..) , PublishOption + , publish + , publishAckDomain + , publishAckDuplicate + , publishAckSequence + , publishAckStream , withExpectedStream , withHeaders , withMsgId , withPublishExpectation ) -import JetStream.Types (Payload, Subject) - --- | Publish operations for JetStream subjects. --- --- The 'publish' field publishes a message and waits for a JetStream publish --- acknowledgement. -newtype PublishAPI = PublishAPI { publish :: Subject -> Payload -> [PublishOption] -> IO (Either JetStreamError PublishAck) } diff --git a/jetstream/JetStream/Publish/Types.hs b/jetstream/JetStream/Publish/Types.hs --- a/jetstream/JetStream/Publish/Types.hs +++ b/jetstream/JetStream/Publish/Types.hs @@ -1,7 +1,8 @@ {-# LANGUAGE OverloadedStrings #-} module JetStream.Publish.Types - ( PublishAck (..) + ( PublishAPI (..) + , PublishAck (..) , PublishExpectation (..) , PublishOption , withMsgId @@ -16,8 +17,18 @@ import qualified Data.ByteString.Char8 as BC import Data.Text.Encoding (encodeUtf8) import Data.Word (Word64) -import JetStream.Types (StreamName) +import JetStream.Error (JetStreamError) +import JetStream.Types + ( JetStreamRequestOption + , Payload + , StreamName + , Subject + ) +-- | Publish operations. The public API hides the constructor so additional +-- publish modes can be introduced without changing downstream code. +newtype PublishAPI = PublishAPI { publish :: Subject -> Payload -> [PublishOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError PublishAck) } + data PublishAck = PublishAck { publishAckStream :: StreamName , publishAckSequence :: Word64 @@ -69,9 +80,9 @@ (key, value) = case expectation of ExpectedLastSequence sequenceNumber -> - ("Nats-Expected-Last-Sequence", renderWord64 sequenceNumber) + ("Nats-Expected-Last-Sequence", renderSequence sequenceNumber) ExpectedLastSubjectSequence sequenceNumber -> - ("Nats-Expected-Last-Subject-Sequence", renderWord64 sequenceNumber) + ("Nats-Expected-Last-Subject-Sequence", renderSequence sequenceNumber) ExpectedLastMsgId msgId -> ("Nats-Expected-Last-Msg-Id", msgId) @@ -85,7 +96,7 @@ publishHeaders :: [PublishOption] -> [(BS.ByteString, BS.ByteString)] publishHeaders options = - publishConfigHeaders (foldr ($) defaultPublishConfig options) + publishConfigHeaders (foldl (flip ($)) defaultPublishConfig options) putHeader :: BS.ByteString -> BS.ByteString -> PublishOption putHeader key value config = @@ -112,3 +123,6 @@ renderWord64 :: Word64 -> BS.ByteString renderWord64 = BC.pack . show + +renderSequence :: Word64 -> BS.ByteString +renderSequence = renderWord64 diff --git a/jetstream/JetStream/Stream.hs b/jetstream/JetStream/Stream.hs --- a/jetstream/JetStream/Stream.hs +++ b/jetstream/JetStream/Stream.hs @@ -11,9 +11,10 @@ import JetStream.Options (JetStreamContext) import qualified JetStream.Protocol.Request as Request import qualified JetStream.Protocol.Subject as Subject -import JetStream.Stream.API (StreamAPI (..)) +import JetStream.Stream.API import JetStream.Stream.Types - ( purgeStreamRequest + ( StreamAPI (..) + , purgeStreamRequest , streamConfigRequest , streamListRequest , streamMessageDeleteRequest @@ -24,52 +25,57 @@ streamAPI :: JetStreamContext -> StreamAPI streamAPI context = StreamAPI - { create = \name subjects options -> + { create = \name subjects options requestOptions -> let config = streamConfigRequest name subjects options in Request.requestJSON context (Subject.streamCreateSubject context name) (Just (toJSON config)) - , createOrUpdate = \name subjects options -> do - updateResult <- update (streamAPI context) name subjects options + requestOptions + , createOrUpdate = \name subjects options requestOptions -> do + updateResult <- update (streamAPI context) name subjects options requestOptions case updateResult of Left err | isStreamNotFound err -> - create (streamAPI context) name subjects options + create (streamAPI context) name subjects options requestOptions other -> pure other - , update = \name subjects options -> + , update = \name subjects options requestOptions -> let config = streamConfigRequest name subjects options in Request.requestJSON context (Subject.streamUpdateSubject context name) (Just (toJSON config)) - , info = \streamName -> + requestOptions + , info = \streamName requestOptions -> Request.requestJSON context (Subject.streamInfoSubject context streamName) Nothing - , getMessage = \streamName selector -> + requestOptions + , getMessage = \streamName selector requestOptions -> Request.requestJSON context (Subject.streamMessageGetSubject context streamName) (Just (toJSON (streamMessageGetRequest selector))) - , deleteMessage = \streamName sequenceNumber mode -> + requestOptions + , deleteMessage = \streamName sequenceNumber mode requestOptions -> Request.requestJSON context (Subject.streamMessageDeleteSubject context streamName) (Just (toJSON (streamMessageDeleteRequest sequenceNumber mode))) - , delete = \streamName -> + requestOptions + , delete = \streamName requestOptions -> Request.requestJSON context (Subject.streamDeleteSubject context streamName) Nothing - , purge = \streamName options -> + requestOptions + , purge = \streamName options requestOptions -> Request.requestJSON context (Subject.streamPurgeSubject context streamName) (Just (toJSON (purgeStreamRequest options))) + requestOptions , list = - Request.requestJSON context - (Subject.streamListSubject context) . - Just . toJSON . streamListRequest + Request.requestJSON context (Subject.streamListSubject context) + . Just . toJSON . streamListRequest , names = - Request.requestJSON context - (Subject.streamNamesSubject context) . - Just . toJSON . streamNamesRequest + Request.requestJSON context (Subject.streamNamesSubject context) + . Just . toJSON . streamNamesRequest } isStreamNotFound :: JetStreamError -> Bool diff --git a/jetstream/JetStream/Stream/API.hs b/jetstream/JetStream/Stream/API.hs --- a/jetstream/JetStream/Stream/API.hs +++ b/jetstream/JetStream/Stream/API.hs @@ -1,9 +1,35 @@ module JetStream.Stream.API - ( StreamAPI (..) + ( StreamAPI + , create + , createOrUpdate + , update + , info + , getMessage + , deleteMessage + , delete + , purge + , list + , names + , Stream + , streamName + , lookupStream + , createStream + , getStreamInfo , RetentionPolicy (..) , StorageType (..) , DiscardPolicy (..) - , StreamConfig (..) + , StreamConfig + , streamConfigName + , streamConfigSubjects + , streamConfigRetention + , streamConfigStorage + , streamConfigDiscard + , streamConfigMaxMessages + , streamConfigMaxBytes + , streamConfigMaxAge + , streamConfigReplicas + , streamConfigDuplicateWindow + , streamConfigAllowDirect , StreamConfigOption , withRetention , withStorage @@ -21,45 +47,100 @@ , StreamListOption , withStreamListOffset , withStreamListSubject - , StreamMessage (..) + , StreamMessage + , streamMessageSubject + , streamMessageSequence + , streamMessageHeadersRaw + , streamMessagePayload + , streamMessageTime , StreamMessageSelector (..) , StreamMessageDeleteMode (..) - , DeleteStreamMessageResponse (..) - , StreamInfo (..) - , StreamState (..) - , StreamCluster (..) - , StreamPeer (..) - , StreamSourceInfo (..) - , DeleteStreamResponse (..) - , PurgeStreamResponse (..) - , StreamListResponse (..) - , StreamNamesResponse (..) + , DeleteStreamMessageResponse + , deleteStreamMessageSuccess + , StreamInfo + , streamInfoConfig + , streamInfoCreated + , streamInfoState + , streamInfoCluster + , streamInfoMirror + , streamInfoSources + , StreamState + , streamStateMessages + , streamStateBytes + , streamStateFirstSequence + , streamStateFirstTime + , streamStateLastSequence + , streamStateLastTime + , streamStateConsumerCount + , streamStateDeleted + , streamStateNumDeleted + , streamStateNumSubjects + , StreamCluster + , streamClusterName + , streamClusterLeader + , streamClusterReplicas + , StreamPeer + , streamPeerName + , streamPeerCurrent + , streamPeerOffline + , streamPeerActive + , streamPeerLag + , StreamSourceInfo + , streamSourceInfoName + , streamSourceInfoFilterSubject + , streamSourceInfoLag + , streamSourceInfoActive + , DeleteStreamResponse + , deleteStreamSuccess + , PurgeStreamResponse + , purgeStreamSuccess + , purgeStreamPurged + , StreamListResponse + , streamListTotal + , streamListOffset + , streamListLimit + , streamListStreams + , StreamNamesResponse + , streamNamesTotal + , streamNamesOffset + , streamNamesLimit + , streamNamesStreams ) where import JetStream.Error (JetStreamError) import JetStream.Stream.Types -import JetStream.Types (StreamName, Subject) +import JetStream.Types + ( JetStreamRequestOption + , StreamName + , Subject + ) --- | Stream management operations. -data StreamAPI = StreamAPI - { create :: StreamName -> [Subject] -> [StreamConfigOption] -> IO (Either JetStreamError StreamInfo) - -- ^ Create a stream. - , createOrUpdate :: StreamName -> [Subject] -> [StreamConfigOption] -> IO (Either JetStreamError StreamInfo) - -- ^ Create a stream or update it if it already exists. - , update :: StreamName -> [Subject] -> [StreamConfigOption] -> IO (Either JetStreamError StreamInfo) - -- ^ Update an existing stream configuration. - , info :: StreamName -> IO (Either JetStreamError StreamInfo) - -- ^ Fetch stream configuration and state. - , getMessage :: StreamName -> StreamMessageSelector -> IO (Either JetStreamError StreamMessage) - -- ^ Fetch a stored stream message by selector. - , deleteMessage :: StreamName -> Integer -> StreamMessageDeleteMode -> IO (Either JetStreamError DeleteStreamMessageResponse) - -- ^ Delete or securely erase a stored stream message. - , delete :: StreamName -> IO (Either JetStreamError DeleteStreamResponse) - -- ^ Delete a stream. - , purge :: StreamName -> [PurgeStreamOption] -> IO (Either JetStreamError PurgeStreamResponse) - -- ^ Purge messages from a stream. - , list :: [StreamListOption] -> IO (Either JetStreamError StreamListResponse) - -- ^ List streams and their metadata. - , names :: [StreamListOption] -> IO (Either JetStreamError StreamNamesResponse) - -- ^ List stream names. - } +-- | Look up a stream and return a stable resource handle. +lookupStream + :: StreamAPI + -> StreamName + -> [JetStreamRequestOption] + -> IO (Either JetStreamError Stream) +lookupStream api name requestOptions = + fmap (fmap (const (Stream name))) (info api name requestOptions) + +-- | Create a stream and return a stable resource handle. +createStream + :: StreamAPI + -> StreamName + -> [Subject] + -> [StreamConfigOption] + -> [JetStreamRequestOption] + -> IO (Either JetStreamError Stream) +createStream api name subjects configOptions requestOptions = + fmap (fmap (const (Stream name))) + (create api name subjects configOptions requestOptions) + +-- | Read the current info for a stream handle. +getStreamInfo + :: StreamAPI + -> Stream + -> [JetStreamRequestOption] + -> IO (Either JetStreamError StreamInfo) +getStreamInfo api stream = + info api (streamName stream) diff --git a/jetstream/JetStream/Stream/Types.hs b/jetstream/JetStream/Stream/Types.hs --- a/jetstream/JetStream/Stream/Types.hs +++ b/jetstream/JetStream/Stream/Types.hs @@ -1,7 +1,9 @@ {-# LANGUAGE OverloadedStrings #-} module JetStream.Stream.Types - ( RetentionPolicy (..) + ( StreamAPI (..) + , Stream (..) + , RetentionPolicy (..) , StorageType (..) , DiscardPolicy (..) , StreamConfig (..) @@ -51,9 +53,12 @@ import qualified Data.ByteString.Base64 as Base64 import Data.Maybe (catMaybes) import Data.Time.Clock (NominalDiffTime, UTCTime) +import Data.Word (Word64) +import JetStream.Error (JetStreamError) import JetStream.Types ( CallOption , DiscardPolicy (..) + , JetStreamRequestOption , Payload , RetentionPolicy (..) , StorageType (..) @@ -64,6 +69,26 @@ , parseByteString ) +-- | Stream management operations. The public module exposes this type +-- abstractly so operations can be added without changing its constructor. +data StreamAPI = StreamAPI + { create :: StreamName -> [Subject] -> [StreamConfigOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamInfo) + , createOrUpdate :: StreamName -> [Subject] -> [StreamConfigOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamInfo) + , update :: StreamName -> [Subject] -> [StreamConfigOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamInfo) + , info :: StreamName -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamInfo) + , getMessage :: StreamName -> StreamMessageSelector -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamMessage) + , deleteMessage :: StreamName -> Word64 -> StreamMessageDeleteMode -> [JetStreamRequestOption] -> IO (Either JetStreamError DeleteStreamMessageResponse) + , delete :: StreamName -> [JetStreamRequestOption] -> IO (Either JetStreamError DeleteStreamResponse) + , purge :: StreamName -> [PurgeStreamOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError PurgeStreamResponse) + , list :: [StreamListOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamListResponse) + , names :: [StreamListOption] -> [JetStreamRequestOption] -> IO (Either JetStreamError StreamNamesResponse) + } + +-- | A stable reference to a stream. It is intentionally just an identity; +-- operations remain on 'StreamAPI', so handles do not retain resources. +newtype Stream = Stream { streamName :: StreamName } + deriving (Eq, Ord, Show) + data StreamConfigRequest = StreamConfigRequest { streamConfigRequestName :: StreamName , streamConfigRequestSubjects :: [Subject] @@ -151,7 +176,7 @@ data PurgeStreamRequest = PurgeStreamRequest { purgeStreamSubject :: Maybe Subject - , purgeStreamSequence :: Maybe Integer + , purgeStreamSequence :: Maybe Word64 , purgeStreamKeep :: Maybe Integer } deriving (Eq, Show) @@ -171,7 +196,7 @@ withPurgeSubject subject request = request { purgeStreamSubject = Just subject } -withPurgeSequence :: Integer -> PurgeStreamOption +withPurgeSequence :: Word64 -> PurgeStreamOption withPurgeSequence sequenceNumber request = request { purgeStreamSequence = Just sequenceNumber } @@ -223,12 +248,12 @@ data StreamState = StreamState { streamStateMessages :: Integer , streamStateBytes :: Integer - , streamStateFirstSequence :: Integer + , streamStateFirstSequence :: Word64 , streamStateFirstTime :: UTCTime - , streamStateLastSequence :: Integer + , streamStateLastSequence :: Word64 , streamStateLastTime :: UTCTime , streamStateConsumerCount :: Int - , streamStateDeleted :: [Integer] + , streamStateDeleted :: [Word64] , streamStateNumDeleted :: Integer , streamStateNumSubjects :: Integer } @@ -285,14 +310,14 @@ data StreamMessage = StreamMessage { streamMessageSubject :: Subject - , streamMessageSequence :: Integer + , streamMessageSequence :: Word64 , streamMessageHeadersRaw :: Maybe BS.ByteString , streamMessagePayload :: Maybe Payload , streamMessageTime :: UTCTime } deriving (Eq, Show) -data StreamMessageSelector = StreamMessageBySequence Integer +data StreamMessageSelector = StreamMessageBySequence Word64 | LastStreamMessageForSubject Subject | NextStreamMessageForSubject Subject deriving (Eq, Show) @@ -308,12 +333,12 @@ deriving (Eq, Show) data StreamMessageDeleteRequest = StreamMessageDeleteRequest - { streamMessageDeleteSequence :: Integer + { streamMessageDeleteSequence :: Word64 , streamMessageDeleteNoErase :: Maybe Bool } deriving (Eq, Show) -streamMessageDeleteRequest :: Integer -> StreamMessageDeleteMode -> StreamMessageDeleteRequest +streamMessageDeleteRequest :: Word64 -> StreamMessageDeleteMode -> StreamMessageDeleteRequest streamMessageDeleteRequest sequenceNumber mode = StreamMessageDeleteRequest { streamMessageDeleteSequence = sequenceNumber diff --git a/jetstream/JetStream/Types.hs b/jetstream/JetStream/Types.hs --- a/jetstream/JetStream/Types.hs +++ b/jetstream/JetStream/Types.hs @@ -11,6 +11,9 @@ , Payload , Headers , CallOption + , JetStreamRequestOption + , applyRequestOptions + , withRequestTimeout , AckPolicy (..) , DeliverPolicy (..) , DiscardPolicy (..) @@ -28,11 +31,12 @@ import qualified Data.Aeson.KeyMap as KeyMap import Data.Aeson.Types (Pair, Parser, typeMismatch) import qualified Data.ByteString as BS -import Data.Maybe (catMaybes) +import Data.Maybe (catMaybes, fromMaybe) import qualified Data.Text as T import qualified Data.Text.Encoding as E import Data.Time.Clock (NominalDiffTime) import qualified Data.Time.Clock as Time +import Data.Word (Word64) type StreamName = BS.ByteString type ConsumerName = BS.ByteString @@ -41,6 +45,29 @@ type Headers = [(BS.ByteString, BS.ByteString)] type CallOption a = a -> a +-- | Options scoped to one JetStream API request. This is distinct from +-- stream, consumer, publish, and fetch configuration. +newtype JetStreamRequestOption = JetStreamRequestOption (JetStreamRequestConfig -> JetStreamRequestConfig) + +newtype JetStreamRequestConfig = JetStreamRequestConfig { requestConfigTimeoutMicros :: Maybe Int } + +applyRequestOptions :: Int -> [JetStreamRequestOption] -> Int +applyRequestOptions defaultTimeout options = + fromMaybe defaultTimeout + (requestConfigTimeoutMicros (foldl apply emptyRequestConfig options)) + where + emptyRequestConfig = JetStreamRequestConfig Nothing + apply value (JetStreamRequestOption option) = option value + +-- | Override the timeout for one JetStream request. +withRequestTimeout :: NominalDiffTime -> JetStreamRequestOption +withRequestTimeout timeoutSeconds = + JetStreamRequestOption $ \config -> + config + { requestConfigTimeoutMicros = + Just (max 1 (floor (realToFrac timeoutSeconds * (1000000 :: Double)))) + } + data AccountInfo = AccountInfo { accountInfoTier :: AccountTier , accountInfoDomain :: Maybe BS.ByteString @@ -82,38 +109,52 @@ applyCallOptions :: [CallOption a] -> a -> a applyCallOptions options value = - foldr ($) value options + foldl (flip ($)) value options -data AckPolicy = AckNone | AckAll | AckExplicit +data AckPolicy = AckNone + | AckAll + | AckExplicit + | AckPolicyUnknown T.Text deriving (Eq, Show) data DeliverPolicy = DeliverAll | DeliverLast | DeliverNew - | DeliverByStartSequence Int + | DeliverByStartSequence Word64 | DeliverByStartTime Time.UTCTime | DeliverLastPerSubject + | DeliverPolicyUnknown T.Text deriving (Eq, Show) -data DiscardPolicy = DiscardOld | DiscardNew +data DiscardPolicy = DiscardOld + | DiscardNew + | DiscardPolicyUnknown T.Text deriving (Eq, Show) -data ReplayPolicy = ReplayInstant | ReplayOriginal +data ReplayPolicy = ReplayInstant + | ReplayOriginal + | ReplayPolicyUnknown T.Text deriving (Eq, Show) -data RetentionPolicy = LimitsPolicy | InterestPolicy | WorkQueuePolicy +data RetentionPolicy = LimitsPolicy + | InterestPolicy + | WorkQueuePolicy + | RetentionPolicyUnknown T.Text deriving (Eq, Show) -data StorageType = FileStorage | MemoryStorage +data StorageType = FileStorage + | MemoryStorage + | StorageTypeUnknown T.Text deriving (Eq, Show) instance ToJSON AckPolicy where toJSON policy = String $ case policy of - AckNone -> "none" - AckAll -> "all" - AckExplicit -> "explicit" + AckNone -> "none" + AckAll -> "all" + AckExplicit -> "explicit" + AckPolicyUnknown value -> value instance FromJSON AckPolicy where parseJSON = withText "AckPolicy" $ \value -> @@ -121,18 +162,19 @@ "none" -> pure AckNone "all" -> pure AckAll "explicit" -> pure AckExplicit - _ -> fail ("unknown ack policy: " ++ T.unpack value) + _ -> pure (AckPolicyUnknown value) instance ToJSON DeliverPolicy where toJSON policy = String $ case policy of - DeliverAll -> "all" - DeliverLast -> "last" - DeliverNew -> "new" - DeliverByStartSequence _ -> "by_start_sequence" - DeliverByStartTime _ -> "by_start_time" - DeliverLastPerSubject -> "last_per_subject" + DeliverAll -> "all" + DeliverLast -> "last" + DeliverNew -> "new" + DeliverByStartSequence _ -> "by_start_sequence" + DeliverByStartTime _ -> "by_start_time" + DeliverLastPerSubject -> "last_per_subject" + DeliverPolicyUnknown value -> value instance FromJSON DeliverPolicy where parseJSON = withText "DeliverPolicy" $ \value -> @@ -143,43 +185,46 @@ "by_start_sequence" -> fail "deliver policy by_start_sequence requires opt_start_seq" "by_start_time" -> fail "deliver policy by_start_time requires opt_start_time" "last_per_subject" -> pure DeliverLastPerSubject - _ -> fail ("unknown deliver policy: " ++ T.unpack value) + _ -> pure (DeliverPolicyUnknown value) instance ToJSON DiscardPolicy where toJSON policy = String $ case policy of - DiscardOld -> "old" - DiscardNew -> "new" + DiscardOld -> "old" + DiscardNew -> "new" + DiscardPolicyUnknown value -> value instance FromJSON DiscardPolicy where parseJSON = withText "DiscardPolicy" $ \value -> case value of "old" -> pure DiscardOld "new" -> pure DiscardNew - _ -> fail ("unknown discard policy: " ++ T.unpack value) + _ -> pure (DiscardPolicyUnknown value) instance ToJSON ReplayPolicy where toJSON policy = String $ case policy of - ReplayInstant -> "instant" - ReplayOriginal -> "original" + ReplayInstant -> "instant" + ReplayOriginal -> "original" + ReplayPolicyUnknown value -> value instance FromJSON ReplayPolicy where parseJSON = withText "ReplayPolicy" $ \value -> case value of "instant" -> pure ReplayInstant "original" -> pure ReplayOriginal - _ -> fail ("unknown replay policy: " ++ T.unpack value) + _ -> pure (ReplayPolicyUnknown value) instance ToJSON RetentionPolicy where toJSON policy = String $ case policy of - LimitsPolicy -> "limits" - InterestPolicy -> "interest" - WorkQueuePolicy -> "workqueue" + LimitsPolicy -> "limits" + InterestPolicy -> "interest" + WorkQueuePolicy -> "workqueue" + RetentionPolicyUnknown value -> value instance FromJSON RetentionPolicy where parseJSON = withText "RetentionPolicy" $ \value -> @@ -187,21 +232,22 @@ "limits" -> pure LimitsPolicy "interest" -> pure InterestPolicy "workqueue" -> pure WorkQueuePolicy - _ -> fail ("unknown retention policy: " ++ T.unpack value) + _ -> pure (RetentionPolicyUnknown value) instance ToJSON StorageType where toJSON storage = String $ case storage of - FileStorage -> "file" - MemoryStorage -> "memory" + FileStorage -> "file" + MemoryStorage -> "memory" + StorageTypeUnknown value -> value instance FromJSON StorageType where parseJSON = withText "StorageType" $ \value -> case value of "file" -> pure FileStorage "memory" -> pure MemoryStorage - _ -> fail ("unknown storage type: " ++ T.unpack value) + _ -> pure (StorageTypeUnknown value) byteStringToJSON :: BS.ByteString -> Value byteStringToJSON = diff --git a/natskell.cabal b/natskell.cabal --- a/natskell.cabal +++ b/natskell.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: natskell -version: 0.3.0.0 +version: 1.0.0.1 synopsis: A NATS client library written in Haskell tested-with: GHC ==8.8, @@ -46,8 +46,14 @@ library import: shared + autogen-modules: + Paths_natskell + other-modules: + Client.Implementation + Paths_natskell build-depends: natskell-internal, + jetstream-internal, stm >= 2.5 && < 2.6, bytestring >= 0.10 && < 0.13, time >= 1.9 && < 1.16 @@ -55,28 +61,22 @@ exposed-modules: API Client - - hs-source-dirs: - client - -library jetstream - import: shared - visibility: public - build-depends: - natskell, - jetstream-internal - - exposed-modules: JetStream.API + JetStream.API.Consumer + JetStream.API.Management + JetStream.API.Message + JetStream.API.Publish + JetStream.API.Stream JetStream.Client hs-source-dirs: - jetstream-client + client library jetstream-internal import: shared + visibility: private build-depends: - natskell, + natskell-internal, aeson >= 2.1 && < 2.3, base64-bytestring >= 1.2 && < 1.3, bytestring >= 0.10 && < 0.13, @@ -92,6 +92,8 @@ JetStream.Message JetStream.Message.API JetStream.Message.Types + JetStream.Management + JetStream.Management.API JetStream.Options JetStream.Protocol.Headers JetStream.Protocol.Request @@ -120,15 +122,20 @@ time >= 1.9 && < 1.16, network >= 3.1 && < 3.3, mtl >= 2.2 && < 2.4, - cryptonite >= 0.29 && < 0.31, + crypton >= 0.34 && < 1.1, memory >= 0.14 && < 0.19, random >= 1.1 && < 1.3, network-simple >= 0.4 && < 0.5, tls >= 1.8 && < 3, + crypton-x509-store >= 1.6 && < 1.10, + crypton-x509-system >= 1.6 && < 1.9, text (>= 1.2 && < 1.3) || (>= 2.0 && < 2.2) exposed-modules: + Client.API Auth.Types + Auth.Config + Auth.Resolver Auth.None Auth.Token Auth.UserPass @@ -158,6 +165,7 @@ Types.Ok Types.Err Types.Info + Types.TLS Types.Pub Types.Connect Types.Sub @@ -183,6 +191,9 @@ Network.Connection.Tls Queue.API Queue.TransactionalQueue + other-modules: + Auth.Credentials + Auth.NKey.Codec hs-source-dirs: internal internal/Types @@ -236,6 +247,7 @@ SidSpec NuidSpec WorkerPoolSpec + SubscriptionStoreSpec WaitGroupSpec JetStream.ConsumerSpec JetStream.MessageSpec @@ -256,6 +268,14 @@ ParserSpec ValidatorsSpec +test-suite public-api-test + import: shared + type: exitcode-stdio-1.0 + hs-source-dirs: test/PublicAPI + main-is: Main.hs + build-depends: + natskell + benchmark parser-bench import: shared type: exitcode-stdio-1.0 @@ -274,10 +294,10 @@ test-suite integration-tests import: shared + ghc-options: -threaded hs-source-dirs: test/Integration build-depends: natskell, - natskell:jetstream, test-common, port-utils >= 0.1 && < 0.3, bytestring, diff --git a/test/Fuzz/ParserSpec.hs b/test/Fuzz/ParserSpec.hs --- a/test/Fuzz/ParserSpec.hs +++ b/test/Fuzz/ParserSpec.hs @@ -35,6 +35,7 @@ , port , proto , server_id + , tls_available , tls_required , version ) @@ -189,6 +190,7 @@ <*> frequency [(2, pure Nothing), (3, Just <$> genSafeText 1 16)] <*> genOptional arbitrary <*> genOptional arbitrary + <*> genOptional arbitrary <*> genOptional (listOf1 (genSafeText 1 24)) <*> genOptional arbitrary <*> genOptional arbitrary @@ -343,6 +345,7 @@ , maybeField "nonce" (quote <$> nonce infoValue) , maybeField "auth_required" (jsonBool <$> auth_required infoValue) , maybeField "tls_required" (jsonBool <$> tls_required infoValue) + , maybeField "tls_available" (jsonBool <$> tls_available infoValue) , maybeField "connect_urls" (jsonArray <$> connect_urls infoValue) , maybeField "ldm" (jsonBool <$> ldm infoValue) , maybeField "headers" (jsonBool <$> headers infoValue) diff --git a/test/Integration/ClientSpec.hs b/test/Integration/ClientSpec.hs --- a/test/Integration/ClientSpec.hs +++ b/test/Integration/ClientSpec.hs @@ -3,15 +3,6 @@ module ClientSpec where import API - ( Client (..) - , MsgView (..) - , withHeaders - , withPayload - , withQueueGroup - , withReplyCallback - , withReplyTo - , withSubscriptionExpiry - ) import Client import Control.Concurrent import Control.Concurrent.STM @@ -42,6 +33,20 @@ defaultINFO = "INFO {\"server_id\": \"some-server\", \"version\": \"semver\", \"go\": \"1.13\", \"host\": \"127.0.0.1\", \"port\": 4222, \"max_payload\": 1024, \"proto\": 3}\r\n" +newClientOrFail servers options = do + result <- newClient servers options + case result of + Left err -> fail ("client failed to connect: " ++ show err) + Right client -> pure client + +pingWithCallback :: Client -> IO () -> IO () +pingWithCallback client action = + void . forkIO $ void (ping client []) >> action + +isValidationFailure :: Either NatsError a -> Bool +isValidationFailure (Left (NatsValidationError _)) = True +isValidationFailure _ = False + tooLongMSG = "MSG a b 5000\r\n" <> BS.replicate 5000 _x <> "\r\n" headerBlock :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString @@ -114,7 +119,7 @@ expectClientBytes :: Socket -> (BS.ByteString -> Bool) -> String -> IO BS.ByteString expectClientBytes sock done failureMessage = do - result <- timeout 1000000 $ + result <- timeout 10000000 $ recvUntil sock done case result of Nothing -> @@ -127,19 +132,24 @@ expectationFailure ("unexpected client bytes: " ++ show bytes) >> fail failureMessage -expectNoClientBytes :: Socket -> (BS.ByteString -> Bool) -> String -> IO () -expectNoClientBytes sock unwanted failureMessage = do - result <- timeout 300000 $ - recvUntil sock unwanted - case result of - Just bytes | unwanted bytes -> - expectationFailure (failureMessage ++ ": " ++ show bytes) - _ -> - pure () +expectNoClientBytesBefore + :: Socket + -> (BS.ByteString -> Bool) + -> BS.ByteString + -> String + -> IO () +expectNoClientBytesBefore sock unwanted barrier failureMessage = do + bytes <- + expectClientBytes + sock + (BS.isInfixOf barrier) + ("client did not send barrier command: " ++ show barrier) + when (unwanted bytes) $ + expectationFailure (failureMessage ++ ": " ++ show bytes) expectClientCommand :: Socket -> BS.ByteString -> IO () expectClientCommand sock command = do - result <- timeout 1000000 $ + result <- timeout 10000000 $ recvUntil sock (BS.isInfixOf command) case result of Nothing -> @@ -150,7 +160,7 @@ expectMVar :: String -> MVar a -> IO a expectMVar failureMessage var = do - result <- timeout 1000000 (takeMVar var) + result <- timeout 10000000 (takeMVar var) case result of Nothing -> expectationFailure failureMessage >> @@ -158,6 +168,17 @@ Just value -> pure value +completeHandshake :: Socket -> IO BS.ByteString +completeHandshake sock = do + sendAll sock defaultINFO + bytes <- + expectClientBytes + sock + (\sent -> BS.isInfixOf "CONNECT " sent && BS.isInfixOf "PING\r\n" sent) + "client did not complete the initial handshake" + sendAll sock "PONG\r\n" + pure bytes + expectQueuedSub :: Socket -> BS.ByteString -> BS.ByteString -> BS.ByteString -> IO () expectQueuedSub sock subject queueGroup sid = expectClientCommand sock (BS.concat ["SUB ", subject, " ", queueGroup, " ", sid, "\r\n"]) @@ -217,17 +238,15 @@ parseCapturedPublish :: BS.ByteString -> Maybe CapturedPublish parseCapturedPublish bytes = do subLine <- findLineWithPrefix "SUB " - pubLine <- findLineWithPrefix "PUB " + pubLine <- findPublishLine let subParts = C.words subLine pubParts = C.words pubLine inbox <- subParts `at` 1 sid <- subParts `at` 2 - subject' <- pubParts `at` 1 - reply <- pubParts `at` 2 - lenBytes <- pubParts `at` 3 - len <- readMaybeInt lenBytes - payloadLine <- payloadAfter pubLine - let payloadBytes = BS.take len payloadLine + (subject', reply, headerLength, totalLength) <- publishLengths pubParts + frameBody <- frameBodyAfter pubLine + let payloadBytes = + BS.take (totalLength - headerLength) (BS.drop headerLength frameBody) pure CapturedPublish { capturedInbox = inbox , capturedSid = sid @@ -243,10 +262,29 @@ case filter (BS.isPrefixOf prefix) normalizedLines of [] -> Nothing (line:_) -> Just line - payloadAfter pubLine = - case dropWhile (/= pubLine) normalizedLines of - (_:payloadLine:_) -> Just payloadLine - _ -> Nothing + findPublishLine = + case filter isPublishLine normalizedLines of + [] -> Nothing + (line:_) -> Just line + isPublishLine line = + "PUB " `BS.isPrefixOf` line || "HPUB " `BS.isPrefixOf` line + publishLengths parts = + case parts of + ["PUB", subject', reply, lengthBytes] -> do + totalLength <- readMaybeInt lengthBytes + pure (subject', reply, 0, totalLength) + ["HPUB", subject', reply, headerLengthBytes, totalLengthBytes] -> do + headerLength <- readMaybeInt headerLengthBytes + totalLength <- readMaybeInt totalLengthBytes + pure (subject', reply, headerLength, totalLength) + _ -> + Nothing + frameBodyAfter line = + let marker = line <> "\r\n" + (_, suffix) = C.breakSubstring marker bytes + in if BS.null suffix + then Nothing + else Just (BS.drop (BS.length marker) suffix) at :: [a] -> Int -> Maybe a at values index = @@ -355,11 +393,11 @@ , withConnectionAttempts 1 , withConnectName "test-client" ] - c <- newClient [("127.0.0.1", p)] configOptions + c <- newClientOrFail [("127.0.0.1", p)] configOptions atomically $ putTMVar tvb c s <- atomically $ takeTMVar tva - sendAll s defaultINFO + void (completeHandshake s) c <- atomically $ takeTMVar tvb return (s, c, sock, exited) @@ -369,13 +407,13 @@ stopClient :: (Socket, Client, Socket, TMVar ClientExitReason) -> IO () stopClient (s, c, sock, _) = do - close c + close c [] Network.Socket.close sock Network.Socket.close s stopClientSafely :: (Socket, Client, Socket, TMVar ClientExitReason) -> IO () stopClientSafely (s, c, sock, _) = do - void (try (close c) :: IO (Either SomeException ())) + void (try (close c []) :: IO (Either SomeException ())) void (try (Network.Socket.close sock) :: IO (Either SomeException ())) void (try (Network.Socket.close s) :: IO (Either SomeException ())) @@ -390,52 +428,28 @@ spec = do describe "client integration" $ do around withClient $ do - it "PING waits for PONG" $ \(serv, client, _, _) -> do - wg <- newWaitGroup 1 - ping client $ done wg + it "PING resolves after PONG" $ \(serv, client, _, _) -> do + resultVar <- newEmptyMVar + void . forkIO $ ping client [] >>= putMVar resultVar + expectClientCommand serv "PING\r\n" sendAll serv "PONG\r\n" - wait wg - it "flush waits for PONG" $ \(serv, client, _, _) -> do - doneVar <- newEmptyMVar - void . forkIO $ do - flush client - putMVar doneVar () - threadDelay 100000 - returnedEarly <- not <$> isEmptyMVar doneVar - when returnedEarly $ - expectationFailure "flush returned before PONG" + expectMVar "PING did not resolve after PONG" resultVar + `shouldReturn` Right () + it "flush resolves after PONG" $ \(serv, client, _, _) -> do + resultVar <- newEmptyMVar + void . forkIO $ flush client [] >>= putMVar resultVar + expectClientCommand serv "PING\r\n" sendAll serv "PONG\r\n" - result <- timeout 1000000 (takeMVar doneVar) - case result of - Nothing -> expectationFailure "flush did not return after PONG" - Just () -> pure () + expectMVar "flush did not resolve after PONG" resultVar + `shouldReturn` Right () it "subscribes with a queue group" $ \(serv, client, _, _) -> do - sid <- subscribe client "JOBS" [withQueueGroup "WORKERS"] (const (pure ())) - expectQueuedSub serv "JOBS" "WORKERS" sid + Right subscription <- subscribe client "JOBS" [withQueueGroup "WORKERS"] (const (pure ())) + expectQueuedSub serv "JOBS" "WORKERS" (subscriptionSid subscription) it "publishes with an explicit reply subject" $ \(serv, client, _, _) -> do - publish client "REQUESTS" [withPayload "hello", withReplyTo "_INBOX.custom"] + publish client "REQUESTS" "hello" [withReplyTo "_INBOX.custom"] `shouldReturn` Right () expectClientCommand serv "PUB REQUESTS _INBOX.custom 5\r\nhello\r\n" - it "PONG resolves one ping" $ \(serv, client, _, _) -> do - first <- newEmptyMVar - second <- newEmptyMVar - ping client (putMVar first ()) - ping client (putMVar second ()) - sendAll serv "PONG\r\n" - firstResult <- timeout 1000000 (takeMVar first) - case firstResult of - Nothing -> expectationFailure "first ping did not resolve" - Just () -> pure () - threadDelay 100000 - secondReady <- not <$> isEmptyMVar second - when secondReady $ - expectationFailure "second ping resolved before second PONG" - sendAll serv "PONG\r\n" - secondResult <- timeout 1000000 (takeMVar second) - case secondResult of - Nothing -> expectationFailure "second ping did not resolve" - Just () -> pure () it "reports user initiated close" $ \(_, client, _, exited) -> do - close client + close client [] result <- atomically $ readTMVar exited result `shouldBe` ExitClosedByUser it "fatal error results in disconnect" $ \(serv, _, _, exited) -> do @@ -447,34 +461,37 @@ it "non fatal error does not result in disconnect" $ \(serv, client, _, _) -> do sendAll serv "-ERR 'Invalid Subject'\r\n" wg <- newWaitGroup 1 - ping client $ done wg + pingWithCallback client $ done wg + expectClientCommand serv "PING\r\n" sendAll serv "PONG\r\n" wait wg it "garbled prefix bytes are ignored" $ \(serv, client, _, _) -> do wg <- newWaitGroup 1 - ping client $ done wg + pingWithCallback client $ done wg + expectClientCommand serv "PING\r\n" sendAll serv "ldkfjajhfklsjhlkajf;alwfPONG\r\n" wait wg it "garbled suffix bytes remove partial prefix" $ \(serv, client, _, _) -> do wg <- newWaitGroup 1 - ping client $ done wg + pingWithCallback client $ done wg + expectClientCommand serv "PING\r\n" sendAll serv "MSGX" sendAll serv "PONG\r\n" wait wg it "messages split over frames are joined" $ \(serv, client, _, _) -> do wg <- newWaitGroup 1 - ping client $ done wg + pingWithCallback client $ done wg + expectClientCommand serv "PING\r\n" sendAll serv "PON" - threadDelay 100000 sendAll serv "G\r\n" wait wg it "MSG subject split over frames is joined" $ \(serv, client, _, _) -> do msgVar <- newEmptyMVar let topic = "SOAK.SUBJECT" payloadValue = "HELLO" - sid <- subscribe client topic [] (putMVar msgVar) + Right subscription <- subscribe client topic [] (putMVar msgVar) + let sid = subscriptionSid subscription sendAll serv "MSG SOAK." - threadDelay 100000 let headerTail = "SUBJECT " <> sid @@ -486,56 +503,55 @@ result <- timeout 1000000 (takeMVar msgVar) case result of Nothing -> expectationFailure "message not received" - Just Nothing -> expectationFailure "received empty message" - Just (Just msg) -> do + Just msg -> do subject msg `shouldBe` topic - payload msg `shouldBe` Just payloadValue + payload msg `shouldBe` payloadValue it "exits when server goes away" $ \(serv, _, _, exited) -> do Network.Socket.close serv result <- atomically $ readTMVar exited case result of ExitRetriesExhausted _ -> return () other -> expectationFailure $ "Unexpected exit reason: " ++ show other - it "drops messages too long for processing" $ \(serv, client, _, _) -> do + it "processes messages larger than the old parser buffer" $ \(serv, client, _, _) -> do wg <- newWaitGroup 1 - ping client $ done wg + pingWithCallback client $ done wg + expectClientCommand serv "PING\r\n" sendAll serv tooLongMSG sendAll serv "PONG\r\n" wait wg - it "unsubscribes after timeout" $ \(_, client, _, _) -> do - wg <- newWaitGroup 1 - _ <- request client "foo" [withSubscriptionExpiry 1] (\x -> do - case x of - Nothing -> done wg - Just _ -> error "should not receive message" - ) - wait wg - it "callback is called when expired" $ \(_, client, _, _) -> do - wg <- newWaitGroup 1 - _ <- request client "foo" [withSubscriptionExpiry 1] (\x -> do - case x of - Nothing -> done wg - Just _ -> error "should not receive message" - ) - close client - wait wg + it "enforces sub-second request timeouts" $ \(_, client, _, _) -> do + result <- timeout 500000 $ + request client "foo" "" [withRequestTimeout 0.05] + result `shouldBe` Just (Left NatsRequestTimedOut) + it "timed out requests do not block close" $ \(_, client, _, _) -> do + request client "foo" "" [withRequestTimeout 0.01] + `shouldReturn` Left NatsRequestTimedOut + timeout 500000 (close client []) `shouldReturn` Just () describe "core NATS fake-server coverage" $ do around withClient $ do it "ignores duplicate request replies after the first message" $ \(serv, client, _, _) -> do - replies <- newIORef ([] :: [Maybe MsgView]) - firstReply <- newEmptyMVar - sid <- request client "SERVICE.ONCE" [] $ \reply -> do + replies <- newIORef ([] :: [Message]) + Right subscription <- subscribeOnce client "SERVICE.ONCE" [] $ \reply -> do atomicModifyIORef' replies $ \seen -> (seen ++ [reply], ()) - void (tryPutMVar firstReply ()) - sendAll serv (msgFrame "SERVICE.ONCE" sid "first") - sendAll serv (msgFrame "SERVICE.ONCE" sid "second") - expectMVar "first request reply did not arrive" firstReply - threadDelay 100000 + let sid = subscriptionSid subscription + expectClientCommand serv (BS.concat ["SUB SERVICE.ONCE ", sid, "\r\n"]) + callbacksDrained <- newEmptyMVar + Right barrier <- subscribe client "CALLBACK.BARRIER" [] $ \_ -> + putMVar callbacksDrained () + let barrierSid = subscriptionSid barrier + expectClientCommand serv (BS.concat ["SUB CALLBACK.BARRIER ", barrierSid, "\r\n"]) + sendAll serv $ + BS.concat + [ msgFrame "SERVICE.ONCE" sid "first" + , msgFrame "SERVICE.ONCE" sid "second" + , msgFrame "CALLBACK.BARRIER" barrierSid "done" + ] + expectMVar "callback barrier did not run" callbacksDrained seen <- readIORef replies case seen of - [Just msg] -> - payload msg `shouldBe` Just "first" + [msg] -> + payload msg `shouldBe` "first" other -> expectationFailure ("unexpected request replies: " ++ show other) @@ -543,79 +559,58 @@ requests <- forM [1 .. 32 :: Int] $ \index -> do replyBox <- newEmptyMVar let requestSubject = C.pack ("SERVICE.MANY." ++ show index) - sid <- request client requestSubject [] (putMVar replyBox) + Right subscription <- subscribeOnce client requestSubject [] (putMVar replyBox) + let sid = subscriptionSid subscription pure (index, requestSubject, sid, replyBox) for_ requests $ \(index, requestSubject, sid, _) -> sendAll serv $ msgFrame requestSubject sid (C.pack ("reply-" ++ show index)) for_ requests $ \(index, _, _, replyBox) -> do reply <- expectMVar ("request reply " ++ show index ++ " did not arrive") replyBox - case reply of - Nothing -> - expectationFailure ("request reply " ++ show index ++ " expired") - Just msg -> - payload msg `shouldBe` Just (C.pack ("reply-" ++ show index)) + payload reply `shouldBe` C.pack ("reply-" ++ show index) it "keeps an active request alive after garbled inbound bytes" $ \(serv, client, _, _) -> do replyBox <- newEmptyMVar - sid <- request client "SERVICE.GARBLED" [] (putMVar replyBox) + Right subscription <- subscribeOnce client "SERVICE.GARBLED" [] (putMVar replyBox) + let sid = subscriptionSid subscription sendAll serv $ BS.concat [ "bad-prefix" , msgFrame "SERVICE.GARBLED" sid "ok" ] reply <- expectMVar "request reply after garbled bytes did not arrive" replyBox - case reply of - Nothing -> - expectationFailure "request expired after garbled bytes" - Just msg -> - payload msg `shouldBe` Just "ok" + payload reply `shouldBe` "ok" it "round trips headers across request publish and reply" $ \(serv, client, _, _) -> do replyBox <- newEmptyMVar - publish client - "HEADERS.SERVICE" - [ withHeaders [("X-Request", "one")] - , withPayload "ping" - , withReplyTo "_INBOX.headers" - , withReplyCallback (putMVar replyBox) - ] - bytes <- expectClientBytes - serv - (\sent -> - BS.isInfixOf "SUB _INBOX.headers " sent - && BS.isInfixOf "HPUB HEADERS.SERVICE _INBOX.headers " sent - && BS.isInfixOf "X-Request:one\r\n" sent - && BS.isInfixOf "\r\n\r\nping\r\n" sent) - "client did not send header request publish" - captured <- case parseCapturedSubscription "_INBOX.headers" bytes of - Nothing -> - expectationFailure ("could not parse header reply subscription: " ++ show bytes) >> - fail "could not parse header reply subscription" - Just sub -> - pure sub + void . forkIO $ + request client + "HEADERS.SERVICE" + "ping" + [withRequestHeaders [("X-Request", "one")]] + >>= putMVar replyBox + captured <- capturePublish serv "HEADERS.SERVICE" sendAll serv $ hmsgFrame - "_INBOX.headers" - (capturedSubSid captured) + (capturedInbox captured) + (capturedSid captured) [("X-Reply", "ok")] "pong" reply <- expectMVar "header reply did not arrive" replyBox case reply of - Nothing -> - expectationFailure "header reply expired" - Just msg -> do - subject msg `shouldBe` "_INBOX.headers" - payload msg `shouldBe` Just "pong" + Left err -> + expectationFailure ("header request failed: " ++ show err) + Right msg -> do + subject msg `shouldBe` capturedInbox captured + payload msg `shouldBe` "pong" headers msg `shouldBe` Just [("X-Reply", "ok")] it "flush sends queued publishes before its PING" $ \(serv, client, _, _) -> do - expectClientCommand serv "CONNECT " doneVar <- newEmptyMVar let publishFrame = "PUB ORDER.FLUSH 3\r\none\r\n" - publish client "ORDER.FLUSH" [withPayload "one"] + publish client "ORDER.FLUSH" "one" [] `shouldReturn` Right () void . forkIO $ do - flush client + flush client [] putMVar doneVar () bytes <- expectClientBytes serv @@ -629,7 +624,8 @@ expectMVar "flush did not complete after PONG" doneVar it "exits when the server disconnects while a subscription is active" $ \(serv, client, _, exited) -> do - sid <- subscribe client "ACTIVE.DISCONNECT" [] (const (pure ())) + Right subscription <- subscribe client "ACTIVE.DISCONNECT" [] (const (pure ())) + let sid = subscriptionSid subscription expectClientCommand serv (BS.concat ["SUB ACTIVE.DISCONNECT ", sid, "\r\n"]) Network.Socket.close serv result <- atomically $ readTMVar exited @@ -640,51 +636,115 @@ expectationFailure ("Unexpected exit reason: " ++ show other) it "does not write publishes above negotiated max_payload" $ \(serv, client, _, _) -> do - expectClientCommand serv "CONNECT " - publish client "BOUNDARY.TOO_BIG" [withPayload (BS.replicate 2048 _x)] - expectNoClientBytes + publish client "BOUNDARY.TOO_BIG" (BS.replicate 2048 _x) [] + `shouldReturn` Left (NatsPayloadTooLarge 2048 1024) + wg <- newWaitGroup 1 + pingWithCallback client (done wg) + expectNoClientBytesBefore serv (BS.isInfixOf "PUB BOUNDARY.TOO_BIG") + "PING\r\n" "oversized publish reached server" - wg <- newWaitGroup 1 - ping client (done wg) - expectClientCommand serv "PING\r\n" sendAll serv "PONG\r\n" wait wg it "does not write invalid publish commands" $ \(serv, client, _, _) -> do - expectClientCommand serv "CONNECT " - publish client "" [withPayload "x"] - expectNoClientBytes + publish client "" "x" [] >>= (`shouldSatisfy` isValidationFailure) + pingBox <- newEmptyMVar + pingWithCallback client (putMVar pingBox ()) + expectNoClientBytesBefore serv (\sent -> BS.isInfixOf "PUB " sent || BS.isInfixOf "HPUB " sent) + "PING\r\n" "invalid publish reached server" + sendAll serv "PONG\r\n" + expectMVar "PING failed after an invalid publish" pingBox + around (withClientWith [withMessageLimit 4]) $ do + it "does not write publishes above the client message limit" $ \(serv, client, _, _) -> do + publish client "BOUNDARY.CLIENT" "12345" [] + `shouldReturn` Left (NatsPayloadTooLarge 5 4) + pingBox <- newEmptyMVar + pingWithCallback client (putMVar pingBox ()) + expectNoClientBytesBefore + serv + (BS.isInfixOf "PUB BOUNDARY.CLIENT") + "PING\r\n" + "publish above the client limit reached server" + sendAll serv "PONG\r\n" + expectMVar "PING failed after a rejected publish" pingBox + + it "accepts a publish exactly at the client message limit" $ \(serv, client, _, _) -> do + publish client "BOUNDARY.EXACT" "1234" [] `shouldReturn` Right () + expectClientCommand serv "PUB BOUNDARY.EXACT 4\r\n1234\r\n" + + it "counts headers toward the client message limit" $ \(serv, client, _, _) -> do + let messageHeaders = [("X", "Y")] + actualSize = BS.length (headerBlock messageHeaders) + publish client "BOUNDARY.HEADERS" "" [withHeaders messageHeaders] + `shouldReturn` Left (NatsPayloadTooLarge actualSize 4) + pingBox <- newEmptyMVar + pingWithCallback client (putMVar pingBox ()) + expectNoClientBytesBefore + serv + (BS.isInfixOf "HPUB BOUNDARY.HEADERS") + "PING\r\n" + "publish whose headers exceed the client limit reached server" + sendAll serv "PONG\r\n" + expectMVar "PING failed after a rejected header publish" pingBox + + it "terminates the connection when an inbound message exceeds the client limit" $ \(serv, client, _, exited) -> do + Right subscription <- subscribe client "BOUNDARY.INBOUND" [] (const (pure ())) + let sid = subscriptionSid subscription + expectClientCommand serv (BS.concat ["SUB BOUNDARY.INBOUND ", sid, "\r\n"]) + sendAll serv (msgFrame "BOUNDARY.INBOUND" sid "12345") + atomically (readTMVar exited) + `shouldReturn` ExitInboundMessageTooLarge 5 4 + + around (withClientWith [withPendingDeliveryLimits 1 1024]) $ do + it "fails a request promptly when the global callback backlog is full" $ \(serv, client, _, _) -> do + callbackStarted <- newEmptyMVar + releaseCallback <- newEmptyMVar + Right subscription <- subscribe client "SLOW.CALLBACK" [] $ \_ -> do + putMVar callbackStarted () + takeMVar releaseCallback + let sid = subscriptionSid subscription + expectClientCommand serv (BS.concat ["SUB SLOW.CALLBACK ", sid, "\r\n"]) + sendAll serv (msgFrame "SLOW.CALLBACK" sid "busy") + expectMVar "blocking callback did not start" callbackStarted + + replyBox <- newEmptyMVar + void . forkIO $ + request client "SERVICE.OVERLOADED" "request" [withRequestTimeout 2] + >>= putMVar replyBox + captured <- capturePublish serv "SERVICE.OVERLOADED" + sendAll serv $ + msgFrame (capturedInbox captured) (capturedSid captured) "reply" + expectMVar "overloaded request did not fail promptly" replyBox + `shouldReturn` Left NatsSlowConsumer + + sendAll serv "PING\r\n" + expectClientCommand serv "PONG\r\n" + putMVar releaseCallback () + it "cleans up no responders replies for expiring requests" $ do bracket startClient stopClientSafely $ \(serv, client, _, _) -> do replyBox <- newEmptyMVar - sid <- request client - "SERVICE.NO_RESPONDERS" - [withSubscriptionExpiry 2] - (putMVar replyBox) + void . forkIO $ + request client "SERVICE.NO_RESPONDERS" "" [withRequestTimeout 2] + >>= putMVar replyBox + captured <- capturePublish serv "SERVICE.NO_RESPONDERS" sendAll serv $ hmsgFrame - "SERVICE.NO_RESPONDERS" - sid + (capturedInbox captured) + (capturedSid captured) [ ("Status", "503") , ("Description", "No Responders") ] "" - reply <- expectMVar "no responders reply callback did not run" replyBox - case reply of - Nothing -> - expectationFailure "no responders reply expired" - Just msg -> - headers msg `shouldBe` Just - [ ("Status", "503") - , ("Description", "No Responders") - ] - result <- timeout 500000 (close client) + reply <- expectMVar "no responders request did not return" replyBox + reply `shouldBe` Left NatsNoResponders + result <- timeout 500000 (close client []) case result of Nothing -> expectationFailure "close blocked after no responders reply" @@ -701,46 +761,39 @@ , withConnectionAttempts 2 , withConnectName "test-client" ] - c <- newClient [("127.0.0.1", p)] configOptions + c <- newClientOrFail [("127.0.0.1", p)] configOptions putMVar clientVar c (firstConn, _) <- accept sock - sendAll firstConn defaultINFO + void (completeHandshake firstConn) clientResult <- timeout 1000000 (takeMVar clientVar) case clientResult of Nothing -> expectationFailure "client did not connect" Just client -> do - replyBox <- newEmptyMVar - publish client - "SERVICE.RECONNECT" - [ withPayload "ping" - , withReplyCallback (putMVar replyBox) - ] + void . forkIO $ + void (request client "SERVICE.RECONNECT" "ping" [withRequestTimeout 5]) captured <- capturePublish firstConn "SERVICE.RECONNECT" Network.Socket.close firstConn (secondConn, _) <- accept sock - sendAll secondConn defaultINFO + void (completeHandshake secondConn) let replySubCommand = BS.concat ["SUB ", capturedInbox captured] - resubscribe <- timeout 300000 $ - recvUntil secondConn (BS.isInfixOf replySubCommand) - case resubscribe of - Just bytes | BS.isInfixOf replySubCommand bytes -> - expectationFailure ("reply inbox was resubscribed: " ++ show bytes) - _ -> - pure () - wg <- newWaitGroup 1 - ping client (done wg) - expectClientCommand secondConn "PING\r\n" + pingBox <- newEmptyMVar + pingWithCallback client (putMVar pingBox ()) + expectNoClientBytesBefore + secondConn + (BS.isInfixOf replySubCommand) + "PING\r\n" + "reply inbox was resubscribed" sendAll secondConn "PONG\r\n" - wait wg - close client + expectMVar "PING failed after reconnect" pingBox + close client [] Network.Socket.close secondConn Network.Socket.close sock describe "jetstream protocol integration" $ do around withClient $ do it "sends stream create options as a JetStream API request" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do result <- JetStream.create (JetStream.streams jetStream) @@ -750,6 +803,7 @@ , JetStream.withMaxMessages 1 , JetStream.withDiscard JetStream.DiscardNew ] + [] putMVar done result captured <- capturePublish serv "$JS.API.STREAM.CREATE.PROTO_STREAM" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"max_msgs\":1" @@ -765,7 +819,7 @@ JetStream.streamConfigName (JetStream.streamInfoConfig info) `shouldBe` "PROTO_STREAM" it "sends only the selected consumer filter representation" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do result <- JetStream.putConsumer @@ -779,6 +833,7 @@ , JetStream.withConsumerFilter (JetStream.ConsumerFilterSubjects ["PROTO.A", "PROTO.B"]) ] + [] putMVar done result captured <- capturePublish serv "$JS.API.CONSUMER.DURABLE.CREATE.PROTO_STREAM.PROTO_CONSUMER" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"filter_subjects\":[\"PROTO.A\",\"PROTO.B\"]" @@ -794,7 +849,7 @@ JetStream.consumerInfoName info `shouldBe` "PROTO_CONSUMER" it "sends push consumer create requests with a delivery subject" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do result <- JetStream.putConsumer @@ -806,6 +861,7 @@ [ JetStream.withConsumerDeliverGroup "workers" , JetStream.withConsumerAckPolicy JetStream.AckExplicit ] + [] putMVar done result captured <- capturePublish serv "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_PUSH" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"action\":\"create\"" @@ -824,23 +880,25 @@ `shouldBe` Just "PROTO.DELIVER" it "subscribes to push consumer delivery subjects with a queue group" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] - subscription <- JetStream.consumePush + let Right jetStream = newJetStream client [] + Right subscription <- JetStream.consumePush (JetStream.messages jetStream) "PROTO.DELIVER" [JetStream.withPushQueueGroup "workers"] + [] (const (pure ())) expectClientCommand serv "SUB PROTO.DELIVER workers " - JetStream.stopPushSubscription subscription + void (JetStream.stopPushSubscription subscription) expectClientCommand serv "UNSUB " it "keeps push consumers alive after garbled inbound bytes" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] delivered <- newEmptyMVar - subscription <- JetStream.consumePush + Right subscription <- JetStream.consumePush (JetStream.messages jetStream) "PROTO.DELIVER" [] + [] (putMVar delivered) captured <- captureSubscription serv "PROTO.DELIVER" sendAll serv $ @@ -849,15 +907,15 @@ , msgFrame "PROTO.DELIVER" (capturedSubSid captured) "payload" ] result <- timeout 1000000 (takeMVar delivered) - JetStream.stopPushSubscription subscription + void (JetStream.stopPushSubscription subscription) case result of Nothing -> expectationFailure "push message was not delivered after garbled bytes" Just message -> - JetStream.messagePayload message `shouldBe` Just "payload" + JetStream.messagePayload message `shouldBe` "payload" it "creates ordered consumers with the ordered defaults" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do result <- JetStream.createOrderedConsumer @@ -866,6 +924,7 @@ [ JetStream.withOrderedConsumerNamePrefix "PROTO_ORDERED" , JetStream.withOrderedConsumerDeliverPolicy (JetStream.DeliverByStartSequence 3) ] + [] putMVar done result captured <- capturePublish serv "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_ORDERED_1" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"name\":\"PROTO_ORDERED_1\"" @@ -886,7 +945,7 @@ pure () it "sends no-wait pull requests and decodes no-message status" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do response <- JetStream.fetch (JetStream.messages jetStream) @@ -895,6 +954,7 @@ [ JetStream.withFetchBatch 1 , JetStream.withFetchWait (JetStream.FetchNoWaitMicros 1000000) ] + [] putMVar done response captured <- capturePublish serv "$JS.API.CONSUMER.MSG.NEXT.PROTO_STREAM.PROTO_CONSUMER" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"no_wait\":true" @@ -913,10 +973,10 @@ `shouldBe` Just (JetStream.PullNoMessages (Just "No Messages")) it "sends account info requests" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do - result <- JetStream.accountInfo jetStream + result <- JetStream.accountInfo (JetStream.management jetStream) [] putMVar done result captured <- capturePublish serv "$JS.API.INFO" capturedPayload captured `shouldBe` "" @@ -931,10 +991,10 @@ JetStream.accountTierStreams (JetStream.accountInfoTier info) `shouldBe` 1 it "returns a JetStream timeout when admin requests receive no response" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [withRequestTimeoutMicros 1000] + let Right jetStream = newJetStream client [withRequestTimeoutMicros 1000] done <- newEmptyMVar void . forkIO $ do - result <- JetStream.accountInfo jetStream + result <- JetStream.accountInfo (JetStream.management jetStream) [] putMVar done result captured <- capturePublish serv "$JS.API.INFO" capturedPayload captured `shouldBe` "" @@ -948,13 +1008,14 @@ expectationFailure ("unexpected account info timeout result: " ++ show other) it "sends stream message get and delete requests" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] fetched <- newEmptyMVar void . forkIO $ do result <- JetStream.getMessage (JetStream.streams jetStream) "PROTO_STREAM" (JetStream.StreamMessageBySequence 2) + [] putMVar fetched result getRequest <- capturePublish serv "$JS.API.STREAM.MSG.GET.PROTO_STREAM" capturedPayload getRequest `shouldSatisfy` BS.isInfixOf "\"seq\":2" @@ -975,6 +1036,7 @@ "PROTO_STREAM" 2 JetStream.DeleteMessage + [] putMVar deleted result deleteRequest <- capturePublish serv "$JS.API.STREAM.MSG.DELETE.PROTO_STREAM" capturedPayload deleteRequest `shouldSatisfy` BS.isInfixOf "\"seq\":2" @@ -990,7 +1052,7 @@ JetStream.deleteStreamMessageSuccess response `shouldBe` True it "sends consumer update actions" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] done <- newEmptyMVar void . forkIO $ do result <- JetStream.putConsumer @@ -1002,6 +1064,7 @@ [ JetStream.withConsumerDescription "updated" , JetStream.withConsumerAckPolicy JetStream.AckExplicit ] + [] putMVar done result captured <- capturePublish serv "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_CONSUMER" capturedPayload captured `shouldSatisfy` BS.isInfixOf "\"action\":\"update\"" @@ -1017,7 +1080,7 @@ JetStream.consumerInfoName info `shouldBe` "PROTO_CONSUMER" it "sends consumer pause and reset requests" $ \(serv, client, _, _) -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] pauseUntil = read "2024-01-01 00:00:01 UTC" paused <- newEmptyMVar void . forkIO $ do @@ -1026,6 +1089,7 @@ "PROTO_STREAM" "PROTO_CONSUMER" pauseUntil + [] putMVar paused result pauseRequest <- capturePublish serv "$JS.API.CONSUMER.PAUSE.PROTO_STREAM.PROTO_CONSUMER" capturedPayload pauseRequest `shouldSatisfy` BS.isInfixOf "\"pause_until\":\"2024-01-01T00:00:01Z\"" @@ -1046,6 +1110,7 @@ "PROTO_STREAM" "PROTO_CONSUMER" [JetStream.withConsumerResetSequence 7] + [] putMVar reset result resetRequest <- capturePublish serv "$JS.API.CONSUMER.RESET.PROTO_STREAM.PROTO_CONSUMER" capturedPayload resetRequest `shouldSatisfy` BS.isInfixOf "\"seq\":7" @@ -1060,26 +1125,60 @@ JetStream.consumerResetResponseSequence response `shouldBe` 7 it "manual unsubscribe does not block close for tracked expiries" $ do bracket startClient stopClientSafely $ \(_, client, _, _) -> do - sid <- request client "foo" [withSubscriptionExpiry 30] (const (pure ())) - unsubscribe client sid - result <- timeout 1000000 (close client) + Right subscription <- subscribeOnce client "foo" [withSubscriptionExpiry 30] (const (pure ())) + unsubscribe client subscription [] `shouldReturn` Right () + result <- timeout 1000000 (close client []) case result of Nothing -> expectationFailure "close blocked after unsubscribe" Just () -> pure () it "early replies do not block close for tracked expiries" $ do bracket startClient stopClientSafely $ \(serv, client, _, _) -> do replyBox <- newEmptyMVar - sid <- request client "foo" [withSubscriptionExpiry 30] (putMVar replyBox) + Right subscription <- subscribeOnce client "foo" [withSubscriptionExpiry 30] (putMVar replyBox) + let sid = subscriptionSid subscription sendAll serv (msgFrame "foo" sid "bar") reply <- timeout 1000000 (takeMVar replyBox) case reply of - Nothing -> expectationFailure "reply callback did not run" - Just Nothing -> expectationFailure "expected a reply message" - Just (Just msg) -> payload msg `shouldBe` Just "bar" - result <- timeout 1000000 (close client) + Nothing -> expectationFailure "reply callback did not run" + Just msg -> payload msg `shouldBe` "bar" + result <- timeout 1000000 (close client []) case result of Nothing -> expectationFailure "close blocked after reply" Just () -> pure () + it "refreshes auth handlers after a rejected connection attempt" $ do + (p, sock) <- openFreePort + listen sock 2 + captures <- newEmptyTMVarIO + secondConnection <- newEmptyTMVarIO + void . forkIO $ do + (firstConn, _) <- accept sock + sendAll firstConn defaultINFO + firstConnect <- recv firstConn 4096 + sendAll firstConn "-ERR 'Authorization Violation'\r\n" + (secondConn, _) <- accept sock + secondConnect <- completeHandshake secondConn + atomically $ do + putTMVar captures (firstConnect, secondConnect) + putTMVar secondConnection secondConn + handlerCalls <- newIORef (0 :: Int) + let tokenHandler = do + call <- atomicModifyIORef' handlerCalls $ \current -> + let next = current + 1 + in (next, next) + pure (Right ("token-" <> C.pack (show call))) + configOptions = + [ withAuthTokenHandler tokenHandler + , withConnectionAttempts 2 + , withConnectName "rotating-auth-client" + ] + client <- newClientOrFail [("127.0.0.1", p)] configOptions + (firstConnect, secondConnect) <- atomically $ takeTMVar captures + firstConnect `shouldSatisfy` BS.isInfixOf "\"token-1\"" + secondConnect `shouldSatisfy` BS.isInfixOf "\"token-2\"" + close client [] + Network.Socket.close =<< atomically (takeTMVar secondConnection) + Network.Socket.close sock + it "retries when the server disconnects before INFO" $ do (p, sock) <- openFreePort listen sock 2 @@ -1089,7 +1188,7 @@ (firstConn, _) <- accept sock Network.Socket.close firstConn (secondConn, _) <- accept sock - sendAll secondConn defaultINFO + void (completeHandshake secondConn) atomically $ putTMVar serverConn secondConn void . forkIO $ do let configOptions = @@ -1097,7 +1196,7 @@ , withConnectionAttempts 2 , withConnectName "test-client" ] - c <- newClient [("127.0.0.1", p)] configOptions + c <- newClientOrFail [("127.0.0.1", p)] configOptions putMVar clientVar c clientResult <- timeout 1000000 (takeMVar clientVar) case clientResult of @@ -1105,10 +1204,11 @@ Just client -> do secondConn <- atomically $ takeTMVar serverConn wg <- newWaitGroup 1 - ping client $ done wg + pingWithCallback client $ done wg + expectClientCommand secondConn "PING\r\n" sendAll secondConn "PONG\r\n" wait wg - close client + close client [] Network.Socket.close secondConn Network.Socket.close sock it "resubscribes with a queue group after reconnect" $ do @@ -1121,21 +1221,22 @@ , withConnectionAttempts 2 , withConnectName "test-client" ] - c <- newClient [("127.0.0.1", p)] configOptions + c <- newClientOrFail [("127.0.0.1", p)] configOptions putMVar clientVar c (firstConn, _) <- accept sock - sendAll firstConn defaultINFO + void (completeHandshake firstConn) clientResult <- timeout 1000000 (takeMVar clientVar) case clientResult of Nothing -> expectationFailure "client did not connect" Just client -> do - sid <- subscribe client "JOBS" [withQueueGroup "WORKERS"] (const (pure ())) + Right subscription <- subscribe client "JOBS" [withQueueGroup "WORKERS"] (const (pure ())) + let sid = subscriptionSid subscription expectQueuedSub firstConn "JOBS" "WORKERS" sid Network.Socket.close firstConn (secondConn, _) <- accept sock - sendAll secondConn defaultINFO + void (completeHandshake secondConn) expectQueuedSub secondConn "JOBS" "WORKERS" sid - close client + close client [] Network.Socket.close secondConn Network.Socket.close sock it "resubscribes JetStream push consumers after reconnect" $ do @@ -1148,26 +1249,27 @@ , withConnectionAttempts 2 , withConnectName "test-client" ] - c <- newClient [("127.0.0.1", p)] configOptions + c <- newClientOrFail [("127.0.0.1", p)] configOptions putMVar clientVar c (firstConn, _) <- accept sock - sendAll firstConn defaultINFO + void (completeHandshake firstConn) clientResult <- timeout 1000000 (takeMVar clientVar) case clientResult of Nothing -> expectationFailure "client did not connect" Just client -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] delivered <- newEmptyMVar - subscription <- JetStream.consumePush + Right subscription <- JetStream.consumePush (JetStream.messages jetStream) "PROTO.DELIVER" [] + [] (putMVar delivered) firstSub <- captureSubscription firstConn "PROTO.DELIVER" Network.Socket.close firstConn (secondConn, _) <- accept sock - sendAll secondConn defaultINFO + void (completeHandshake secondConn) secondSub <- captureSubscription secondConn "PROTO.DELIVER" capturedSubSid secondSub `shouldBe` capturedSubSid firstSub sendAll secondConn $ @@ -1177,9 +1279,9 @@ Nothing -> expectationFailure "push message was not delivered after reconnect" Just message -> - JetStream.messagePayload message `shouldBe` Just "after-reconnect" - JetStream.stopPushSubscription subscription - close client + JetStream.messagePayload message `shouldBe` "after-reconnect" + void (JetStream.stopPushSubscription subscription) + close client [] Network.Socket.close secondConn Network.Socket.close sock it "ordered fetch returns after disconnect and recovers on reconnect" $ do @@ -1192,22 +1294,23 @@ , withConnectionAttempts 2 , withConnectName "test-client" ] - c <- newClient [("127.0.0.1", p)] configOptions + c <- newClientOrFail [("127.0.0.1", p)] configOptions putMVar clientVar c (firstConn, _) <- accept sock - sendAll firstConn defaultINFO + void (completeHandshake firstConn) clientResult <- timeout 1000000 (takeMVar clientVar) case clientResult of Nothing -> expectationFailure "client did not connect" Just client -> do - let jetStream = newJetStream client [] + let Right jetStream = newJetStream client [] orderedVar <- newEmptyMVar void . forkIO $ do result <- JetStream.createOrderedConsumer (JetStream.messages jetStream) "PROTO_STREAM" [JetStream.withOrderedConsumerNamePrefix "PROTO_ORDERED"] + [] putMVar orderedVar result createInitial <- capturePublish firstConn "$JS.API.CONSUMER.CREATE.PROTO_STREAM.PROTO_ORDERED_1" replyToCapturedPublish firstConn createInitial $ @@ -1229,6 +1332,7 @@ [ JetStream.withFetchBatch 1 , JetStream.withFetchWait (JetStream.FetchExpiresMicros 100000) ] + [] putMVar fetchVar result deleteCurrent <- capturePublish firstConn "$JS.API.CONSUMER.DELETE.PROTO_STREAM.PROTO_ORDERED_1" replyToCapturedPublish firstConn deleteCurrent protoDeleteConsumerResponse @@ -1247,10 +1351,10 @@ expectationFailure ("unexpected ordered fetch result after disconnect: " ++ show other) (secondConn, _) <- accept sock - sendAll secondConn defaultINFO + void (completeHandshake secondConn) infoVar <- newEmptyMVar void . forkIO $ do - result <- JetStream.orderedConsumerInfo ordered + result <- JetStream.orderedConsumerInfo ordered [] putMVar infoVar result infoRequest <- capturePublish secondConn "$JS.API.CONSUMER.INFO.PROTO_STREAM.PROTO_ORDERED_2" replyToCapturedPublish secondConn infoRequest $ @@ -1263,10 +1367,10 @@ expectationFailure ("ordered consumer info failed after reconnect: " ++ show err) Just (Right info) -> JetStream.consumerInfoName info `shouldBe` "PROTO_ORDERED_2" - close client + close client [] Network.Socket.close secondConn Network.Socket.close sock - around (withClientWith [withBufferLimit (64 * 1024), withCallbackConcurrency 4]) $ do + around (withClientWith [withMessageLimit (64 * 1024), withCallbackConcurrency 4]) $ do it "soak: parses a large buffer of MSG and HMSG frames" $ \(serv, client, _, _) -> do let subject = "SOAK.SUBJECT" smallPayloadSize = 512 @@ -1295,7 +1399,7 @@ Nothing -> (Just err, ()) Just _ -> (current, ()) let handleMsg msg = do - let payloadLen = maybe 0 BS.length (payload msg) + let payloadLen = BS.length (payload msg) case headers msg of Just hs -> do atomicModifyIORef' headerCounter $ \count -> (count + 1, ()) @@ -1314,7 +1418,8 @@ in (nextCount, nextCount) when (count == expectedTotal) $ void (tryPutMVar done ()) - sid <- subscribe client subject [] (maybe (pure ()) handleMsg) + Right subscription <- subscribe client subject [] handleMsg + let sid = subscriptionSid subscription let msg = msgFrame subject sid smallPayload let hmsg = hmsgFrame subject sid headerPairs largePayload let buffer = BS.concat (replicate smallCount msg ++ replicate largeCount hmsg) diff --git a/test/PublicAPI/Main.hs b/test/PublicAPI/Main.hs new file mode 100644 --- /dev/null +++ b/test/PublicAPI/Main.hs @@ -0,0 +1,162 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- This deliberately depends only on the public library. The declarations are +-- representative downstream calls; evaluating them is unnecessary because +-- compiling this executable is the compatibility check. +module Main (main) where + +import qualified API as Nats +import qualified Client +import qualified JetStream.API as JetStream +import qualified JetStream.API.Consumer as Consumer +import qualified JetStream.API.Management as Management +import qualified JetStream.API.Message as JetStreamMessage +import qualified JetStream.API.Publish as JetStreamPublish +import qualified JetStream.API.Stream as Stream +import qualified JetStream.Client as JetStreamClient + +connect :: IO (Either Client.ConnectError Nats.Client) +connect = + Client.connect + [configuredServer] + [ Client.withConnectName "public-api-compile-test" + , Client.withConnectionAttempts 1 + , Client.withMessageLimit (1024 * 1024) + , Client.withPendingDeliveryLimits 65536 (64 * 1024 * 1024) + , Client.withErrorHandler (const (pure ())) + ] + +configuredServer :: Client.Server +configuredServer = + case Client.server "127.0.0.1" 4222 of + Left err -> error (show err) + Right endpoint -> endpoint + +serverBuilders + :: ( Either Client.ServerConfigError Client.Server + , Either Client.ServerConfigError Client.Server + ) +serverBuilders = + (Client.server "nats.example" 4222, Client.serverWithDefaultPort "nats.example") + +coreOperations :: Nats.Client -> IO () +coreOperations client = do + _ <- Nats.publish + client + "events.created" + "payload" + [Nats.withHeaders [("content-type", "text/plain")]] + subscription <- Nats.subscribe + client + "events.*" + [Nats.withQueueGroup "workers"] + observeMessage + case subscription of + Left _ -> pure () + Right handle -> do + _ <- Nats.unsubscribe client handle [] + pure () + _ <- Nats.request + client + "service.echo" + "request" + [Nats.withRequestTimeout 1] + _ <- Nats.ping client [] + _ <- Nats.flush client [] + Nats.close client [] + +observeMessage :: Nats.Message -> IO () +observeMessage message = do + let _ = Nats.subject message + _ = Nats.payload message + _ = Nats.replyTo message + _ = Nats.headers message + pure () + +jetStreamContext + :: Nats.Client + -> Either JetStreamClient.JetStreamConfigError JetStream.JetStream +jetStreamContext client = + JetStreamClient.newJetStream + client + [ JetStreamClient.withDomain "hub" + , JetStreamClient.withRequestTimeoutMicros 1000000 + ] + +jetStreamOperations :: JetStream.JetStream -> IO () +jetStreamOperations jetStream = do + _ <- Management.accountInfo + (JetStream.management jetStream) + [JetStream.withRequestTimeout 1] + created <- Stream.createStream + (JetStream.streams jetStream) + "ORDERS" + ["orders.>"] + [ Stream.withStorage Stream.FileStorage + , Stream.withMaxMessages 1000 + ] + [] + case created of + Left _ -> pure () + Right handle -> do + _ <- Stream.getStreamInfo (JetStream.streams jetStream) handle [] + pure () + _ <- JetStreamPublish.publish + (JetStream.publisher jetStream) + "orders.created" + "order" + [JetStreamPublish.withMsgId "api-compile-test"] + [] + consumer <- Consumer.putConsumerHandle + (JetStream.consumers jetStream) + "ORDERS" + Consumer.ConsumerCreate + (Consumer.DurableConsumer "processor") + Consumer.PullConsumer + [Consumer.withConsumerAckPolicy Consumer.AckExplicit] + [] + case consumer of + Left _ -> pure () + Right handle -> do + _ <- Consumer.getConsumerInfo (JetStream.consumers jetStream) handle [] + pure () + _ <- JetStreamMessage.fetch + (JetStream.messages jetStream) + "ORDERS" + "processor" + [JetStreamMessage.withFetchBatch 10] + [] + _ <- JetStreamMessage.consumePush + (JetStream.messages jetStream) + "deliver.orders" + [] + [] + (const (pure ())) + pure () + +messageOperations + :: JetStream.MessageAPI + -> JetStreamMessage.Message + -> IO () +messageOperations messageAPI message = do + JetStreamMessage.messageMetadata message `seq` pure () + _ <- JetStreamMessage.ack messageAPI message [] + _ <- JetStreamMessage.nak messageAPI message [] + pure () + +inspectJetStreamApiError :: JetStream.JetStreamApiError -> IO () +inspectJetStreamApiError err = do + JetStream.apiErrorCode err `seq` pure () + JetStream.apiErrorCodeDetail err `seq` pure () + JetStream.apiErrorDescription err `seq` pure () + +main :: IO () +main = + connect + `seq` serverBuilders + `seq` coreOperations + `seq` jetStreamContext + `seq` jetStreamOperations + `seq` messageOperations + `seq` inspectJetStreamApiError + `seq` pure () diff --git a/test/Unit/ConnectSpec.hs b/test/Unit/ConnectSpec.hs --- a/test/Unit/ConnectSpec.hs +++ b/test/Unit/ConnectSpec.hs @@ -6,6 +6,7 @@ import Data.Aeson import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS +import Lib.CallOption (applyCallOptions) import Test.Hspec import Text.Printf import Transformers.Transformers @@ -14,6 +15,9 @@ spec :: Spec spec = do + describe "call options" $ do + it "applies options from left to right so the last override wins" $ do + applyCallOptions [(+ 1), (* 2)] (3 :: Int) `shouldBe` 8 transformationCases validationCases diff --git a/test/Unit/ConnectionSpec.hs b/test/Unit/ConnectionSpec.hs --- a/test/Unit/ConnectionSpec.hs +++ b/test/Unit/ConnectionSpec.hs @@ -3,12 +3,16 @@ module ConnectionSpec (spec) where import Auth.None (auth) +import qualified Client import Control.Concurrent import Control.Concurrent.STM import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS import Data.IORef (newIORef, readIORef, writeIORef) -import Handshake.Nats (performHandshake) +import Handshake.Nats + ( HandshakeError (HandshakeAuthError, HandshakeTLSError, HandshakeTimeout) + , performHandshake + ) import Lib.Logger ( LogLevel (Debug) , LoggerConfig (LoggerConfig) @@ -27,6 +31,7 @@ , ParsedMessage (ParsedInfo, ParsedPing, ParsedPong) , ParserAPI (ParserAPI) ) +import qualified Parser.Attoparsec as Attoparsec import qualified Queue.API as Queue import Queue.TransactionalQueue (newQueue) import Router.Nats @@ -41,6 +46,7 @@ ) import State.Types (ClientConfig (..)) import Subscription.Store (newSubscriptionStore) +import Subscription.Types (defaultPendingLimits) import System.Timeout (timeout) import Test.Hspec import Transformers.Transformers (Transformer (transform)) @@ -48,9 +54,21 @@ import Types.Info (Info (Info)) import Types.Ping (Ping (Ping)) import Types.Pong (Pong (Pong)) +import Types.TLS (TLSConfig (..), defaultTLSConfig) spec :: Spec spec = do + describe "TLS configuration" $ do + it "does not expose client private keys when rendered" $ do + let config = defaultTLSConfig + { tlsClientCertificate = Just ("certificate", "private-key") + , tlsRootCertificates = ["private-root"] + } + + show config `shouldNotContain` "private-key" + show config `shouldNotContain` "private-root" + show config `shouldContain` "tlsRootCertificates = 1 configured" + describe "Connection reader" $ do it "unblocks a blocking read when closeReader is called" $ do conn <- newConn connectionApi @@ -74,7 +92,7 @@ describe "Router ping lifecycle" $ do it "responds to server PING by enqueueing PONG" $ do state <- newTestState - store <- newSubscriptionStore + store <- newSubscriptionStore defaultPendingLimits (pure ()) directive <- routeMessage state store (ParsedPing Ping) @@ -86,17 +104,32 @@ Left err -> expectationFailure ("expected queued PONG, got queue error: " ++ err) - it "runs the next pending ping action when PONG is routed" $ do + it "runs exactly one pending ping action for each PONG" $ do state <- newTestState - store <- newSubscriptionStore - actionRan <- newIORef False - pushPingAction state (writeIORef actionRan True) + store <- newSubscriptionStore defaultPendingLimits (pure ()) + firstActionRan <- newIORef False + secondActionRan <- newIORef False + pushPingAction state (writeIORef firstActionRan True) + pushPingAction state (writeIORef secondActionRan True) - directive <- routeMessage state store (ParsedPong Pong) + firstDirective <- routeMessage state store (ParsedPong Pong) - directive `shouldBe` RouteContinue - readIORef actionRan `shouldReturn` True + firstDirective `shouldBe` RouteContinue + readIORef firstActionRan `shouldReturn` True + readIORef secondActionRan `shouldReturn` False + + secondDirective <- routeMessage state store (ParsedPong Pong) + + secondDirective `shouldBe` RouteContinue + readIORef secondActionRan `shouldReturn` True describe "Handshake" $ do + it "returns a typed error when no servers are configured" $ do + result <- Client.newClient [] [Client.withConnectionAttempts 1] + case result of + Left Client.ConnectNoServers -> pure () + Left err -> expectationFailure ("unexpected connection error: " ++ show err) + Right _ -> expectationFailure "client connected without a server" + it "accepts an incremental parser backend for the initial INFO frame" $ do state <- newTestState conn <- newConn connectionApi @@ -105,6 +138,7 @@ newScriptedTransport [ "INF" , "O {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n" + , "PONG\r\n" ] writes pointTransport conn transport @@ -113,7 +147,8 @@ result `shouldBe` Right () readServerInfo state `shouldReturn` Just testInfo - readTVarIO writes `shouldReturn` [LBS.toStrict (transform testConnect)] + readTVarIO writes `shouldReturn` + [LBS.toStrict (transform testConnect <> transform Ping)] it "accepts a parser backend that drops an invalid prefix before INFO" $ do state <- newTestState @@ -123,6 +158,7 @@ newScriptedTransport [ "XIN" , "FO {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n" + , "PONG\r\n" ] writes pointTransport conn transport @@ -131,14 +167,85 @@ result `shouldBe` Right () readServerInfo state `shouldReturn` Just testInfo - readTVarIO writes `shouldReturn` [LBS.toStrict (transform testConnect)] + readTVarIO writes `shouldReturn` + [LBS.toStrict (transform testConnect <> transform Ping)] + it "does not report readiness until the server sends PONG" $ do + state <- newTestStateWithTimeout 20000 + conn <- newConn connectionApi + writes <- newTVarIO [] + blocker <- newEmptyMVar + let transport = Transport + { transportRead = \_ -> do + first <- atomically $ do + written <- readTVar writes + pure (null written) + if first then pure infoFrame else takeMVar blocker + , transportWrite = \bytes -> + atomically $ modifyTVar' writes (<> [bytes]) + , transportWriteLazy = \bytes -> + atomically $ modifyTVar' writes (<> [LBS.toStrict bytes]) + , transportFlush = pure () + , transportClose = pure () + , transportUpgrade = Nothing + } + pointTransport conn transport + + result <- performHandshake connectionApi incrementalInfoParser state auth conn "127.0.0.1" + + result `shouldBe` Left HandshakeTimeout + + it "accepts +OK before the handshake PONG" $ do + state <- newTestState + conn <- newConn connectionApi + writes <- newTVarIO [] + transport <- newScriptedTransport [infoFrame, "+OK\r\nPONG\r\n"] writes + pointTransport conn transport + + result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1" + + result `shouldBe` Right () + + it "returns authentication errors received before PONG" $ do + state <- newTestState + conn <- newConn connectionApi + writes <- newTVarIO [] + transport <- newScriptedTransport + [infoFrame, "-ERR 'Authorization Violation'\r\n"] writes + pointTransport conn transport + + result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1" + + case result of + Left (HandshakeAuthError _) -> pure () + other -> expectationFailure ("expected auth error, got: " ++ show other) + + it "rejects requested TLS when the server does not advertise it" $ do + state <- newTestStateWithConfig $ \cfg -> + cfg { tlsConfig = Just defaultTLSConfig } + conn <- newConn connectionApi + writes <- newTVarIO [] + transport <- newScriptedTransport [infoFrame] writes + pointTransport conn transport + + result <- performHandshake connectionApi Attoparsec.parserApi state auth conn "127.0.0.1" + + case result of + Left (HandshakeTLSError _) -> pure () + other -> expectationFailure ("expected TLS error, got: " ++ show other) + newTestState = do + newTestStateWithConfig id + +newTestStateWithTimeout timeoutMicros = do + newTestStateWithConfig $ \cfg -> cfg { connectTimeoutMicros = timeoutMicros } + +newTestStateWithConfig updateConfig = do queue <- newQueue ctx <- newLogContext conn <- newConn connectionApi logger <- newSilentLogger - newClientState (testConfig logger) queue conn ctx + newClientState (updateConfig (testConfig logger)) queue conn ctx newSilentLogger :: IO LoggerConfig newSilentLogger = do @@ -149,11 +256,12 @@ testConfig logger = ClientConfig { connectionAttempts = 1 + , connectTimeoutMicros = 1000000 , callbackConcurrency = 1 - , bufferLimit = 4096 + , messageLimit = 1024 * 1024 , connectConfig = testConnect , loggerConfig = logger - , tlsCert = Nothing + , tlsConfig = Nothing , exitAction = const (pure ()) , connectOptions = [] } @@ -196,6 +304,7 @@ Nothing Nothing Nothing + Nothing incrementalInfoParser :: ParserAPI ParsedMessage incrementalInfoParser = ParserAPI parseInfo @@ -205,6 +314,8 @@ NeedMore | bytes == infoFrame = Emit (ParsedInfo testInfo) "" + | bytes == "PONG\r\n" = + Emit (ParsedPong Pong) "" | otherwise = Reject ("unexpected input: " ++ show bytes) @@ -216,6 +327,8 @@ DropPrefix 1 "invalid prefix" | bytes == infoFrame = Emit (ParsedInfo testInfo) "" + | bytes == "PONG\r\n" = + Emit (ParsedPong Pong) "" | otherwise = Reject ("unexpected input: " ++ show bytes) diff --git a/test/Unit/ErrSpec.hs b/test/Unit/ErrSpec.hs --- a/test/Unit/ErrSpec.hs +++ b/test/Unit/ErrSpec.hs @@ -21,6 +21,9 @@ ("-ERR 'Attempted To Connect To Route Port'\r\n", ErrRoutePortConn "Attempted To Connect To Route Port"), ("-ERR 'Authorization Violation'\r\n", ErrAuthViolation "Authorization Violation"), ("-ERR 'Authorization Timeout'\r\n", ErrAuthTimeout "Authorization Timeout"), + ("-ERR 'User Authentication Expired: jwt is expired'\r\n", ErrAuthExpired "User Authentication Expired: jwt is expired"), + ("-ERR 'User Authentication Revoked'\r\n", ErrAuthRevoked "User Authentication Revoked"), + ("-ERR 'Account Authentication Expired'\r\n", ErrAccountAuthExpired "Account Authentication Expired"), ("-ERR 'Invalid Client Protocol'\r\n", ErrInvalidProtocol "Invalid Client Protocol"), ("-ERR 'Maximum Control Line Exceeded'\r\n", ErrMaxControlLineEx "Maximum Control Line Exceeded"), ("-ERR 'Secure Connection - TLS Required'\r\n", ErrTlsRequired "Secure Connection - TLS Required"), @@ -30,7 +33,8 @@ ("-ERR 'Maximum Payload Violation'\r\n", ErrMaxPayload "Maximum Payload Violation"), ("-ERR 'Invalid Subject'\r\n", ErrInvalidSubject "Invalid Subject"), ("-ERR 'Permissions Violation For Subscription To FOO.'\r\n", ErrPermViolation "Permissions Violation For Subscription To FOO."), - ("-ERR 'Permissions Violation For Publish To FOO.'\r\n", ErrPermViolation "Permissions Violation For Publish To FOO.") + ("-ERR 'Permissions Violation For Publish To FOO.'\r\n", ErrPermViolation "Permissions Violation For Publish To FOO."), + ("-ERR 'A New Server Error'\r\n", ErrErr "A New Server Error") ] generatedCases = diff --git a/test/Unit/InfoSpec.hs b/test/Unit/InfoSpec.hs --- a/test/Unit/InfoSpec.hs +++ b/test/Unit/InfoSpec.hs @@ -25,15 +25,19 @@ explicitCases = [ ( "INFO {\"server_id\": \"some-server\", \"version\": \"semver\", \"go\": \"1.13\", \"host\": \"127.0.0.1\", \"port\": 4222, \"max_payload\": 1024, \"proto\": 3, \"client_id\": 1, \"nonce\": \"nonce-123\", \"auth_required\": true, \"tls_required\": true, \"connect_urls\": [\"https://127.0.0.1:4222\"], \"ldm\": true, \"headers\": true}\r\n", - Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 (Just 1) (Just "nonce-123") (Just True) (Just True) (Just ["https://127.0.0.1:4222"]) (Just True) (Just True) + Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 (Just 1) (Just "nonce-123") (Just True) (Just True) Nothing (Just ["https://127.0.0.1:4222"]) (Just True) (Just True) ), ( "INFO {\"server_id\": \"some-server\", \"version\": \"semver\", \"go\": \"1.13\", \"host\": \"127.0.0.1\", \"port\": 4222, \"max_payload\": 1024, \"proto\": 3}\r\n", - Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 Nothing Nothing Nothing Nothing Nothing Nothing Nothing + Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing ), ( "INFO {\"server_id\": \"some-server\", \"version\": \"semver\", \"go\": \"1.13\", \"host\": \"127.0.0.1\", \"port\": 4222, \"max_payload\": 1024, \"proto\": 3, \"client_id\": 1, \"auth_required\": true, \"tls_required\": true, \"connect_urls\": [\"https://127.0.0.1:4222\", \"https://192.168.9.7:4222\"], \"ldm\": true, \"headers\": false}\r\n", - Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 (Just 1) Nothing (Just True) (Just True) (Just ["https://127.0.0.1:4222", "https://192.168.9.7:4222"]) (Just True) (Just False) + Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 (Just 1) Nothing (Just True) (Just True) Nothing (Just ["https://127.0.0.1:4222", "https://192.168.9.7:4222"]) (Just True) (Just False) + ), + ( + "INFO {\"server_id\": \"some-server\", \"version\": \"semver\", \"go\": \"1.13\", \"host\": \"127.0.0.1\", \"port\": 4222, \"max_payload\": 1024, \"proto\": 3, \"tls_available\": true}\r\n", + Info "some-server" "semver" "1.13" "127.0.0.1" 4222 1024 3 Nothing Nothing Nothing Nothing (Just True) Nothing Nothing Nothing ) ] @@ -52,6 +56,7 @@ <*> maybeify nonceCases <*> maybeify boolCases <*> maybeify boolCases + <*> maybeify boolCases <*> maybeify connectStringCases <*> maybeify boolCases <*> maybeify boolCases @@ -93,6 +98,8 @@ newField "auth_required" (fmap boolToJSON . auth_required $ m), maybeComma (tls_required m), newField "tls_required" (fmap boolToJSON . tls_required $ m), + maybeComma (tls_available m), + newField "tls_available" (fmap boolToJSON . tls_available $ m), maybeComma (connect_urls m), newField "connect_urls" (fmap arrayToJSON . connect_urls $ m), maybeComma (ldm m), diff --git a/test/Unit/JetStream/ConsumerSpec.hs b/test/Unit/JetStream/ConsumerSpec.hs --- a/test/Unit/JetStream/ConsumerSpec.hs +++ b/test/Unit/JetStream/ConsumerSpec.hs @@ -26,6 +26,14 @@ eitherDecode (encode targetKindConsumerConfigRequest) `shouldBe` Right targetKindConsumerConfigValue + it "applies consumer options from left to right" $ do + let request = consumerConfigRequest + [ withConsumerDeliverPolicy DeliverAll + , withConsumerDeliverPolicy DeliverLast + ] + eitherDecode (encode request) + `shouldBe` Right (object ["deliver_policy" .= ("last" :: String)]) + describe "Consumer reset request JSON" $ do it "encodes an optional reset sequence" $ do eitherDecode (encode (consumerResetRequest [withConsumerResetSequence 7])) @@ -33,7 +41,7 @@ describe "ConsumerInfo JSON" . it "decodes server responses into concrete fields" $ do - eitherDecode consumerInfoJSON `shouldBe` Right consumerInfo + eitherDecode consumerInfoJSON `shouldBe` Right consumerInfoFixture describe "ConsumerResetResponse JSON" . it "decodes reset metadata with the updated consumer info" $ do @@ -47,7 +55,7 @@ { consumerListTotal = 1 , consumerListOffset = 0 , consumerListLimit = 1024 - , consumerListConsumers = [consumerInfo] + , consumerListConsumers = [consumerInfoFixture] } it "decodes consumer name responses" $ do @@ -148,8 +156,8 @@ consumerNamesJSON = "{\"type\":\"io.nats.jetstream.api.v1.consumer_names_response\",\"total\":2,\"offset\":0,\"limit\":1024,\"consumers\":[\"orders-puller\",\"orders-worker\"]}" -consumerInfo :: ConsumerInfo -consumerInfo = ConsumerInfo +consumerInfoFixture :: ConsumerInfo +consumerInfoFixture = ConsumerInfo { consumerInfoStreamName = "ORDERS" , consumerInfoName = "orders-puller" , consumerInfoCreated = timestamp diff --git a/test/Unit/JetStream/MessageSpec.hs b/test/Unit/JetStream/MessageSpec.hs --- a/test/Unit/JetStream/MessageSpec.hs +++ b/test/Unit/JetStream/MessageSpec.hs @@ -2,13 +2,14 @@ module JetStream.MessageSpec (spec) where -import qualified API as Nats +import qualified Client.API as Nats import Data.IORef import JetStream.Error (JetStreamError (..)) import JetStream.Message (fetchMessages) import JetStream.Message.Types import JetStream.Options (newJetStreamContext) import Publish (defaultPublishConfig) +import Publish.Config (publishReplyTo) import Test.Hspec import Types.Msg (Payload, SID, Subject) @@ -41,7 +42,7 @@ describe "message metadata" $ do it "parses v1 JetStream ack reply subjects" $ do let metadata = messageMetadata $ - Message "ORDERS.created" (Just "payload") Nothing + Message "ORDERS.created" "payload" Nothing (Just "$JS.ACK.ORDERS.WORKER.3.10.2.123000000000.4") Nothing fmap messageMetadataStream metadata `shouldBe` Just "ORDERS" @@ -55,14 +56,14 @@ it "parses v2 JetStream ack reply subjects with domains" $ do let metadata = messageMetadata $ - Message "ORDERS.created" (Just "payload") Nothing + Message "ORDERS.created" "payload" Nothing (Just "$JS.ACK.HUB.ACCOUNT.ORDERS.WORKER.1.11.3.124000000000.0") Nothing fmap messageMetadataDomain metadata `shouldBe` Just (Just "HUB") fmap messageMetadataStreamSequence metadata `shouldBe` Just 11 it "rejects non-JetStream reply subjects" $ do - messageMetadata (Message "ORDERS.created" Nothing Nothing (Just "_INBOX.reply") Nothing) + messageMetadata (Message "ORDERS.created" "" Nothing (Just "_INBOX.reply") Nothing) `shouldBe` Nothing describe "pull request payloads" $ do @@ -82,12 +83,12 @@ callbackRef <- newIORef Nothing let client = fakeClient publishCalls unsubscribeCalls callbackRef request = pullRequest [withFetchBatch 2] - result <- fetchMessages (newJetStreamContext client []) "ORDERS" "WORKER" request + result <- fetchMessages (newJetStreamContext client []) "ORDERS" "WORKER" request [] publishCalls' <- readIORef publishCalls publishCalls' `shouldBe` [ ( "$JS.API.CONSUMER.MSG.NEXT.ORDERS.WORKER" - , Just "{\"batch\":2,\"expires\":1000000000}" + , "{\"batch\":2,\"expires\":1000000000}" , Just "_INBOX.batch" ) ] @@ -97,7 +98,7 @@ Left err -> expectationFailure ("fetch failed: " ++ show err) Right response -> do - fmap messagePayload (pullResponseMessages response) `shouldBe` [Just "one", Just "two"] + fmap messagePayload (pullResponseMessages response) `shouldBe` ["one", "two"] pullResponseStatus response `shouldBe` Nothing it "returns a JetStream timeout when no pull response arrives" $ do @@ -105,13 +106,13 @@ unsubscribeCalls <- newIORef [] let client = silentFakeClient publishCalls unsubscribeCalls request = pullRequest [withFetchWait (FetchNoWaitMicros 1000)] - result <- fetchMessages (newJetStreamContext client []) "ORDERS" "WORKER" request + result <- fetchMessages (newJetStreamContext client []) "ORDERS" "WORKER" request [] result `shouldBe` Left JetStreamTimeout publishCalls' <- readIORef publishCalls publishCalls' `shouldBe` [ ( "$JS.API.CONSUMER.MSG.NEXT.ORDERS.WORKER" - , Just "{\"batch\":1,\"no_wait\":true}" + , "{\"batch\":1,\"no_wait\":true}" , Just "_INBOX.timeout" ) ] @@ -119,62 +120,68 @@ unsubscribeCalls' `shouldBe` ["sid-timeout"] fakeClient - :: IORef [(Subject, Maybe Payload, Maybe Subject)] + :: IORef [(Subject, Payload, Maybe Subject)] -> IORef [SID] - -> IORef (Maybe (Maybe Nats.MsgView -> IO ())) + -> IORef (Maybe (Nats.MsgView -> IO ())) -> Nats.Client fakeClient publishCalls unsubscribeCalls callbackRef = Nats.Client - { Nats.publish = \subject options -> do - let (body, _, _, replyTo') = foldr ($) defaultPublishConfig options + { Nats.publish = \subject body options -> do + let replyTo' = publishReplyTo (foldl (flip ($)) defaultPublishConfig options) modifyIORef' publishCalls (++ [(subject, body, replyTo')]) callback <- readIORef callbackRef case callback of Nothing -> pure () Just deliver -> do - deliver (Just (fakeMsg "one")) - deliver (Just (fakeMsg "two")) + deliver (fakeMsg "one") + deliver (fakeMsg "two") + pure (Right ()) , Nats.subscribe = \_ _ callback -> do writeIORef callbackRef (Just callback) - pure "sid-1" - , Nats.request = \_ _ _ -> pure "sid-request" - , Nats.unsubscribe = \sid -> + pure (Right (Nats.Subscription "sid-1")) + , Nats.subscribeOnce = \_ _ _ -> pure (Right (Nats.Subscription "sid-once")) + , Nats.request = \_ _ _ -> pure (Left Nats.NatsRequestTimedOut) + , Nats.unsubscribe = \(Nats.Subscription sid) _ -> do modifyIORef' unsubscribeCalls (++ [sid]) + pure (Right ()) , Nats.newInbox = pure "_INBOX.batch" - , Nats.ping = id - , Nats.flush = pure () - , Nats.reset = pure () - , Nats.close = pure () + , Nats.ping = \_ -> pure (Right ()) + , Nats.flush = \_ -> pure (Right ()) + , Nats.reset = \_ -> pure () + , Nats.close = \_ -> pure () } -fakeMsg :: Payload -> Nats.MsgView +fakeMsg :: Payload -> Nats.Message fakeMsg body = - Nats.MsgView + Nats.Message { Nats.subject = "ORDERS.created" , Nats.sid = "sid-1" , Nats.replyTo = Just "$JS.ACK.ORDERS.WORKER.1.1.1" - , Nats.payload = Just body + , Nats.payload = body , Nats.headers = Nothing } silentFakeClient - :: IORef [(Subject, Maybe Payload, Maybe Subject)] + :: IORef [(Subject, Payload, Maybe Subject)] -> IORef [SID] -> Nats.Client silentFakeClient publishCalls unsubscribeCalls = Nats.Client - { Nats.publish = \subject options -> do - let (body, _, _, replyTo') = foldr ($) defaultPublishConfig options + { Nats.publish = \subject body options -> do + let replyTo' = publishReplyTo (foldl (flip ($)) defaultPublishConfig options) modifyIORef' publishCalls (++ [(subject, body, replyTo')]) + pure (Right ()) , Nats.subscribe = \_ _ _ -> - pure "sid-timeout" - , Nats.request = \_ _ _ -> pure "sid-request" - , Nats.unsubscribe = \sid -> + pure (Right (Nats.Subscription "sid-timeout")) + , Nats.subscribeOnce = \_ _ _ -> pure (Right (Nats.Subscription "sid-once")) + , Nats.request = \_ _ _ -> pure (Left Nats.NatsRequestTimedOut) + , Nats.unsubscribe = \(Nats.Subscription sid) _ -> do modifyIORef' unsubscribeCalls (++ [sid]) + pure (Right ()) , Nats.newInbox = pure "_INBOX.timeout" - , Nats.ping = id - , Nats.flush = pure () - , Nats.reset = pure () - , Nats.close = pure () + , Nats.ping = \_ -> pure (Right ()) + , Nats.flush = \_ -> pure (Right ()) + , Nats.reset = \_ -> pure () + , Nats.close = \_ -> pure () } diff --git a/test/Unit/JetStream/ProtocolSpec.hs b/test/Unit/JetStream/ProtocolSpec.hs --- a/test/Unit/JetStream/ProtocolSpec.hs +++ b/test/Unit/JetStream/ProtocolSpec.hs @@ -2,11 +2,17 @@ module JetStream.ProtocolSpec (spec) where -import qualified API as Nats +import qualified Client.API as Nats import Data.Aeson (eitherDecodeStrict) import Data.Either (isLeft) import JetStream.Error (JetStreamError (..)) -import JetStream.Options (newJetStreamContext, withDomain) +import JetStream.Options + ( JetStreamConfigError (..) + , newJetStreamContext + , requestTimeoutMicros + , tryNewJetStreamContext + , withDomain + ) import JetStream.Protocol.Headers (statusCode, statusDescription) import JetStream.Protocol.Request (decodeJetStreamResponse) import JetStream.Protocol.Subject @@ -14,7 +20,9 @@ import JetStream.Types ( AccountInfo (..) , AccountTier (..) + , AckPolicy (..) , DeliverPolicy + , withRequestTimeout ) import Test.Hspec @@ -35,6 +43,17 @@ consumerPauseSubject testContext "ORDERS" "WORKER" `shouldBe` "$JS.API.CONSUMER.PAUSE.ORDERS.WORKER" consumerResetSubject testContext "ORDERS" "WORKER" `shouldBe` "$JS.API.CONSUMER.RESET.ORDERS.WORKER" + describe "JetStream options" $ do + it "rejects invalid construction options" $ do + case tryNewJetStreamContext fakeClient [withDomain ""] of + Left EmptyJetStreamDomain -> pure () + _ -> expectationFailure "expected an empty-domain error" + + it "applies request options left to right" $ do + requestTimeoutMicros testContext + [withRequestTimeout 2, withRequestTimeout 0.25] + `shouldBe` 250000 + describe "JetStream status headers" $ do it "extracts status and description from normal headers" $ do let msg = fakeMsg { Nats.headers = Just [("Status", "404"), ("Description", "No Messages")] } @@ -70,6 +89,10 @@ it "rejects parameterised deliver policies without their companion fields" $ do (eitherDecodeStrict "\"by_start_sequence\"" :: Either String DeliverPolicy) `shouldSatisfy` isLeft + + it "preserves unknown wire enum values" $ do + eitherDecodeStrict "\"future_ack_policy\"" + `shouldBe` Right (AckPolicyUnknown "future_ack_policy") (eitherDecodeStrict "\"by_start_time\"" :: Either String DeliverPolicy) `shouldSatisfy` isLeft @@ -85,23 +108,24 @@ fakeClient :: Nats.Client fakeClient = Nats.Client - { Nats.publish = \_ _ -> pure () - , Nats.subscribe = \_ _ _ -> pure "0" - , Nats.request = \_ _ _ -> pure "0" - , Nats.unsubscribe = \_ -> pure () + { Nats.publish = \_ _ _ -> pure (Right ()) + , Nats.subscribe = \_ _ _ -> pure (Right (Nats.Subscription "0")) + , Nats.subscribeOnce = \_ _ _ -> pure (Right (Nats.Subscription "0")) + , Nats.request = \_ _ _ -> pure (Left Nats.NatsRequestTimedOut) + , Nats.unsubscribe = \_ _ -> pure (Right ()) , Nats.newInbox = pure "_INBOX.TEST" - , Nats.ping = id - , Nats.flush = pure () - , Nats.reset = pure () - , Nats.close = pure () + , Nats.ping = \_ -> pure (Right ()) + , Nats.flush = \_ -> pure (Right ()) + , Nats.reset = \_ -> pure () + , Nats.close = \_ -> pure () } -fakeMsg :: Nats.MsgView +fakeMsg :: Nats.Message fakeMsg = - Nats.MsgView + Nats.Message { Nats.subject = "subject" , Nats.sid = "sid" , Nats.replyTo = Nothing - , Nats.payload = Nothing + , Nats.payload = "" , Nats.headers = Nothing } diff --git a/test/Unit/JetStream/PublishSpec.hs b/test/Unit/JetStream/PublishSpec.hs --- a/test/Unit/JetStream/PublishSpec.hs +++ b/test/Unit/JetStream/PublishSpec.hs @@ -17,10 +17,10 @@ , withPublishExpectation (ExpectedLastSequence 12) , withHeaders [("Nats-Expected-Test", "custom")] ] - `shouldBe` [ ("Nats-Msg-Id", "msg-1") - , ("Nats-Expected-Stream", "ORDERS") + `shouldBe` [ ("Nats-Expected-Test", "custom") , ("Nats-Expected-Last-Sequence", "12") - , ("Nats-Expected-Test", "custom") + , ("Nats-Expected-Stream", "ORDERS") + , ("Nats-Msg-Id", "msg-1") ] it "keeps publish expectations mutually exclusive" $ do @@ -29,7 +29,7 @@ , withPublishExpectation (ExpectedLastSubjectSequence 7) , withPublishExpectation (ExpectedLastMsgId "msg-0") ] - `shouldBe` [("Nats-Expected-Last-Sequence", "12")] + `shouldBe` [("Nats-Expected-Last-Msg-Id", "msg-0")] describe "PublishAck" $ do it "decodes required and optional publish ack fields" $ do diff --git a/test/Unit/NKeySpec.hs b/test/Unit/NKeySpec.hs --- a/test/Unit/NKeySpec.hs +++ b/test/Unit/NKeySpec.hs @@ -2,9 +2,27 @@ module NKeySpec (spec) where -import qualified Auth.Jwt as Jwt -import qualified Auth.NKey as NKey -import qualified Data.ByteString as BS +import Auth.Config (mergeAuth) +import qualified Auth.Jwt as Jwt +import qualified Auth.NKey as NKey +import Auth.Resolver (buildAuthPatch) +import qualified Auth.Token as Token +import Auth.Types + ( AuthContext (AuthContext) + , AuthError (AuthError) + , AuthPatch (..) + , emptyAuthPatch + ) +import qualified Auth.UserPass as UserPass +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BC +import Data.Either (isLeft) +import Data.IORef + ( atomicModifyIORef' + , newIORef + , readIORef + , writeIORef + ) import Test.Hspec spec :: Spec @@ -23,6 +41,24 @@ it "returns Nothing when blocks are missing" $ do Jwt.parseJwtBundle "missing blocks" `shouldBe` Nothing + it "rejects a credentials bundle with a missing end marker" $ do + let creds = + "-----BEGIN NATS USER JWT-----\n\ + \jwt-token\n\ + \-----BEGIN USER NKEY SEED-----\n\ + \seed-token\n\ + \------END USER NKEY SEED------" + Jwt.parseJwtBundle creds `shouldBe` Nothing + + it "rejects empty credential blocks" $ do + let creds = + "-----BEGIN NATS USER JWT-----\n\ + \------END NATS USER JWT------\n\ + \-----BEGIN USER NKEY SEED-----\n\ + \seed-token\n\ + \------END USER NKEY SEED------" + Jwt.parseJwtBundle creds `shouldBe` Nothing + describe "signNonceWithSeed" $ do it "derives the expected public key from a seed" $ do let seed = "SUAHR6JNS2HKJQEAQFHYPOXFXWE4JXBPKUWFX3IMYU72UHOGXT3ZMVFHXI" @@ -32,3 +68,83 @@ Right (pub, sig) -> do pub `shouldBe` expectedPub BS.length sig `shouldBe` 86 + + it "rejects a seed with non-canonical trailing base32 data" $ do + let seed = "SUAHR6JNS2HKJQEAQFHYPOXFXWE4JXBPKUWFX3IMYU72UHOGXT3ZMVFHXI" + + NKey.signNonceWithSeed (seed <> "A") "nonce-123" `shouldSatisfy` isLeft + + it "rejects lower-case seed encodings" $ do + let seed = "suahr6jns2hkjqeaqfhypoxfxwe4jxbpkuwfx3imyu72uhogxt3zmvfhxi" + + NKey.signNonceWithSeed seed "nonce-123" `shouldSatisfy` isLeft + + it "rejects seeds with an invalid checksum" $ do + let seed = "SUAHR6JNS2HKJQEAQFHYPOXFXWE4JXBPKUWFX3IMYU72UHOGXT3ZMVFHXJ" + + NKey.signNonceWithSeed seed "nonce-123" `shouldSatisfy` isLeft + + describe "secret rendering" $ do + it "redacts credentials and authentication patches" $ do + let bundle = Jwt.JwtBundle "jwt-secret" "seed-secret" + patch = emptyAuthPatch + { patchPass = Just "password-secret" + , patchSig = Just "signature-secret" + } + + show bundle `shouldNotContain` "jwt-secret" + show bundle `shouldNotContain` "seed-secret" + show patch `shouldNotContain` "password-secret" + show patch `shouldNotContain` "signature-secret" + + describe "auth handlers" $ do + it "fetches a fresh token for every connection attempt" $ do + calls <- newIORef (0 :: Int) + let handler = do + call <- atomicModifyIORef' calls $ \current -> + let next = current + 1 + in (next, next) + pure (Right ("token-" <> BC.pack (show call))) + auth = Token.authHandler handler + + first <- buildAuthPatch auth (AuthContext Nothing) + second <- buildAuthPatch auth (AuthContext Nothing) + + fmap patchAuthToken first `shouldBe` Right (Just "token-1") + fmap patchAuthToken second `shouldBe` Right (Just "token-2") + + it "allows token and user/pass credentials to coexist" $ do + let auth = mergeAuth (Token.auth "token") (UserPass.auth ("user", "pass")) + + result <- buildAuthPatch auth (AuthContext Nothing) + + fmap patchAuthToken result `shouldBe` Right (Just "token") + fmap patchUser result `shouldBe` Right (Just "user") + fmap patchPass result `shouldBe` Right (Just "pass") + + it "base64url encodes raw signatures returned by handlers" $ do + let publicKey = "UAB7EFDOTOBBMPOCK4SXFA62FVZOADQDZOU2W4IUDCGFKJXYVOK3LV7X" + auth = NKey.authHandler publicKey $ \_ -> + pure (Right (BS.replicate 64 0)) + + result <- buildAuthPatch auth (AuthContext (Just "nonce")) + + fmap (fmap BS.length . patchSig) result `shouldBe` Right (Just 86) + + it "validates public nkeys before invoking signature handlers" $ do + called <- newIORef False + let auth = NKey.authHandler "USER_PUBLIC_KEY" $ \_ -> do + writeIORef called True + pure (Right (BS.replicate 64 0)) + + result <- buildAuthPatch auth (AuthContext (Just "nonce")) + + result `shouldSatisfy` isLeft + readIORef called `shouldReturn` False + + it "returns handler failures as typed authentication errors" $ do + let auth = Token.authHandler (pure (Left (AuthError "token unavailable"))) + + result <- buildAuthPatch auth (AuthContext Nothing) + + result `shouldBe` Left (AuthError "token unavailable") diff --git a/test/Unit/ParserSpec.hs b/test/Unit/ParserSpec.hs --- a/test/Unit/ParserSpec.hs +++ b/test/Unit/ParserSpec.hs @@ -4,10 +4,10 @@ import Parser.API ( ParseStep (DropPrefix, Emit, NeedMore) - , ParsedMessage (ParsedMsg, ParsedPing) + , ParsedMessage (ParsedMessageTooLarge, ParsedMsg, ParsedPing) , parse ) -import Parser.Attoparsec (parserApi) +import Parser.Attoparsec (parserApi, parserApiWithMessageLimit) import Test.Hspec import qualified Types.Msg as Msg import Types.Ping (Ping (Ping)) @@ -20,6 +20,21 @@ it "suggests pulling more data for truncated input" $ do parse parserApi "MSG FOO 1 5\r\nHEL" `shouldBe` NeedMore + + it "rejects an oversized MSG before waiting for its payload" $ do + parse (parserApiWithMessageLimit 4) "MSG FOO 1 5\r\n" + `shouldBe` Emit (ParsedMessageTooLarge 5 4) "" + + it "rejects an oversized HMSG before waiting for its body" $ do + parse (parserApiWithMessageLimit 15) "HMSG FOO 1 12 16\r\n" + `shouldBe` Emit (ParsedMessageTooLarge 16 15) "" + + it "accepts a message exactly at the configured limit" $ do + parse (parserApiWithMessageLimit 5) "MSG FOO 1 5\r\nHELLO\r\n" + `shouldBe` + Emit + (ParsedMsg (Msg.Msg "FOO" "1" Nothing (Just "HELLO") Nothing)) + "" it "suggests dropping invalid prefix bytes" $ do case parse parserApi "LOL" of diff --git a/test/Unit/StreamingSpec.hs b/test/Unit/StreamingSpec.hs --- a/test/Unit/StreamingSpec.hs +++ b/test/Unit/StreamingSpec.hs @@ -143,7 +143,7 @@ hClose server hClose client describe "Broadcasting" $ do - it "drops messages larger than the buffer limit" $ do + it "writes transformed messages without applying a parser buffer limit" $ do dl <- defaultLogger' ctx <- newLogContext q <- newQueue @@ -152,8 +152,8 @@ let BroadcastingAPI runBroadcasting = broadcastingApi enqueue q (QueueItem pub) `shouldReturn` Right () close q - runWithLogger dl ctx (runBroadcasting 5 q captureWriterApi (CaptureWriter output) :: AppM ()) - readTVarIO output `shouldReturn` [] + runWithLogger dl ctx (runBroadcasting q captureWriterApi (CaptureWriter output) :: AppM ()) + readTVarIO output `shouldReturn` ["PUB FOO 10\r\n0123456789\r\n"] ensureTVarIsEmpty :: TVar ByteString -> STM () ensureTVarIsEmpty tvar = do diff --git a/test/Unit/SubscriptionStoreSpec.hs b/test/Unit/SubscriptionStoreSpec.hs new file mode 100644 --- /dev/null +++ b/test/Unit/SubscriptionStoreSpec.hs @@ -0,0 +1,149 @@ +{-# LANGUAGE OverloadedStrings #-} + +module SubscriptionStoreSpec (spec) where + +import Control.Concurrent +import Control.Concurrent.STM +import Control.Exception +import Control.Monad +import Data.IORef +import Subscription.Store +import Subscription.Types +import Test.Hspec +import Types.Msg + +spec :: Spec +spec = do + describe "global pending delivery limits" $ do + it "shares the message limit across subscriptions and coalesces notification" $ do + slowEvents <- newIORef (0 :: Int) + firstOverflow <- newIORef (0 :: Int) + secondOverflow <- newIORef (0 :: Int) + store <- + newSubscriptionStore + (PendingLimits 1 1024) + (modifyIORef' slowEvents (+ 1)) + registerSubscription store "1" "A" (modifyIORef' firstOverflow (+ 1)) + registerSubscription store "2" "B" (modifyIORef' secondOverflow (+ 1)) + + dispatchMessage store (message "1" "one") + `shouldReturn` DispatchQueued + dispatchMessage store (message "2" "two") + `shouldReturn` DispatchDropped True + dispatchMessage store (message "2" "three") + `shouldReturn` DispatchDropped False + + readIORef slowEvents `shouldReturn` 0 + readIORef firstOverflow `shouldReturn` 0 + readIORef secondOverflow `shouldReturn` 2 + + it "rejects a delivery that exceeds the global byte limit" $ do + store <- newSubscriptionStore (PendingLimits 10 4) (pure ()) + registerSubscription store "1" "A" (pure ()) + + dispatchMessage store (message "1" "12345") + `shouldReturn` DispatchDropped True + + it "re-enables notification after the backlog recovers" $ do + slowEvents <- newIORef (0 :: Int) + callbackStarted <- newEmptyMVar + releaseCallback <- newEmptyMVar + stopping <- newTVarIO False + store <- + newSubscriptionStore + (PendingLimits 1 1024) + (modifyIORef' slowEvents (+ 1)) + register + store + "1" + (SubscriptionMeta "A" Nothing False) + (SubscribeConfig Nothing Nothing) + (\_ -> putMVar callbackStarted () >> takeMVar releaseCallback) + (pure ()) + startWorkers + 1 + store + (readTVar stopping >>= check) + (const (pure ())) + + dispatchMessage store (message "1" "one") + `shouldReturn` DispatchQueued + takeMVar callbackStarted + dispatchMessage store (message "1" "two") + `shouldReturn` DispatchDropped True + putMVar releaseCallback () + atomically (awaitCallbackDrain store) + readIORef slowEvents `shouldReturn` 1 + + dispatchMessage store (message "1" "three") + `shouldReturn` DispatchQueued + takeMVar callbackStarted + dispatchMessage store (message "1" "four") + `shouldReturn` DispatchDropped True + putMVar releaseCallback () + atomically (awaitCallbackDrain store) + readIORef slowEvents `shouldReturn` 2 + atomically (writeTVar stopping True) + + it "releases capacity when a callback throws" $ do + callbackFailed <- newEmptyMVar + stopping <- newTVarIO False + store <- newSubscriptionStore (PendingLimits 1 1024) (pure ()) + register + store + "1" + (SubscriptionMeta "A" Nothing False) + (SubscribeConfig Nothing Nothing) + (const (throwIO (userError "callback failed"))) + (pure ()) + startWorkers + 1 + store + (readTVar stopping >>= check) + (putMVar callbackFailed) + + dispatchMessage store (message "1" "one") + `shouldReturn` DispatchQueued + void (takeMVar callbackFailed) + atomically (awaitCallbackDrain store) + dispatchMessage store (message "1" "two") + `shouldReturn` DispatchQueued + void (takeMVar callbackFailed) + atomically (awaitCallbackDrain store) + atomically (writeTVar stopping True) + + it "removes a one-shot subscription whose delivery is dropped" $ do + store <- newSubscriptionStore (PendingLimits 1 1024) (pure ()) + registerSubscription store "1" "A" (pure ()) + register + store + "2" + (SubscriptionMeta "B" Nothing True) + (SubscribeConfig Nothing Nothing) + (const (pure ())) + (pure ()) + + dispatchMessage store (message "1" "one") + `shouldReturn` DispatchQueued + dispatchMessage store (message "2" "two") + `shouldReturn` DispatchDropped True + dispatchMessage store (message "2" "two") + `shouldReturn` DispatchMissing + +registerSubscription + :: SubscriptionStore + -> SID + -> Subject + -> IO () + -> IO () +registerSubscription store sidValue subjectValue = + register + store + sidValue + (SubscriptionMeta subjectValue Nothing False) + (SubscribeConfig Nothing Nothing) + (const (pure ())) + +message :: SID -> Payload -> Msg +message sidValue payloadValue = + Msg "A" sidValue Nothing (Just payloadValue) Nothing