natskell (empty) → 0.0.0.1
raw patch · 85 files changed
Files
- LICENSE +21/−0
- bench/Main.hs +557/−0
- client/API.hs +118/−0
- client/Client.hs +440/−0
- internal/Lib/CallOption.hs +6/−0
- internal/Lib/Logger.hs +78/−0
- internal/Lib/Logger/Types.hs +57/−0
- internal/Lib/LoggerAPI.hs +15/−0
- internal/Lib/WaitGroup.hs +36/−0
- internal/Lib/WorkerPool.hs +41/−0
- internal/Plumbing/Network/Connection.hs +29/−0
- internal/Plumbing/Network/Connection/Core.hs +229/−0
- internal/Plumbing/Network/Connection/Tcp.hs +45/−0
- internal/Plumbing/Network/Connection/Tls.hs +110/−0
- internal/Plumbing/Network/Connection/Types.hs +42/−0
- internal/Plumbing/Network/ConnectionAPI.hs +41/−0
- internal/Plumbing/Parser/API.hs +29/−0
- internal/Plumbing/Parser/Attoparsec.hs +292/−0
- internal/Plumbing/Pipeline/Broadcasting.hs +22/−0
- internal/Plumbing/Pipeline/Broadcasting/API.hs +12/−0
- internal/Plumbing/Pipeline/Broadcasting/Sink.hs +24/−0
- internal/Plumbing/Pipeline/Broadcasting/Source.hs +16/−0
- internal/Plumbing/Pipeline/Broadcasting/Transformer.hs +18/−0
- internal/Plumbing/Pipeline/Streaming.hs +27/−0
- internal/Plumbing/Pipeline/Streaming/API.hs +12/−0
- internal/Plumbing/Pipeline/Streaming/Parser.hs +63/−0
- internal/Plumbing/Pipeline/Streaming/Sink.hs +9/−0
- internal/Plumbing/Pipeline/Streaming/Source.hs +26/−0
- internal/Plumbing/Queue/API.hs +20/−0
- internal/Plumbing/Queue/TransactionalQueue.hs +56/−0
- internal/Plumbing/Transformers/Transformers.hs +8/−0
- internal/Plumbing/Validators/Validators.hs +7/−0
- internal/Policy/Auth/Jwt.hs +10/−0
- internal/Policy/Auth/NKey.hs +9/−0
- internal/Policy/Auth/None.hs +8/−0
- internal/Policy/Auth/Token.hs +8/−0
- internal/Policy/Auth/Types.hs +395/−0
- internal/Policy/Auth/UserPass.hs +8/−0
- internal/Policy/Engine.hs +263/−0
- internal/Policy/Handshake/Nats.hs +116/−0
- internal/Policy/Publish.hs +8/−0
- internal/Policy/Publish/Config.hs +10/−0
- internal/Policy/Router/Nats.hs +66/−0
- internal/Policy/State/Store.hs +224/−0
- internal/Policy/State/Types.hs +41/−0
- internal/Policy/Subscription/Store.hs +183/−0
- internal/Policy/Subscription/Types.hs +16/−0
- internal/Types/Connect.hs +171/−0
- internal/Types/Err.hs +30/−0
- internal/Types/Info.hs +79/−0
- internal/Types/Msg.hs +23/−0
- internal/Types/Nuid.hs +95/−0
- internal/Types/Ok.hs +4/−0
- internal/Types/Ping.hs +15/−0
- internal/Types/Pong.hs +15/−0
- internal/Types/Pub.hs +85/−0
- internal/Types/Sid.hs +18/−0
- internal/Types/Sub.hs +43/−0
- internal/Types/Unsub.hs +31/−0
- natskell.cabal +241/−0
- test/Fuzz/ParserSpec.hs +370/−0
- test/Fuzz/Spec.hs +1/−0
- test/Fuzz/ValidatorsSpec.hs +131/−0
- test/Integration/ClientSpec.hs +365/−0
- test/Integration/Skip.hs +4/−0
- test/Unit/ConnectSpec.hs +75/−0
- test/Unit/ConnectionSpec.hs +208/−0
- test/Unit/ErrSpec.hs +54/−0
- test/Unit/Fixtures.hs +67/−0
- test/Unit/InfoSpec.hs +131/−0
- test/Unit/MsgSpec.hs +107/−0
- test/Unit/NKeySpec.hs +34/−0
- test/Unit/NuidSpec.hs +20/−0
- test/Unit/OkSpec.hs +29/−0
- test/Unit/ParserSpec.hs +60/−0
- test/Unit/PingSpec.hs +48/−0
- test/Unit/PongSpec.hs +45/−0
- test/Unit/PubSpec.hs +56/−0
- test/Unit/SidSpec.hs +15/−0
- test/Unit/Spec.hs +1/−0
- test/Unit/StreamingSpec.hs +212/−0
- test/Unit/SubSpec.hs +44/−0
- test/Unit/UnsubSpec.hs +43/−0
- test/Unit/WaitGroupSpec.hs +14/−0
- test/Unit/WorkerPoolSpec.hs +59/−0
+ LICENSE view
@@ -0,0 +1,21 @@+MIT License++Copyright (c) 2022 Sam White++Permission is hereby granted, free of charge, to any person obtaining a copy+of this software and associated documentation files (the "Software"), to deal+in the Software without restriction, including without limitation the rights+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell+copies of the Software, and to permit persons to whom the Software is+furnished to do so, subject to the following conditions:++The above copyright notice and this permission notice shall be included in all+copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE+SOFTWARE.
+ bench/Main.hs view
@@ -0,0 +1,557 @@+{-# LANGUAGE OverloadedStrings #-}++module Main (main) where++import Control.Exception (evaluate)+import Control.Monad (when)+import Data.Bits (xor)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as B8+import Data.Int (Int64)+import Data.List (foldl')+import Data.Maybe (fromMaybe)+import Data.Word (Word32, Word64)+import GHC.Clock (getMonotonicTimeNSec)+import GHC.Stats+import Numeric (showFFloat)+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore, Reject)+ , ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong)+ , ParserAPI+ , parse+ )+import Parser.Attoparsec (parserApi)+import System.Environment (getArgs)+import System.Exit (die)+import System.Mem (performGC)+import Text.Printf (printf)+import Types.Err (Err (..))+import Types.Info+ ( client_id+ , connect_urls+ , go+ , host+ , max_payload+ , nonce+ , port+ , proto+ , server_id+ , version+ )+import Types.Msg (Msg (..))++data Config = Config+ { scenarioName :: String+ , repetitions :: Int+ , scale :: Int+ }++data Scenario = Scenario+ { name :: String+ , description :: String+ , chunks :: [BS.ByteString]+ , frameCount :: Int+ , inputBytes :: Int+ }++data RunSummary = RunSummary+ { parsedFrames :: !Int+ , droppedBytes :: !Int+ , checksum :: !Int+ }++data RuntimeDelta = RuntimeDelta+ { allocatedBytes :: !Word64+ , copiedBytes :: !Word64+ , peakLiveBytes :: !Word64+ , peakMemBytes :: !Word64+ , gcCount :: !Word32+ , mutatorCpuNs :: !Int64+ , gcCpuNs :: !Int64+ }++data Aggregate = Aggregate+ { totalFrames :: !Int+ , totalDropped :: !Int+ , combinedChecksums :: !Int+ }++defaultConfig :: Config+defaultConfig =+ Config+ { scenarioName = "mixed-chunked"+ , repetitions = 5+ , scale = 256+ }++main :: IO ()+main = do+ args <- getArgs+ if "--list-scenarios" `elem` args+ then listScenarios+ else do+ config <- parseArgs defaultConfig args+ scenario <- resolveScenario config+ benchmarkScenario parserApi config scenario++listScenarios :: IO ()+listScenarios =+ mapM_ printScenario availableScenarioNames+ where+ printScenario (scenarioKey, scenarioDescription) =+ putStrLn (scenarioKey ++ " - " ++ scenarioDescription)++availableScenarioNames :: [(String, String)]+availableScenarioNames =+ [ ("mixed-aligned", "Mixed INFO/PING/PONG/MSG/HMSG/ERR stream with frame-aligned input.")+ , ("mixed-chunked", "The same mixed stream split across irregular chunk boundaries.")+ , ("mixed-chunked-noisy", "Chunked mixed stream with occasional invalid prefix bytes.")+ ]++parseArgs :: Config -> [String] -> IO Config+parseArgs config args =+ case args of+ [] ->+ pure config+ "--scenario":value:rest ->+ parseArgs config { scenarioName = value } rest+ "--repetitions":value:rest ->+ parsePositiveInt "--repetitions" value >>= \parsed ->+ parseArgs config { repetitions = parsed } rest+ "--scale":value:rest ->+ parsePositiveInt "--scale" value >>= \parsed ->+ parseArgs config { scale = parsed } rest+ flag:_ | "--" `B8.isPrefixOf` B8.pack flag ->+ die ("unknown benchmark flag: " ++ flag)+ unexpected:_ ->+ die ("unexpected benchmark argument: " ++ unexpected)++parsePositiveInt :: String -> String -> IO Int+parsePositiveInt flag value =+ case reads value of+ [(parsed, "")]+ | parsed > 0 ->+ pure parsed+ _ ->+ die (flag ++ " expects a positive integer, got: " ++ value)++resolveScenario :: Config -> IO Scenario+resolveScenario config =+ case scenarioName config of+ "mixed-aligned" ->+ pure (alignedScenario (scale config))+ "mixed-chunked" ->+ pure (chunkedScenario (scale config))+ "mixed-chunked-noisy" ->+ pure (noisyScenario (scale config))+ unknown ->+ die ("unknown scenario: " ++ unknown)++benchmarkScenario :: ParserAPI ParsedMessage -> Config -> Scenario -> IO ()+benchmarkScenario selectedParser config scenario = do+ putStrLn "parser backend: attoparsec"+ putStrLn ("scenario: " ++ name scenario)+ putStrLn ("description: " ++ description scenario)+ putStrLn ("repetitions: " ++ show (repetitions config))+ putStrLn ("frames per repetition: " ++ show (frameCount scenario))+ putStrLn ("input bytes per repetition: " ++ show (inputBytes scenario))++ warmup <- runScenario selectedParser scenario+ validateScenario scenario warmup++ performGC+ beforeStats <- snapshotStats+ startNs <- getMonotonicTimeNSec+ aggregate <- runRepetitions (repetitions config) (runScenario selectedParser scenario)+ endNs <- getMonotonicTimeNSec+ performGC+ afterStats <- snapshotStats++ let elapsedNs = endNs - startNs+ elapsedSeconds = fromIntegral elapsedNs / 1e9 :: Double+ averageSeconds = elapsedSeconds / fromIntegral (repetitions config)+ totalParsedFrames = totalFrames aggregate+ totalInputBytes = repetitions config * inputBytes scenario+ totalDroppedBytes = totalDropped aggregate+ averageRate = fromIntegral totalParsedFrames / elapsedSeconds+ averageMiBPerSecond = fromIntegral totalInputBytes / (1024 * 1024) / elapsedSeconds :: Double+ finalChecksum = combinedChecksums aggregate+ delta = diffStats beforeStats afterStats+ allocatedPerFrame = fromIntegral (allocatedBytes delta) / fromIntegral totalParsedFrames :: Double+ copiedPerFrame = fromIntegral (copiedBytes delta) / fromIntegral totalParsedFrames :: Double++ printf "average delivery rate: %.2f frames/s\n" averageRate+ printf "average parser throughput: %.2f MiB/s\n" averageMiBPerSecond+ printf "average time per repetition: %.6fs\n" averageSeconds+ putStrLn ("dropped invalid prefix bytes: " ++ show totalDroppedBytes)+ putStrLn ("checksum: " ++ show finalChecksum)+ putStrLn ("allocated bytes: " ++ show (allocatedBytes delta))+ printf "allocated bytes per frame: %.2f\n" allocatedPerFrame+ putStrLn ("copied bytes: " ++ show (copiedBytes delta))+ printf "copied bytes per frame: %.2f\n" copiedPerFrame+ putStrLn ("peak live bytes: " ++ show (peakLiveBytes delta))+ putStrLn ("peak memory in use bytes: " ++ show (peakMemBytes delta))+ putStrLn ("GC count: " ++ show (gcCount delta))+ putStrLn ("mutator CPU ms: " ++ formatMillis (mutatorCpuNs delta))+ putStrLn ("GC CPU ms: " ++ formatMillis (gcCpuNs delta))+ putStrLn "For the full RTS breakdown, rerun with `+RTS -s`."++validateScenario :: Scenario -> RunSummary -> IO ()+validateScenario scenario summary =+ when+ (parsedFrames summary /= frameCount scenario)+ ( die+ ( "benchmark corpus parse mismatch: expected "+ ++ show (frameCount scenario)+ ++ " frames but parsed "+ ++ show (parsedFrames summary)+ )+ )++runScenario :: ParserAPI ParsedMessage -> Scenario -> IO RunSummary+runScenario parserForScenario scenario =+ pure (consumeChunks parserForScenario (chunks scenario))++runRepetitions :: Int -> IO RunSummary -> IO Aggregate+runRepetitions repetitionCount action =+ go 0 (Aggregate 0 0 0)+ where+ go completed aggregate+ | completed >= repetitionCount =+ pure aggregate+ | otherwise = do+ summary <- action+ let aggregate' =+ Aggregate+ { totalFrames = totalFrames aggregate + parsedFrames summary+ , totalDropped = totalDropped aggregate + droppedBytes summary+ , combinedChecksums = combinedChecksums aggregate * 16777619 + checksum summary+ }+ evaluate (combinedChecksums aggregate')+ go (completed + 1) aggregate'++consumeChunks :: ParserAPI ParsedMessage -> [BS.ByteString] -> RunSummary+consumeChunks parserForScenario =+ go BS.empty (RunSummary 0 0 0)+ where+ go buffer summary remainingChunks =+ case remainingChunks of+ [] ->+ drain buffer summary []+ nextChunk:rest ->+ drain (buffer <> nextChunk) summary rest++ drain buffer summary remainingChunks+ | BS.null buffer =+ case remainingChunks of+ [] ->+ summary+ _ ->+ go buffer summary remainingChunks+ | otherwise =+ case parse parserForScenario buffer of+ Emit message rest+ | BS.length rest == BS.length buffer ->+ error "parser emitted without consuming input"+ | otherwise ->+ drain+ rest+ summary+ { parsedFrames = parsedFrames summary + 1+ , checksum = checksum summary `xor` fingerprint message+ }+ remainingChunks+ NeedMore ->+ case remainingChunks of+ [] ->+ error "parser requested more input after the corpus ended"+ _ ->+ go buffer summary remainingChunks+ DropPrefix bytesToDrop _+ | bytesToDrop <= 0 ->+ error "parser requested a non-positive prefix drop"+ | bytesToDrop > BS.length buffer ->+ error "parser requested dropping more bytes than available"+ | otherwise ->+ drain+ (BS.drop bytesToDrop buffer)+ summary { droppedBytes = droppedBytes summary + bytesToDrop }+ remainingChunks+ Reject reason ->+ error ("parser rejected benchmark corpus: " ++ reason)++fingerprint :: ParsedMessage -> Int+fingerprint parsedMessage =+ case parsedMessage of+ ParsedPing _ ->+ 1+ ParsedPong _ ->+ 2+ ParsedOk _ ->+ 3+ ParsedErr err ->+ 11 + BS.length (errReason err)+ ParsedInfo info ->+ 17+ + BS.length (server_id info)+ + BS.length (version info)+ + BS.length (go info)+ + BS.length (host info)+ + port info+ + max_payload info+ + proto info+ + fromMaybe 0 (client_id info)+ + maybe 0 BS.length (nonce info)+ + maybe 0 length (connect_urls info)+ ParsedMsg message ->+ 23+ + BS.length (subject message)+ + BS.length (sid message)+ + maybe 0 BS.length (replyTo message)+ + maybe 0 BS.length (payload message)+ + headersWeight (headers message)++headersWeight :: Maybe [(BS.ByteString, BS.ByteString)] -> Int+headersWeight Nothing = 0+headersWeight (Just pairs) =+ foldl' (\total (headerKey, headerValue) -> total + BS.length headerKey + BS.length headerValue) 0 pairs++errReason :: Err -> BS.ByteString+errReason err =+ case err of+ ErrUnknownOp reason -> reason+ ErrRoutePortConn reason -> reason+ ErrAuthViolation reason -> reason+ ErrAuthTimeout 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++snapshotStats :: IO RTSStats+snapshotStats = do+ enabled <- getRTSStatsEnabled+ if enabled+ then getRTSStats+ else die "RTS stats are not enabled; the benchmark must be built with -T"++diffStats :: RTSStats -> RTSStats -> RuntimeDelta+diffStats before after =+ RuntimeDelta+ { allocatedBytes = allocated_bytes after - allocated_bytes before+ , copiedBytes = copied_bytes after - copied_bytes before+ , peakLiveBytes = max_live_bytes after+ , peakMemBytes = max_mem_in_use_bytes after+ , gcCount = gcs after - gcs before+ , mutatorCpuNs = mutator_cpu_ns after - mutator_cpu_ns before+ , gcCpuNs = gc_cpu_ns after - gc_cpu_ns before+ }++formatMillis :: Int64 -> String+formatMillis nanoseconds =+ showFFloat (Just 3) (fromIntegral nanoseconds / 1e6 :: Double) "ms"++alignedScenario :: Int -> Scenario+alignedScenario scenarioScale =+ Scenario+ { name = "mixed-aligned"+ , description = "Mixed control and delivery traffic with each frame presented whole."+ , chunks = frames+ , frameCount = length frames+ , inputBytes = sum (map BS.length frames)+ }+ where+ frames = buildMixedFrames scenarioScale++chunkedScenario :: Int -> Scenario+chunkedScenario scenarioScale =+ Scenario+ { name = "mixed-chunked"+ , description = "The same mixed traffic split across irregular byte chunks."+ , chunks = chunkStream chunkPlan payload+ , frameCount = length frames+ , inputBytes = BS.length payload+ }+ where+ frames = buildMixedFrames scenarioScale+ payload = BS.concat frames++noisyScenario :: Int -> Scenario+noisyScenario scenarioScale =+ Scenario+ { name = "mixed-chunked-noisy"+ , description = "Chunked mixed traffic with occasional invalid prefix bytes to exercise resynchronization."+ , chunks = chunkStream chunkPlan payload+ , frameCount = length baseFrames+ , inputBytes = BS.length payload+ }+ where+ baseFrames = buildMixedFrames scenarioScale+ frames = injectNoise baseFrames+ payload = BS.concat frames++chunkPlan :: [Int]+chunkPlan = [17, 31, 64, 127, 257, 511, 1024, 41, 89, 2048, 13]++chunkStream :: [Int] -> BS.ByteString -> [BS.ByteString]+chunkStream plan = go (cycle plan)+ where+ go _ remaining+ | BS.null remaining =+ []+ go (nextChunkSize:rest) remaining =+ let (chunkBytes, remainingBytes) = BS.splitAt nextChunkSize remaining+ in chunkBytes : go rest remainingBytes+ go [] _ =+ error "chunk plan unexpectedly ended"++injectNoise :: [BS.ByteString] -> [BS.ByteString]+injectNoise =+ snd . foldl' inject (0 :: Int, [])+ where+ inject (index, acc) frame+ | index `mod` 29 == 0 =+ (index + 1, acc ++ ["Z", frame])+ | otherwise =+ (index + 1, acc ++ [frame])++buildMixedFrames :: Int -> [BS.ByteString]+buildMixedFrames scenarioScale =+ concatMap buildCycle [0 .. scenarioScale - 1]++buildCycle :: Int -> [BS.ByteString]+buildCycle cycleIndex =+ [ buildInfoFrame cycleIndex+ , "PING\r\n"+ , buildMsgFrame cycleIndex "bench.alpha" Nothing 0 Nothing+ , buildMsgFrame cycleIndex "bench.beta" (Just "_INBOX.reply") 32 Nothing+ , buildMsgFrame cycleIndex "bench.gamma" Nothing 128 Nothing+ , buildHMsgFrame cycleIndex "bench.delta" (Just "_INBOX.headers") 256 defaultHeaders+ , "PONG\r\n"+ , buildMsgFrame cycleIndex "bench.epsilon" Nothing 512 Nothing+ , "+OK\r\n"+ , buildMsgFrame cycleIndex "bench.zeta" (Just "_INBOX.zeta") 1024 Nothing+ , buildHMsgFrame cycleIndex "bench.eta" Nothing 1536 extendedHeaders+ , buildMsgFrame cycleIndex "bench.theta" Nothing 2048 Nothing+ , "-ERR 'Permissions Violation For Publish To bench.zeta.'\r\n"+ , buildMsgFrame cycleIndex "bench.iota" Nothing 64 Nothing+ ]++defaultHeaders :: [(BS.ByteString, BS.ByteString)]+defaultHeaders =+ [ ("Nats-Msg-Id", "msg-123456")+ , ("Trace-Id", "trace-a1b2c3d4")+ ]++extendedHeaders :: [(BS.ByteString, BS.ByteString)]+extendedHeaders =+ [ ("Nats-Msg-Id", "msg-654321")+ , ("Trace-Id", "trace-z9y8x7w6")+ , ("Content-Type", "application/json")+ , ("Client", "bench-runner")+ ]++buildInfoFrame :: Int -> BS.ByteString+buildInfoFrame cycleIndex =+ BS.concat+ [ "INFO {"+ , field "server_id" (quote "bench-server")+ , ","+ , field "version" (quote "2.10.22")+ , ","+ , field "go" (quote "1.23.0")+ , ","+ , field "host" (quote "127.0.0.1")+ , ","+ , field "port" (decimalBytes 4222)+ , ","+ , field "max_payload" (decimalBytes 1048576)+ , ","+ , field "proto" (decimalBytes 1)+ , ","+ , field "client_id" (decimalBytes (1000 + cycleIndex))+ , ","+ , field "nonce" (quote ("nonce-" <> decimalBytes cycleIndex))+ , ","+ , field "headers" "true"+ , "}\r\n"+ ]++field :: BS.ByteString -> BS.ByteString -> BS.ByteString+field key value = BS.concat [quote key, ":", value]++quote :: BS.ByteString -> BS.ByteString+quote bytes = BS.concat ["\"", bytes, "\""]++buildMsgFrame :: Int -> BS.ByteString -> Maybe BS.ByteString -> Int -> Maybe [(BS.ByteString, BS.ByteString)] -> BS.ByteString+buildMsgFrame cycleIndex subjectName replySubject payloadBytes maybeHeaders =+ case maybeHeaders of+ Nothing ->+ BS.concat+ [ "MSG "+ , subjectName+ , " "+ , sidBytes cycleIndex+ , " "+ , maybeReply replySubject+ , decimalBytes payloadBytes+ , "\r\n"+ , payloadOf cycleIndex payloadBytes+ , "\r\n"+ ]+ Just headersList ->+ buildHMsgFrame cycleIndex subjectName replySubject payloadBytes headersList++buildHMsgFrame :: Int -> BS.ByteString -> Maybe BS.ByteString -> Int -> [(BS.ByteString, BS.ByteString)] -> BS.ByteString+buildHMsgFrame cycleIndex subjectName replySubject payloadBytes headersList =+ BS.concat+ [ "HMSG "+ , subjectName+ , " "+ , sidBytes cycleIndex+ , " "+ , maybeReply replySubject+ , decimalBytes headerBytes+ , " "+ , decimalBytes (headerBytes + payloadBytes)+ , "\r\n"+ , headerBlock+ , payloadOf cycleIndex payloadBytes+ , "\r\n"+ ]+ where+ headerBlock = BS.concat ["NATS/1.0\r\n", renderHeaders headersList, "\r\n"]+ headerBytes = BS.length headerBlock++renderHeaders :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString+renderHeaders =+ foldMap (\(headerKey, headerValue) -> BS.concat [headerKey, ": ", headerValue, "\r\n"])++maybeReply :: Maybe BS.ByteString -> BS.ByteString+maybeReply Nothing = ""+maybeReply (Just replySubject) = BS.concat [replySubject, " "]++sidBytes :: Int -> BS.ByteString+sidBytes cycleIndex = decimalBytes (100000 + cycleIndex)++payloadOf :: Int -> Int -> BS.ByteString+payloadOf cycleIndex payloadBytes =+ BS.take payloadBytes (BS.concat (replicate repeats unit))+ where+ unit =+ BS.concat+ [ "payload-"+ , decimalBytes cycleIndex+ , "-abcdefghijklmnopqrstuvwxyz0123456789"+ ]+ repeats =+ max 1 ((payloadBytes `div` BS.length unit) + 1)++decimalBytes :: Int -> BS.ByteString+decimalBytes = B8.pack . show
+ client/API.hs view
@@ -0,0 +1,118 @@+-- | Capability record for the NATS client surface.+module API+ (+ Client (..)+ , MsgView (..)+ , PublishOption+ , SubscribeOption+ , withSubscriptionExpiry+ , withPayload+ , withReplyCallback+ , withHeaders+ ) 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.+ , 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 }++-- | 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) = (Just payload, callback, headers)++-- | 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) =+ (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)++-- | 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, _) = (payload, callback, Just headers)
+ client/Client.hs view
@@ -0,0 +1,440 @@+{-# LANGUAGE OverloadedStrings #-}++-- | High-level client implementation for NATS.+module Client+ ( newClient+ , ConfigOption+ , withConnectName+ , withEcho+ , withAuthToken+ , withUserPass+ , withNKey+ , withJWT+ , withTLSCert+ , withMinimumLogLevel+ , withLogAction+ , withConnectionAttempts+ , withCallbackConcurrency+ , withBufferLimit+ , withExitAction+ , LogLevel (..)+ , LogEntry (..)+ , renderLogEntry+ , AuthTokenData+ , UserPassData+ , NKeyData+ , JWTTokenData+ , TLSPublicKey+ , TLSPrivateKey+ , TLSCertData+ , ClientExitReason (..)+ ) 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 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+ , readStatus+ , runClient+ , setConnectName+ , waitForClosed+ , waitForNotRunning+ , waitForServerInfo+ )+import State.Types+ ( ClientConfig (..)+ , ClientExitReason (..)+ , ClientStatus (..)+ , TLSCertData+ , TLSPrivateKey+ , TLSPublicKey+ )+import Subscription.Store+ ( SubscriptionStore+ , awaitNoTrackedExpiries+ , hasTrackedExpiries+ , newSubscriptionStore+ , register+ , startExpiryWorker+ , startWorkers+ , unregister+ )+import Subscription.Types+ ( SubscribeConfig (SubscribeConfig)+ , SubscriptionMeta (SubscriptionMeta)+ )+import qualified Types.Connect as Connect+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++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 $+ waitForServerInfo clientState+ `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+ , 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++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) = do+ runClient client $+ logMessage Debug ("publishing to subject: " ++ show subject)+ replyTo <- case callback of+ Nothing ->+ pure Nothing+ Just replyCallback -> do+ inbox <- nextInbox client+ sid <- nextSid client+ let meta =+ SubscriptionMeta inbox Nothing True+ register store sid meta defaultSubscribeConfig replyCallback+ enqueue client $+ QueueItem+ Sub.Sub+ { Sub.subject = inbox+ , Sub.queueGroup = Nothing+ , Sub.sid = sid+ }+ enqueue client $+ QueueItem+ Unsub.Unsub+ { Unsub.sid = sid+ , Unsub.maxMsg = Just 1+ }+ pure (Just inbox)+ enqueue client $+ QueueItem+ Pub.Pub+ { Pub.subject = subject+ , Pub.payload = payload+ , Pub.replyTo = replyTo+ , Pub.headers = headers+ }++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 meta =+ SubscriptionMeta subject Nothing isReply+ register store sid meta cfg callback+ enqueue client $+ QueueItem+ Sub.Sub+ { Sub.subject = subject+ , Sub.queueGroup = Nothing+ , Sub.sid = sid+ }+ 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
+ internal/Lib/CallOption.hs view
@@ -0,0 +1,6 @@+module Lib.CallOption where++type CallOption a = (a -> a)++applyCallOptions :: [CallOption a] -> a -> a+applyCallOptions options value = foldr ($) value options
+ internal/Lib/Logger.hs view
@@ -0,0 +1,78 @@+module Lib.Logger+ ( module Lib.Logger.Types+ , defaultLogContext+ , newLogContext+ , updateLogContext+ , renderLogEntry+ , defaultLogger+ , withLogLock+ , loggerApi+ ) where++import Control.Concurrent.STM+import Control.Exception+import Control.Monad+import Control.Monad.Reader+import Data.Maybe (catMaybes, fromMaybe)+import Lib.Logger.Types+import Lib.LoggerAPI (LoggerAPI (LoggerAPI))++defaultLogContext :: LogContext+defaultLogContext = LogContext Nothing Nothing Nothing++newLogContext :: IO (TVar LogContext)+newLogContext = newTVarIO defaultLogContext++updateLogContext :: TVar LogContext -> (LogContext -> LogContext) -> IO ()+updateLogContext ctx f = atomically $ modifyTVar' ctx f++instance MonadWithLogger AppM where+ runWithLogger cfg ctx = flip runReaderT (LoggerEnv cfg ctx) . runAppM++instance MonadLogger AppM where+ logEntry entry = AppM . ReaderT $ \(LoggerEnv (LoggerConfig minLvl out lock) _) ->+ when (leLevel entry >= minLvl) (withLogLock lock (out entry))+ getLogContext = AppM . ReaderT $ \(LoggerEnv _ ctxVar) ->+ readTVarIO ctxVar+ logMessage lvl msg = AppM . ReaderT $ \(LoggerEnv (LoggerConfig minLvl out lock) ctxVar) ->+ when (lvl >= minLvl) $ do+ ctx <- readTVarIO ctxVar+ withLogLock lock $+ out LogEntry+ { leLevel = lvl+ , leMessage = msg+ , leClientId = lcClientId ctx+ , leConnectName = lcConnectName ctx+ , leServer = lcServer ctx+ }++renderLogEntry :: LogEntry -> String+renderLogEntry entry =+ let cidStr = maybe "unset" show (leClientId entry)+ cnStr = fromMaybe "unset" (leConnectName entry)+ label = "cid-" ++ cidStr ++ ":cn-" ++ cnStr+ ctxParts = catMaybes+ [ Just ("client=" ++ label)+ , fmap ("server=" ++) (leServer entry)+ ]+ ctxBlock = "[" ++ unwords ctxParts ++ "]"+ in "[" ++ levelTag (leLevel entry) ++ "] " ++ ctxBlock ++ " " ++ leMessage entry++levelTag :: LogLevel -> String+levelTag Debug = "D"+levelTag Info = "I"+levelTag Warn = "W"+levelTag Error = "E"+levelTag Fatal = "F"++defaultLogger :: IO LoggerConfig+defaultLogger = do+ lock <- newTMVarIO ()+ pure $ LoggerConfig Info (putStrLn . renderLogEntry) lock++withLogLock :: TMVar () -> IO a -> IO a+withLogLock lock =+ bracket_ (atomically $ takeTMVar lock) (atomically $ putTMVar lock ())++loggerApi :: LoggerAPI+loggerApi = LoggerAPI runWithLogger updateLogContext logMessage
+ internal/Lib/Logger/Types.hs view
@@ -0,0 +1,57 @@+{-# LANGUAGE GeneralizedNewtypeDeriving #-}++module Lib.Logger.Types+ ( MonadLogger (..)+ , MonadWithLogger (..)+ , LogLevel (..)+ , LoggerConfig (..)+ , LogContext (..)+ , LoggerEnv (..)+ , AppM (..)+ , LogEntry (..)+ ) where++import Control.Concurrent.STM (TMVar, TVar)+import Control.Monad.IO.Class (MonadIO)+import Control.Monad.Reader (MonadReader, ReaderT)++class MonadIO m => MonadLogger m where+ logEntry :: LogEntry -> m ()+ getLogContext :: m LogContext+ logMessage :: LogLevel -> String -> m ()++class (MonadLogger m, MonadIO m) => MonadWithLogger m where+ runWithLogger :: LoggerConfig -> TVar LogContext -> m a -> IO a++data LogLevel = Debug | Info | Warn | Error | Fatal+ deriving (Eq, Ord, Show)++data LoggerConfig = LoggerConfig+ { minLogLevel :: LogLevel+ , logFn :: LogEntry -> IO ()+ , logLock :: TMVar ()+ }++data LogContext = LogContext+ { lcClientId :: Maybe Int+ , lcConnectName :: Maybe String+ , lcServer :: Maybe String+ }+ deriving (Eq, Show)++data LoggerEnv = LoggerEnv+ { envConfig :: LoggerConfig+ , envContext :: TVar LogContext+ }++newtype AppM a = AppM { runAppM :: ReaderT LoggerEnv IO a }+ deriving (Applicative, Functor, Monad, MonadIO, MonadReader LoggerEnv)++data LogEntry = LogEntry+ { leLevel :: LogLevel+ , leMessage :: String+ , leClientId :: Maybe Int+ , leConnectName :: Maybe String+ , leServer :: Maybe String+ }+ deriving (Eq, Show)
+ internal/Lib/LoggerAPI.hs view
@@ -0,0 +1,15 @@+{-# LANGUAGE RankNTypes #-}++module Lib.LoggerAPI+ ( LoggerAPI (..)+ ) where++import Control.Concurrent.STM (TVar)+import Lib.Logger.Types++-- | Thin API wrapper for logger capabilities.+data LoggerAPI = LoggerAPI+ { runWithLogger :: forall a. LoggerConfig -> TVar LogContext -> AppM a -> IO a+ , updateLogContext :: TVar LogContext -> (LogContext -> LogContext) -> IO ()+ , logMessage :: LogLevel -> String -> AppM ()+ }
+ internal/Lib/WaitGroup.hs view
@@ -0,0 +1,36 @@+module WaitGroup where++import Control.Concurrent (ThreadId, forkIO)+import Control.Concurrent.STM+import Control.Exception (finally)++data WaitGroup = WaitGroup+ { count :: TVar Int+ , lock :: TMVar ()+ }++newWaitGroup :: Int -> IO WaitGroup+newWaitGroup n = do+ c <- newTVarIO n+ WaitGroup c <$> newEmptyTMVarIO++add :: WaitGroup -> Int -> IO ()+add wg n = atomically $ modifyTVar' (count wg) (+n)++done :: WaitGroup -> IO ()+done wg = atomically $ do+ -- decrement the count+ c <- readTVar (count wg)+ case c of+ 0 -> return ()+ 1 -> modifyTVar' (count wg) (subtract 1) >> putTMVar (lock wg) ()+ _ -> modifyTVar' (count wg) (subtract 1)++wait :: WaitGroup -> IO ()+wait wg = atomically $ do+ -- wait for the lock to be released, then replace it+ l <- takeTMVar (lock wg)+ putTMVar (lock wg) l++forkWaitGroup :: WaitGroup -> IO () -> IO ThreadId+forkWaitGroup wg action = forkIO (action `finally` done wg)
+ internal/Lib/WorkerPool.hs view
@@ -0,0 +1,41 @@+module Lib.WorkerPool+ ( startWorkerPool+ ) where++import Control.Concurrent (forkIO)+import Control.Concurrent.STM+ ( STM+ , TQueue+ , atomically+ , check+ , isEmptyTQueue+ , orElse+ , readTQueue+ )+import Control.Exception (SomeException, try)+import Control.Monad (replicateM_, void)++startWorkerPool :: Int -> TQueue (IO ()) -> STM () -> (SomeException -> IO ()) -> IO ()+startWorkerPool concurrency queue stopSignal onError = do+ let workerCount = max 1 concurrency+ replicateM_ workerCount (void (forkIO worker))+ where+ worker = do+ let loop = do+ action <- atomically $+ (Just <$> readTQueue queue)+ `orElse`+ (do+ stopSignal+ empty <- isEmptyTQueue queue+ check empty+ return Nothing)+ case action of+ Nothing -> return ()+ Just job -> do+ result <- (try job :: IO (Either SomeException ()))+ case result of+ Left err -> onError err+ Right () -> pure ()+ loop+ loop
+ internal/Plumbing/Network/Connection.hs view
@@ -0,0 +1,29 @@+module Network.Connection+ ( connectionApi+ ) where++import qualified Network.Connection.Core as Core+import qualified Network.Connection.Tcp as Tcp+import qualified Network.Connection.Tls as Tls+import Network.ConnectionAPI++connectionApi :: ConnectionAPI+connectionApi =+ ConnectionAPI+ { newConn = Core.newConn+ , reader = ReaderAPI+ { readData = Core.readData+ , closeReader = Core.closeReader+ , openReader = Core.openReader+ }+ , writer = WriterAPI+ { writeData = Core.writeData+ , writeDataLazy = Core.writeDataLazy+ , closeWriter = Core.closeWriter+ , openWriter = Core.openWriter+ }+ , open = Core.openConn+ , close = Core.closeConn+ , connectTcp = Tcp.connectTcp+ , configure = Tls.configureTransport+ }
+ internal/Plumbing/Network/Connection/Core.hs view
@@ -0,0 +1,229 @@+{-# LANGUAGE TypeApplications #-}++module Network.Connection.Core+ ( ReadError+ , WriteError+ , Transport (..)+ , TransportOption (..)+ , Conn+ , newConn+ , readData+ , closeReader+ , openReader+ , writeData+ , writeDataLazy+ , closeWriter+ , openWriter+ , openConn+ , closeConn+ , pointTransport+ , currentTransport+ , bufferRead+ , enableReadWorker+ ) where++import Control.Concurrent (forkIO)+import Control.Concurrent.STM+import Control.Exception+import Control.Monad+import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Data.Maybe (isJust, isNothing)+import Network.Connection.Types++newConn :: IO Conn+newConn =+ Conn+ <$> newEmptyTMVarIO+ <*> newEmptyTMVarIO+ <*> newEmptyTMVarIO+ <*> newTVarIO mempty+ <*> newTBQueueIO 1000+ <*> newTVarIO False+ <*> newTVarIO False++readChunkSize :: Int+readChunkSize = 4096++startReadWorker :: Conn -> IO ()+startReadWorker conn = do+ shouldStart <- atomically $ do+ enabled <- readTVar (readWorkerEnabled conn)+ running <- readTVar (readWorkerRunning conn)+ if not enabled || running+ then return False+ else do+ writeTVar (readWorkerRunning conn) True+ return True+ when shouldStart . void $ forkIO (readWorkerLoop conn)++readWorkerLoop :: Conn -> IO ()+readWorkerLoop conn = do+ let cleanup = atomically $ writeTVar (readWorkerRunning conn) False+ finally loop cleanup+ where+ loop = do+ (blocked, enabled) <- atomically $ do+ blocked <- tryReadTMVar (readBlock conn)+ enabled <- readTVar (readWorkerEnabled conn)+ return (blocked, enabled)+ case (blocked, enabled) of+ (Just _, _) -> return ()+ (_, False) -> return ()+ _ -> do+ current <- currentTransport conn+ case current of+ Nothing -> return ()+ Just currentTransport' -> do+ result <- try @SomeException (transportRead currentTransport' readChunkSize)+ shouldContinue <- enqueueReadResult conn result+ case result of+ Left _ -> return ()+ Right _ -> when shouldContinue loop++enqueueReadResult :: Conn -> Either SomeException ByteString -> IO Bool+enqueueReadResult conn result = atomically $+ (do+ blocked <- tryReadTMVar (readBlock conn)+ check (isNothing blocked)+ writeTBQueue (readQueue conn) (either (Left . show) Right result)+ return True)+ `orElse`+ (do+ _ <- readTMVar (readBlock conn)+ return False)++enableReadWorker :: Conn -> IO ()+enableReadWorker conn = atomically $ writeTVar (readWorkerEnabled conn) True++readData :: Conn -> Int -> IO (Either ReadError ByteString)+readData conn n = do+ blocked <- atomically $ tryReadTMVar (readBlock conn)+ if isJust blocked+ then return $ Left "Read operation is blocked"+ else do+ buffered <- atomically $ do+ buf <- readTVar (readBuffer conn)+ if BS.null buf+ then return Nothing+ else do+ let (chunk, rest) = BS.splitAt n buf+ writeTVar (readBuffer conn) rest+ return (Just chunk)+ case buffered of+ Just chunk -> return $ Right chunk+ Nothing -> do+ current <- currentTransport conn+ case current of+ Nothing -> return $ Left "Transport not initialized"+ Just currentTransport' -> do+ enabled <- readTVarIO (readWorkerEnabled conn)+ if enabled+ then do+ startReadWorker conn+ result <- atomically $+ (Left "Read operation is blocked" <$ readTMVar (readBlock conn))+ `orElse`+ readTBQueue (readQueue conn)+ case result of+ Left err -> return $ Left err+ Right bytes -> do+ let (chunk, rest) = BS.splitAt n bytes+ unless (BS.null rest)+ (atomically $ modifyTVar' (readBuffer conn) (<> rest))+ return $ Right chunk+ else do+ resultVar <- newEmptyTMVarIO+ _ <- forkIO $ do+ result <- try @SomeException (transportRead currentTransport' n)+ atomically . void $ tryPutTMVar resultVar result+ result <- atomically $+ (Left "Read operation is blocked" <$ readTMVar (readBlock conn))+ `orElse`+ (Right <$> readTMVar resultVar)+ case result of+ Left err -> return $ Left err+ Right (Left err) -> return $ Left (show err)+ Right (Right bytes) -> do+ let (chunk, rest) = BS.splitAt n bytes+ unless (BS.null rest)+ (atomically $ modifyTVar' (readBuffer conn) (<> rest))+ return $ Right chunk++closeReader :: Conn -> IO ()+closeReader conn = void . atomically $ tryPutTMVar (readBlock conn) ()++openReader :: Conn -> IO ()+openReader conn = void . atomically $ tryTakeTMVar (readBlock conn)++writeData :: Conn -> ByteString -> IO (Either WriteError ())+writeData conn bytes = do+ blocked <- atomically $ tryReadTMVar (writeBlock conn)+ if isJust blocked+ then return $ Left "Write operation is blocked"+ else do+ current <- currentTransport conn+ case current of+ Nothing -> return $ Left "Transport not initialized"+ Just currentTransport' -> do+ result <- try @SomeException $ do+ transportWrite currentTransport' bytes+ transportFlush currentTransport'+ case result of+ Left err -> return $ Left (show err)+ Right _ -> return $ Right ()++writeDataLazy :: Conn -> LBS.ByteString -> IO (Either WriteError ())+writeDataLazy conn bytes = do+ blocked <- atomically $ tryReadTMVar (writeBlock conn)+ if isJust blocked+ then return $ Left "Write operation is blocked"+ else do+ current <- currentTransport conn+ case current of+ Nothing -> return $ Left "Transport not initialized"+ Just currentTransport' -> do+ result <- try @SomeException $ do+ transportWriteLazy currentTransport' bytes+ transportFlush currentTransport'+ case result of+ Left err -> return $ Left (show err)+ Right _ -> return $ Right ()++closeWriter :: Conn -> IO ()+closeWriter conn = void . atomically $ tryPutTMVar (writeBlock conn) ()++openWriter :: Conn -> IO ()+openWriter conn = void . atomically $ tryTakeTMVar (writeBlock conn)++pointTransport :: Conn -> Transport -> IO ()+pointTransport conn newTransport = atomically $ do+ _ <- tryTakeTMVar (transport conn)+ putTMVar (transport conn) newTransport++currentTransport :: Conn -> IO (Maybe Transport)+currentTransport conn = atomically $ tryReadTMVar (transport conn)++closeConn :: Conn -> IO ()+closeConn conn = do+ closeReader conn+ closeWriter conn+ current <- currentTransport conn+ case current of+ Nothing -> return ()+ Just currentTransport' -> void $ try @SomeException (transportClose currentTransport')++openConn :: Conn -> IO ()+openConn conn = do+ openReader conn+ openWriter conn+ atomically $ do+ writeTVar (readBuffer conn) mempty+ void $ flushTBQueue (readQueue conn)+ writeTVar (readWorkerEnabled conn) False++bufferRead :: Conn -> ByteString -> IO ()+bufferRead conn bytes =+ unless (BS.null bytes)+ (atomically $ modifyTVar' (readBuffer conn) (bytes <>))
+ internal/Plumbing/Network/Connection/Tcp.hs view
@@ -0,0 +1,45 @@+{-# LANGUAGE TypeApplications #-}++module Network.Connection.Tcp+ ( connectTcp+ ) where++import Control.Exception+import qualified Data.ByteString.Lazy as LBS+import Network.Connection.Core (pointTransport)+import qualified Network.Connection.Tls as Tls+import Network.Connection.Types (Conn, Transport (..))+import qualified Network.Simple.TCP as TCP+import qualified Network.Socket as NS+import qualified Network.Socket.ByteString as NSB++tcpTransport :: NS.Socket -> Transport+tcpTransport sock =+ Transport+ { transportRead = NSB.recv sock+ , transportWrite = NSB.sendAll sock+ , transportWriteLazy = NSB.sendMany sock . LBS.toChunks+ , transportFlush = pure ()+ , transportClose = NS.close sock+ , transportUpgrade = Just (Tls.upgradeTcp sock)+ }++openTcpTransport :: String -> Int -> IO (Either String Transport)+openTcpTransport host port = do+ result <- try @SomeException $ do+ (sock, _) <- TCP.connectSock host (show port)+ NS.setSocketOption sock NS.NoDelay 1+ NS.setSocketOption sock NS.Cork 0+ pure (tcpTransport sock)+ case result of+ Left err -> return (Left (show err))+ Right transport -> return (Right transport)++connectTcp :: Conn -> String -> Int -> IO (Either String ())+connectTcp conn host port = do+ result <- openTcpTransport host port+ case result of+ Left err -> return (Left err)+ Right transport -> do+ pointTransport conn transport+ return (Right ())
+ internal/Plumbing/Network/Connection/Tls.hs view
@@ -0,0 +1,110 @@+{-# LANGUAGE TypeApplications #-}++module Network.Connection.Tls+ ( configureTransport+ , upgradeTcp+ ) where++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 Network.Connection.Core+ ( bufferRead+ , currentTransport+ , enableReadWorker+ , pointTransport+ )+import Network.Connection.Types+ ( Conn+ , Transport (..)+ , TransportOption (..)+ )+import qualified Network.Socket as NS+import qualified Network.Socket.ByteString as NSB+import qualified Network.TLS as TLS++tlsTransport :: TLS.Context -> Transport+tlsTransport ctx =+ Transport+ { transportRead = const (TLS.recvData ctx)+ , transportWrite = TLS.sendData ctx . LBS.fromStrict+ , transportWriteLazy = TLS.sendData ctx+ , transportFlush = TLS.contextFlush ctx+ , transportClose = do+ void $ try @SomeException (TLS.bye ctx)+ void $ try @SomeException (TLS.contextClose ctx)+ , transportUpgrade = Nothing+ }++upgradeTcp :: NS.Socket -> TLS.ClientParams -> IO (Either String Transport)+upgradeTcp sock params = do+ let backend = TLS.Backend+ { TLS.backendSend = NSB.sendAll sock+ , TLS.backendRecv = NSB.recv sock+ , TLS.backendFlush = pure ()+ , TLS.backendClose = NS.close sock+ }+ result <- try @SomeException $ do+ ctx <- TLS.contextNew backend params+ TLS.handshake ctx+ pure (tlsTransport ctx)+ case result of+ Left err -> return $ Left (show err)+ Right transport -> return $ Right transport++upgradeToTLS :: Conn -> TLS.ClientParams -> IO (Either String ())+upgradeToTLS conn params = do+ current <- currentTransport conn+ case current of+ Nothing -> return $ Left "Transport not initialized"+ Just currentTransport' ->+ case transportUpgrade currentTransport' of+ Nothing -> return $ Right ()+ Just upgrade -> do+ result <- upgrade params+ case result of+ Left err -> return $ Left err+ Right newTransport -> do+ pointTransport conn newTransport+ return $ Right ()++upgradeToTLSWithConfig :: Conn -> String -> Maybe (ByteString, ByteString) -> IO (Either String ())+upgradeToTLSWithConfig conn host tlsConfig = do+ paramsResult <- buildTlsParams host tlsConfig+ case paramsResult of+ Left err -> return (Left err)+ Right params -> upgradeToTLS conn params++configureTransport :: Conn -> TransportOption -> IO (Either String ())+configureTransport conn transportOption = do+ 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 ())+ else do+ bufferRead conn (transportInitialBytes transportOption)+ enableReadWorker conn+ return (Right ())++buildTlsParams :: String -> Maybe (ByteString, ByteString) -> 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 })
+ internal/Plumbing/Network/Connection/Types.hs view
@@ -0,0 +1,42 @@+module Network.Connection.Types+ ( ReadError+ , WriteError+ , Transport (..)+ , TransportOption (..)+ , Conn (..)+ ) where++import Control.Concurrent.STM+import Data.ByteString (ByteString)+import qualified Data.ByteString.Lazy as LBS+import qualified Network.TLS as TLS++type ReadError = String+type WriteError = String++data Transport = Transport+ { transportRead :: Int -> IO ByteString+ , transportWrite :: ByteString -> IO ()+ , transportWriteLazy :: LBS.ByteString -> IO ()+ , transportFlush :: IO ()+ , transportClose :: IO ()+ , transportUpgrade :: Maybe (TLS.ClientParams -> IO (Either String Transport))+ }++data TransportOption = TransportOption+ { transportHost :: String+ , transportTlsRequired :: Bool+ , transportTlsRequested :: Bool+ , transportTlsCert :: Maybe (ByteString, ByteString)+ , transportInitialBytes :: ByteString+ }++data Conn = Conn+ { transport :: TMVar Transport+ , readBlock :: TMVar ()+ , writeBlock :: TMVar ()+ , readBuffer :: TVar ByteString+ , readQueue :: TBQueue (Either ReadError ByteString)+ , readWorkerRunning :: TVar Bool+ , readWorkerEnabled :: TVar Bool+ }
+ internal/Plumbing/Network/ConnectionAPI.hs view
@@ -0,0 +1,41 @@+module Network.ConnectionAPI+ ( ReadError+ , WriteError+ , ReaderAPI (..)+ , WriterAPI (..)+ , Conn+ , TransportOption (..)+ , ConnectionAPI (..)+ ) where++import Data.ByteString (ByteString)+import qualified Data.ByteString.Lazy as LBS+import Network.Connection.Types+ ( Conn+ , ReadError+ , TransportOption (..)+ , WriteError+ )++data ReaderAPI reader = ReaderAPI+ { readData :: reader -> Int -> IO (Either ReadError ByteString)+ , closeReader :: reader -> IO ()+ , openReader :: reader -> IO ()+ }++data WriterAPI writer = WriterAPI+ { writeData :: writer -> ByteString -> IO (Either WriteError ())+ , writeDataLazy :: writer -> LBS.ByteString -> IO (Either WriteError ())+ , closeWriter :: writer -> IO ()+ , openWriter :: writer -> IO ()+ }++data ConnectionAPI = ConnectionAPI+ { newConn :: IO Conn+ , reader :: ReaderAPI Conn+ , writer :: WriterAPI Conn+ , open :: Conn -> IO ()+ , close :: Conn -> IO ()+ , connectTcp :: Conn -> String -> Int -> IO (Either String ())+ , configure :: Conn -> TransportOption -> IO (Either String ())+ }
+ internal/Plumbing/Parser/API.hs view
@@ -0,0 +1,29 @@+module Parser.API+ ( ParseStep (..)+ , ParsedMessage (..)+ , ParserAPI (..)+ ) where++import Data.ByteString (ByteString)+import Types.Err (Err)+import Types.Info (Info)+import Types.Msg (Msg)+import Types.Ok (Ok)+import Types.Ping (Ping)+import Types.Pong (Pong)++data ParseStep a = Emit a ByteString+ | NeedMore+ | DropPrefix Int String+ | Reject String+ deriving (Eq, Show)++data ParsedMessage = ParsedPing Ping+ | ParsedPong Pong+ | ParsedOk Ok+ | ParsedErr Err+ | ParsedInfo Info+ | ParsedMsg Msg+ deriving (Eq, Show)++newtype ParserAPI a = ParserAPI { parse :: ByteString -> ParseStep a }
+ internal/Plumbing/Parser/Attoparsec.hs view
@@ -0,0 +1,292 @@+{-# LANGUAGE OverloadedStrings #-}++module Parser.Attoparsec+ ( parserApi+ ) where++import Control.Applicative ((<|>))+import qualified Data.Aeson as Aeson+import qualified Data.Attoparsec.ByteString as A+import qualified Data.Attoparsec.ByteString.Char8 as AC+import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as B8+import qualified Data.ByteString.Lazy as BSL+import Data.Word (Word8)+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore)+ , ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong)+ , ParserAPI (ParserAPI)+ )+import Types.Err+ ( Err (ErrAuthTimeout, ErrAuthViolation, ErrErr, ErrInvalidProtocol, ErrInvalidSubject, ErrMaxConnsEx, ErrMaxControlLineEx, ErrMaxPayload, ErrPermViolation, ErrRoutePortConn, ErrSlowConsumer, ErrStaleConn, ErrTlsRequired, ErrUnknownOp)+ )+import Types.Msg (Msg (Msg))+import Types.Ok (Ok (Ok))+import Types.Ping (Ping (Ping))+import Types.Pong (Pong (Pong))++parserApi :: ParserAPI ParsedMessage+parserApi = ParserAPI parseStep++parseStep :: BS.ByteString -> ParseStep ParsedMessage+parseStep bytes =+ case A.parse parsedMessageParser bytes of+ A.Done rest parsedMessage ->+ Emit parsedMessage rest+ A.Partial _ ->+ NeedMore+ A.Fail rest _ reason ->+ DropPrefix (dropBytes bytes rest) reason++parsedMessageParser :: A.Parser ParsedMessage+parsedMessageParser = do+ prefix <- A.peekWord8'+ case prefix of+ 43 ->+ okParser+ 45 ->+ errParser+ 72 ->+ hmsgParser+ 73 ->+ infoParser+ 77 ->+ msgParser+ 80 ->+ pingOrPongParser+ _ ->+ fail ("unknown protocol prefix: " ++ show prefix)++okParser :: A.Parser ParsedMessage+okParser = do+ _ <- A.string "+OK\r\n"+ pure (ParsedOk Ok)++errParser :: A.Parser ParsedMessage+errParser = do+ _ <- A.string "-ERR"+ skipHorizontalSpace1+ _ <- A.word8 singleQuote+ reason <- A.takeTill (== singleQuote)+ _ <- A.word8 singleQuote+ _ <- A.string "\r\n"+ case classifyErr reason of+ Left parseReason ->+ fail parseReason+ Right err ->+ pure (ParsedErr err)++infoParser :: A.Parser ParsedMessage+infoParser = do+ _ <- A.string "INFO"+ skipHorizontalSpace1+ rawInfo <- lineParser+ case Aeson.eitherDecode (BSL.fromStrict rawInfo) of+ Left parseReason ->+ fail ("invalid INFO frame: " ++ parseReason)+ Right info ->+ pure (ParsedInfo info)++msgParser :: A.Parser ParsedMessage+msgParser = do+ _ <- A.string "MSG"+ skipHorizontalSpace1+ controlLine <- lineParser+ case B8.words controlLine of+ [subjectName, sidValue, payloadSizeBytes] ->+ parseMsgPayload subjectName sidValue Nothing payloadSizeBytes+ [subjectName, sidValue, replySubject, payloadSizeBytes] ->+ parseMsgPayload subjectName sidValue (Just replySubject) payloadSizeBytes+ _ ->+ fail ("invalid MSG control line: " ++ B8.unpack controlLine)++hmsgParser :: A.Parser ParsedMessage+hmsgParser = do+ _ <- A.string "HMSG"+ skipHorizontalSpace1+ controlLine <- lineParser+ case B8.words controlLine of+ [subjectName, sidValue, headerSizeBytes, totalSizeBytes] ->+ parseHMsgPayload subjectName sidValue Nothing headerSizeBytes totalSizeBytes+ [subjectName, sidValue, replySubject, headerSizeBytes, totalSizeBytes] ->+ parseHMsgPayload subjectName sidValue (Just replySubject) headerSizeBytes totalSizeBytes+ _ ->+ fail ("invalid HMSG control line: " ++ B8.unpack controlLine)++pingOrPongParser :: A.Parser ParsedMessage+pingOrPongParser =+ pingParser <|> pongParser++pingParser :: A.Parser ParsedMessage+pingParser = do+ _ <- A.string "PING\r\n"+ pure (ParsedPing Ping)++pongParser :: A.Parser ParsedMessage+pongParser = do+ _ <- 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 =+ 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+ )+ )++parseHMsgPayload :: BS.ByteString -> BS.ByteString -> Maybe BS.ByteString -> BS.ByteString -> BS.ByteString -> A.Parser ParsedMessage+parseHMsgPayload subjectName sidValue replySubject headerSizeBytes totalSizeBytes =+ case (integerBytes "HMSG header size" headerSizeBytes, integerBytes "HMSG total size" totalSizeBytes) of+ (Left parseReason, _) ->+ fail parseReason+ (_, Left parseReason) ->+ fail parseReason+ (Right headerSize, Right totalSize)+ | totalSize < headerSize ->+ fail+ ( "invalid HMSG sizes: total size "+ ++ show totalSize+ ++ " is smaller than header size "+ ++ show headerSize+ )+ | otherwise -> do+ headerBytes <- A.take headerSize+ payloadBytes <- A.take (totalSize - headerSize)+ _ <- A.string "\r\n"+ case A.parseOnly (headerBlockParser <* A.endOfInput) headerBytes of+ Left parseReason ->+ fail ("invalid HMSG headers: " ++ parseReason)+ Right parsedHeaders ->+ pure+ ( ParsedMsg+ ( Msg+ subjectName+ sidValue+ replySubject+ (Just payloadBytes)+ (Just parsedHeaders)+ )+ )++headerBlockParser :: A.Parser [(BS.ByteString, BS.ByteString)]+headerBlockParser = do+ _ <- A.string "NATS/"+ _ <- lineParser+ headerPairsParser []++headerPairsParser :: [(BS.ByteString, BS.ByteString)] -> A.Parser [(BS.ByteString, BS.ByteString)]+headerPairsParser reversedPairs = do+ nextByte <- A.peekWord8'+ if nextByte == carriageReturn+ then do+ _ <- A.string "\r\n"+ pure (reverse reversedPairs)+ else do+ headerKey <- A.takeTill (== colon)+ _ <- A.word8 colon+ headerValue <- lineParser+ headerPairsParser ((stripHorizontalSpace headerKey, stripHorizontalSpace headerValue) : reversedPairs)++classifyErr :: BS.ByteString -> Either String Err+classifyErr reason+ | reason == "Unknown Protocol Operation" =+ Right (ErrUnknownOp reason)+ | reason == "Attempted To Connect To Route Port" =+ Right (ErrRoutePortConn reason)+ | reason == "Authorization Violation" =+ Right (ErrAuthViolation reason)+ | reason == "Authorization Timeout" =+ Right (ErrAuthTimeout reason)+ | reason == "Invalid Client Protocol" =+ Right (ErrInvalidProtocol reason)+ | reason == "Maximum Control Line Exceeded" =+ Right (ErrMaxControlLineEx reason)+ | reason == "Parser Error" =+ Right (ErrErr reason)+ | reason == "Secure Connection - TLS Required" =+ Right (ErrTlsRequired reason)+ | reason == "Stale Connection" =+ Right (ErrStaleConn reason)+ | reason == "Maximum Connections Exceeded" =+ Right (ErrMaxConnsEx reason)+ | reason == "Slow Consumer" =+ Right (ErrSlowConsumer reason)+ | reason == "Maximum Payload Violation" =+ Right (ErrMaxPayload reason)+ | reason == "Invalid Subject" =+ Right (ErrInvalidSubject reason)+ | "Permissions Violation" `BS.isPrefixOf` reason =+ Right (ErrPermViolation reason)+ | otherwise =+ Left ("unknown -ERR reason: " ++ B8.unpack reason)++lineParser :: A.Parser BS.ByteString+lineParser = do+ lineBytes <- A.takeTill (== carriageReturn)+ _ <- A.string "\r\n"+ pure lineBytes++skipHorizontalSpace1 :: A.Parser ()+skipHorizontalSpace1 = do+ _ <- A.satisfy isHorizontalSpaceByte+ A.skipWhile isHorizontalSpaceByte++integerBytes :: String -> BS.ByteString -> Either String Int+integerBytes label digits =+ case A.parseOnly (AC.decimal <* A.endOfInput) digits of+ Left _ ->+ Left (label ++ " is not an integer: " ++ B8.unpack digits)+ Right value ->+ Right value++dropBytes :: BS.ByteString -> BS.ByteString -> Int+dropBytes input remaining =+ min (BS.length input) (consumedBytes + 1)+ where+ consumedBytes = BS.length input - BS.length remaining++nonEmpty :: BS.ByteString -> Maybe BS.ByteString+nonEmpty bytes+ | BS.null bytes =+ Nothing+ | otherwise =+ Just bytes++stripHorizontalSpace :: BS.ByteString -> BS.ByteString+stripHorizontalSpace =+ BS.dropWhile isHorizontalSpaceByte . dropWhileEnd isHorizontalSpaceByte++dropWhileEnd :: (Word8 -> Bool) -> BS.ByteString -> BS.ByteString+dropWhileEnd predicate =+ BS.reverse . BS.dropWhile predicate . BS.reverse++isHorizontalSpaceByte :: Word8 -> Bool+isHorizontalSpaceByte word8Value =+ word8Value == space || word8Value == horizontalTab++carriageReturn :: Word8+carriageReturn = 13++singleQuote :: Word8+singleQuote = 39++colon :: Word8+colon = 58++space :: Word8+space = 32++horizontalTab :: Word8+horizontalTab = 9
+ internal/Plumbing/Pipeline/Broadcasting.hs view
@@ -0,0 +1,22 @@+module Pipeline.Broadcasting+ ( broadcastingApi+ ) where++import Conduit+import Lib.Logger.Types (MonadLogger)+import Network.ConnectionAPI (WriterAPI)+import Pipeline.Broadcasting.API (BroadcastingAPI (..))+import Pipeline.Broadcasting.Sink+import Pipeline.Broadcasting.Source+import Pipeline.Broadcasting.Transformer+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++broadcastingApi :: BroadcastingAPI+broadcastingApi = BroadcastingAPI+ { run = runBroadcasting+ }
+ internal/Plumbing/Pipeline/Broadcasting/API.hs view
@@ -0,0 +1,12 @@+{-# LANGUAGE RankNTypes #-}++module Pipeline.Broadcasting.API+ ( BroadcastingAPI (..)+ ) where++import Control.Monad.IO.Class (MonadIO)+import Lib.Logger.Types (MonadLogger)+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 () }
+ internal/Plumbing/Pipeline/Broadcasting/Sink.hs view
@@ -0,0 +1,24 @@+module Pipeline.Broadcasting.Sink where++import Conduit+import qualified Data.ByteString.Lazy as LBS+import Lib.Logger.Types (LogLevel (..), MonadLogger (..))+import Network.ConnectionAPI (WriterAPI (..))++sink :: (MonadLogger m , MonadIO m)+ => WriterAPI writer+ -> writer+ -> ConduitT LBS.ByteString Void m ()+sink writerApi writer = do+ bs <- await+ case bs of+ Nothing -> do+ lift . logMessage Debug $ "no more data to write; stopping sink"+ return ()+ Just bs -> do+ res <- liftIO $ writeDataLazy writerApi writer bs+ case res of+ Left err -> do+ lift . logMessage Error $ ("write failed: " ++ err)+ return ()+ Right () -> sink writerApi writer
+ internal/Plumbing/Pipeline/Broadcasting/Source.hs view
@@ -0,0 +1,16 @@+module Pipeline.Broadcasting.Source where+import Conduit+import Lib.Logger.Types (LogLevel (..), MonadLogger (..))+import Queue.API (Queue, QueueItem, dequeue)++source :: (MonadLogger m , MonadIO m)+ => Queue+ -> ConduitT () QueueItem m ()+source q = do+ x <- liftIO (dequeue q)+ case x of+ Left err -> do+ lift . logMessage Info $ ("dequeue failed: " ++ err)+ Right i -> do+ yield i+ source q
+ internal/Plumbing/Pipeline/Broadcasting/Transformer.hs view
@@ -0,0 +1,18 @@+module Pipeline.Broadcasting.Transformer where++import Conduit+import qualified Data.ByteString.Lazy as LBS+import Lib.Logger.Types (LogLevel (..), 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
+ internal/Plumbing/Pipeline/Streaming.hs view
@@ -0,0 +1,27 @@+module Pipeline.Streaming+ ( streamingApi+ ) where++import Conduit+import Lib.Logger.Types (MonadLogger)+import Network.ConnectionAPI (ReaderAPI)+import Parser.API (ParserAPI)+import Pipeline.Streaming.API (StreamingAPI (..))+import Pipeline.Streaming.Parser+import Pipeline.Streaming.Sink+import Pipeline.Streaming.Source++runStreaming :: (MonadLogger m , MonadIO m)+ => Int+ -> ReaderAPI reader+ -> reader+ -> ParserAPI a+ -> (a -> IO ())+ -> m ()+runStreaming bufferLimit readerApi reader parserApi router =+ runConduit $ source readerApi reader .| parser bufferLimit parserApi .| sink router++streamingApi :: StreamingAPI+streamingApi = StreamingAPI+ { run = runStreaming+ }
+ internal/Plumbing/Pipeline/Streaming/API.hs view
@@ -0,0 +1,12 @@+{-# LANGUAGE RankNTypes #-}++module Pipeline.Streaming.API+ ( StreamingAPI (..)+ ) where++import Control.Monad.IO.Class (MonadIO)+import Lib.Logger.Types (MonadLogger)+import Network.ConnectionAPI (ReaderAPI)+import Parser.API (ParserAPI)++newtype StreamingAPI = StreamingAPI { run :: forall m reader a. (MonadLogger m, MonadIO m) => Int -> ReaderAPI reader -> reader -> ParserAPI a -> (a -> IO ()) -> m () }
+ internal/Plumbing/Pipeline/Streaming/Parser.hs view
@@ -0,0 +1,63 @@+module Pipeline.Streaming.Parser where+import Conduit+import Data.ByteString+import qualified Data.ByteString as BS+import Lib.Logger.Types (LogLevel (..), MonadLogger (..))+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore, Reject)+ , ParserAPI+ , parse+ )+import Prelude hiding (drop, length, take)++parser :: (MonadLogger m , MonadIO m)+ => Int+ -> ParserAPI result+ -> ConduitT ByteString result m ()+parser bufferLimit parserApi = loop empty+ where+ loop acc = do+ bs <- await+ case bs of+ Nothing -> return ()+ Just chunk -> handleChunk $ append acc chunk+ handleChunk bs+ | BS.null bs = parser bufferLimit parserApi+ | otherwise = do+ lift . logMessage Debug $ "parsing chunk"+ case parse parserApi bs of+ NeedMore -> do+ let bsLen = length bs+ if bsLen > bufferLimit+ then do+ lift . logMessage Error $+ "overloaded buffer: "+ ++ show bsLen+ ++ " bytes exceeds limit "+ ++ show bufferLimit+ lift . logMessage Debug $ ("invalid prefix: " ++ show (take bufferLimit bs))+ handleChunk (drop bufferLimit bs)+ else do+ lift . logMessage Debug $ "message spans frame, waiting for more data"+ loop bs+ DropPrefix n reason -> do+ lift . logMessage Error $ ("dropping invalid prefix: " ++ reason)+ lift . logMessage Debug $ ("invalid prefix: " ++ show bs)+ lift . logMessage Error $ ("dropping " ++ show n ++ " bytes")+ handleChunk (drop n bs)+ Reject reason ->+ 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
+ internal/Plumbing/Pipeline/Streaming/Sink.hs view
@@ -0,0 +1,9 @@+module Pipeline.Streaming.Sink where+import Conduit+import Lib.Logger.Types (LogLevel (..), MonadLogger (..))++sink :: (MonadIO m, MonadLogger m) => (a -> IO ()) -> ConduitT a Void m ()+sink action = do+ awaitForever $ \ma -> do+ liftIO $ action ma+ lift . logMessage Debug $ "executed action on message"
+ internal/Plumbing/Pipeline/Streaming/Source.hs view
@@ -0,0 +1,26 @@+module Pipeline.Streaming.Source where+import Conduit+import Data.ByteString+import Lib.Logger.Types (LogLevel (..), MonadLogger (..))+import Network.ConnectionAPI (ReaderAPI (..))+import Prelude hiding (length, null)++source :: (MonadLogger m , MonadIO m)+ => ReaderAPI reader+ -> reader+ -> ConduitT () ByteString m ()+source readerApi reader = do+ lift . logMessage Debug $ "reading from connection"+ result <- liftIO $ readData readerApi reader 4096+ case result of+ Left err -> do+ lift . logMessage Info $ ("read failed: " ++ err)+ Right chunk -> do+ case null chunk of+ True -> do+ lift . logMessage Info $ "read returned empty chunk; treating as disconnect"+ return ()+ _ -> do+ lift . logMessage Debug $ ("read " ++ show (length chunk) ++ " bytes")+ yield chunk+ source readerApi reader
+ internal/Plumbing/Queue/API.hs view
@@ -0,0 +1,20 @@+{-# LANGUAGE GADTs #-}++module Queue.API+ ( Queue (..)+ , QueueItem (..)+ ) where++import Transformers.Transformers (Transformer (..))++data QueueItem where QueueItem :: Transformer m => m -> QueueItem++instance Transformer QueueItem where+ transform (QueueItem item) = transform item++data Queue = Queue+ { enqueue :: QueueItem -> IO (Either String ())+ , dequeue :: IO (Either String QueueItem)+ , close :: IO ()+ , open :: IO ()+ }
+ internal/Plumbing/Queue/TransactionalQueue.hs view
@@ -0,0 +1,56 @@+module Queue.TransactionalQueue+ ( newQueue+ ) where++import Control.Concurrent.STM+import qualified Control.Monad+import Queue.API (Queue (Queue), QueueItem)++data TransactionalQueue = TransactionalQueue+ { transactionalQueueItems :: TBQueue QueueItem+ , transactionalQueueClosed :: TMVar ()+ }++newQueue :: IO Queue+newQueue = do+ queueItems <- newTBQueueIO 1000+ queueClosed <- newEmptyTMVarIO+ let transactionalQueue =+ TransactionalQueue+ { transactionalQueueItems = queueItems+ , transactionalQueueClosed = queueClosed+ }+ pure $+ Queue+ (enqueue transactionalQueue)+ (dequeue transactionalQueue)+ (close transactionalQueue)+ (open transactionalQueue)++enqueue :: TransactionalQueue -> QueueItem -> IO (Either String ())+enqueue transactionalQueue item = do+ isOpen <- atomically $ isEmptyTMVar (transactionalQueueClosed transactionalQueue)+ if isOpen+ then atomically $ do+ writeTBQueue (transactionalQueueItems transactionalQueue) item+ pure (Right ())+ else pure (Left "Queue is closed")++dequeue :: TransactionalQueue -> IO (Either String QueueItem)+dequeue transactionalQueue =+ atomically $+ (Right <$> readTBQueue (transactionalQueueItems transactionalQueue))+ `orElse`+ (do+ isOpen <- isEmptyTMVar (transactionalQueueClosed transactionalQueue)+ check (not isOpen)+ pure (Left "Queue is closed"))++close :: TransactionalQueue -> IO ()+close transactionalQueue = atomically $ do+ _ <- tryTakeTMVar (transactionalQueueClosed transactionalQueue)+ putTMVar (transactionalQueueClosed transactionalQueue) ()++open :: TransactionalQueue -> IO ()+open transactionalQueue =+ Control.Monad.void (atomically (tryTakeTMVar (transactionalQueueClosed transactionalQueue)))
+ internal/Plumbing/Transformers/Transformers.hs view
@@ -0,0 +1,8 @@+module Transformers.Transformers+ ( Transformer (..)+ ) where++import qualified Data.ByteString.Lazy as LBS++class Transformer a where+ transform :: a -> LBS.ByteString
+ internal/Plumbing/Validators/Validators.hs view
@@ -0,0 +1,7 @@+module Validators.Validators where++import Data.ByteString++class Validator a where+ validate :: a -> Either ByteString ()+
+ internal/Policy/Auth/Jwt.hs view
@@ -0,0 +1,10 @@+module Auth.Jwt+ ( JwtBundle (..)+ , auth+ , parseJwtBundle+ ) where++import Auth.Types++auth :: JWTTokenData -> Auth+auth = AuthJWT
+ internal/Policy/Auth/NKey.hs view
@@ -0,0 +1,9 @@+module Auth.NKey+ ( auth+ , signNonceWithSeed+ ) where++import Auth.Types++auth :: NKeyData -> Auth+auth = AuthNKey
+ internal/Policy/Auth/None.hs view
@@ -0,0 +1,8 @@+module Auth.None+ ( auth+ ) where++import Auth.Types++auth :: Auth+auth = AuthNone
+ internal/Policy/Auth/Token.hs view
@@ -0,0 +1,8 @@+module Auth.Token+ ( auth+ ) where++import Auth.Types++auth :: AuthTokenData -> Auth+auth = AuthToken
+ internal/Policy/Auth/Types.hs view
@@ -0,0 +1,395 @@+{-# LANGUAGE OverloadedStrings #-}++module Auth.Types+ ( User+ , Pass+ , UserPassData+ , NKeyData+ , AuthTokenData+ , JWTTokenData+ , JwtBundle (..)+ , AuthError (..)+ , AuthContext (..)+ , AuthPatch (..)+ , Auth (..)+ , 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++type User = BS.ByteString+type Pass = BS.ByteString+type UserPassData = (User, Pass)++type NKeyData = BS.ByteString++type AuthTokenData = BS.ByteString++type JWTTokenData = BS.ByteString++data JwtBundle = JwtBundle+ { jwtToken :: BS.ByteString+ , jwtSeed :: BS.ByteString+ }+ deriving (Eq, Show)++newtype AuthError = AuthError String+ deriving (Eq, Show)++newtype AuthContext = AuthContext { authNonce :: Maybe BS.ByteString }+ deriving (Eq, Show)++data AuthPatch = AuthPatch+ { patchAuthToken :: Maybe BS.ByteString+ , patchUser :: Maybe BS.ByteString+ , patchPass :: Maybe BS.ByteString+ , patchJwt :: Maybe BS.ByteString+ , patchNKey :: Maybe BS.ByteString+ , patchSig :: Maybe BS.ByteString+ }+ deriving (Eq, Show)++data Auth = AuthNone+ | AuthToken AuthTokenData+ | AuthUserPass UserPassData+ | AuthNKey NKeyData+ | AuthJWT JWTTokenData+ deriving (Eq, Show)++emptyAuthPatch :: AuthPatch+emptyAuthPatch =+ AuthPatch+ { patchAuthToken = Nothing+ , patchUser = Nothing+ , patchPass = Nothing+ , patchJwt = Nothing+ , 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
+ internal/Policy/Auth/UserPass.hs view
@@ -0,0 +1,8 @@+module Auth.UserPass+ ( auth+ ) where++import Auth.Types++auth :: UserPassData -> Auth+auth = AuthUserPass
+ internal/Policy/Engine.hs view
@@ -0,0 +1,263 @@+module Engine+ ( closeClient+ , resetClient+ , runEngine+ ) where++import Auth.Types+import Control.Concurrent (forkIO)+import Control.Concurrent.STM+import Control.Monad (forM_, unless, void, when)+import Data.Foldable (for_)+import Handshake.Nats (performHandshake)+import Lib.Logger (LogLevel (..), MonadLogger (..))+import Network.ConnectionAPI+ ( Conn+ , ConnectionAPI+ , close+ , closeReader+ , closeWriter+ , connectTcp+ , open+ , reader+ , writer+ )+import Parser.API (ParsedMessage, ParserAPI)+import Pipeline.Broadcasting.API (BroadcastingAPI (BroadcastingAPI))+import Pipeline.Streaming.API (StreamingAPI (StreamingAPI))+import Queue.API (QueueItem (QueueItem))+import Router.Nats (RouteDirective (..), routeMessage)+import State.Store+ ( ClientState+ , closeQueue+ , config+ , connection+ , enqueue+ , incrementAttemptIndex+ , markClosed+ , markConnected+ , openQueue+ , queue+ , readAttemptIndex+ , readStatus+ , runClient+ , setClosing+ , setEndpoint+ , waitForClosed+ )+import State.Types+ ( ClientConfig (..)+ , ClientExitReason (..)+ , ClientStatus (..)+ )+import Subscription.Store+ ( SubscriptionStore+ , active+ , awaitCallbackDrain+ , awaitNoTrackedExpiries+ )+import Subscription.Types (SubscriptionMeta (SubscriptionMeta))+import qualified Types.Sub as Sub++data ConnectionRunResult = ConnectionDisconnected+ | ConnectionExit ClientExitReason++runEngine+ :: ConnectionAPI+ -> StreamingAPI+ -> BroadcastingAPI+ -> ParserAPI ParsedMessage+ -> ClientState+ -> SubscriptionStore+ -> Auth+ -> IO ()+runEngine connectionApi streamingApi broadcastingApi parserApi state store auth =+ loop (connectionAttempts (config state)) Nothing+ where+ StreamingAPI runStreaming = streamingApi+ BroadcastingAPI runBroadcasting = broadcastingApi++ loop remaining lastErr+ | remaining <= 0 = do+ runClient state $+ logMessage Info "retries exhausted; exiting"+ finalize (ExitRetriesExhausted lastErr)+ | otherwise = do+ status <- readStatus state+ case status of+ Running -> do+ attemptErr <- runAttempt+ nextStatus <- readStatus state+ case nextStatus of+ Running -> do+ when (remaining > 1) $ do+ runClient state $+ logMessage Info "retrying client connection"+ incrementAttemptIndex state+ loop (remaining - 1) (Just attemptErr)+ Closing reason ->+ finalize reason+ Closed _ ->+ pure ()+ Closing reason ->+ finalize reason+ Closed _ ->+ pure ()++ runAttempt = do+ openQueue state+ open connectionApi (connection state)+ transportResult <- acquireTransport+ 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+ 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)+ Right () -> do+ resubscribeIfNeeded+ connectionResult <- runConnection conn+ close connectionApi conn+ case connectionResult of+ ConnectionDisconnected -> do+ runClient state $+ logMessage Info "connection disconnected"+ pure "connection disconnected"+ ConnectionExit reason -> do+ setClosing state reason+ pure (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))++ resubscribeIfNeeded = do+ alreadyConnected <- markConnected state+ when alreadyConnected $ do+ activeSubscriptions <- active store+ let resumable = filter isResumable activeSubscriptions+ unless (null resumable) $+ runClient state+ (logMessage Info ("resubscribing " ++ show (length resumable) ++ " subscriptions"))+ forM_ resumable $ \(sid, SubscriptionMeta subject queueGroup _) ->+ enqueue state $+ QueueItem+ Sub.Sub+ { Sub.subject = subject+ , Sub.queueGroup = queueGroup+ , Sub.sid = sid+ }++ finalize reason = do+ result <- markClosed state reason+ for_ result (exitAction (config state))++ runConnection :: Conn -> IO ConnectionRunResult+ runConnection conn = do+ exitVar <- newEmptyTMVarIO+ readerDone <- newEmptyTMVarIO+ writerDone <- newEmptyTMVarIO++ let stopReader = closeReader (reader connectionApi) conn+ stopWriter = closeWriter (writer connectionApi) conn+ stopQueue = closeQueue state+ signalReaderDone = atomically (void (tryPutTMVar readerDone ()))+ signalWriterDone = atomically (void (tryPutTMVar writerDone ()))+ signalExit reason = atomically (void (tryPutTMVar exitVar reason))++ void . forkIO $ do+ runClient state $ do+ logMessage Debug "starting broadcasting thread"+ runBroadcasting+ (bufferLimit (config state))+ (queue state)+ (writer connectionApi)+ conn+ logMessage Debug "broadcasting thread exited"+ stopQueue+ stopReader+ signalWriterDone++ void . forkIO $ do+ runClient state $ do+ logMessage Debug "starting streaming thread"+ runStreaming+ (bufferLimit (config state))+ (reader connectionApi)+ conn+ parserApi+ (\message -> do+ directive <- routeMessage state store message+ case directive of+ RouteContinue ->+ pure ()+ RouteExit reason -> do+ setClosing state reason+ signalExit reason+ stopQueue+ stopReader+ stopWriter)+ logMessage Debug "streaming thread exited"+ stopQueue+ stopWriter+ signalReaderDone++ outcome <- atomically $+ (Left <$> readTMVar exitVar)+ `orElse` (Right () <$ readTMVar readerDone)+ `orElse` (Right () <$ readTMVar writerDone)+ atomically $ do+ readTMVar readerDone+ readTMVar writerDone+ case outcome of+ Left reason -> pure (ConnectionExit reason)+ Right () -> pure ConnectionDisconnected++ isResumable :: (a, SubscriptionMeta) -> Bool+ isResumable (_, SubscriptionMeta _ _ isReply) = not isReply++closeClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO ()+closeClient connectionApi state store =+ shutdownClient connectionApi state store ExitClosedByUser "closing client connection"++resetClient :: ConnectionAPI -> ClientState -> SubscriptionStore -> IO ()+resetClient connectionApi state store =+ shutdownClient connectionApi state store ExitResetRequested "resetting client connection"++shutdownClient+ :: ConnectionAPI+ -> ClientState+ -> SubscriptionStore+ -> ClientExitReason+ -> String+ -> IO ()+shutdownClient connectionApi state store reason message = do+ setClosing state reason+ runClient state $+ logMessage Info message+ closeQueue state+ closeReader (reader connectionApi) (connection state)+ closeWriter (writer connectionApi) (connection state)+ atomically $ waitForClosed state+ atomically $ awaitNoTrackedExpiries store+ atomically $ awaitCallbackDrain store
+ internal/Policy/Handshake/Nats.hs view
@@ -0,0 +1,116 @@+module Handshake.Nats+ ( HandshakeError (..)+ , performHandshake+ ) where++import Auth.Types+import Control.Monad (unless)+import qualified Data.ByteString as BS+import Data.Maybe (fromMaybe, isJust)+import Network.ConnectionAPI+ ( Conn+ , ConnectionAPI+ , TransportOption (..)+ , configure+ , readData+ , reader+ , writeDataLazy+ , writer+ )+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore, Reject)+ , ParsedMessage (..)+ , ParserAPI+ , parse+ )+import State.Store+ ( ClientState+ , config+ , setServerInfo+ , updateLogContextFromInfo+ )+import State.Types (ClientConfig (..))+import Transformers.Transformers (Transformer (transform))+import qualified Types.Connect as Connect+import qualified Types.Info as I++data HandshakeError = HandshakeTransportError String+ | HandshakeProtocolError String+ | HandshakeAuthError AuthError+ 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+ 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+ }+ case validateAuth auth of+ Left err ->+ pure (Left (HandshakeAuthError err))+ Right () ->+ case buildAuthPatch auth authContext of+ Left err ->+ pure (Left (HandshakeAuthError err))+ Right patch -> do+ writeResult <-+ writeDataLazy+ (writer connectionApi)+ conn+ (transform (applyAuthPatch patch connectPayload))+ case writeResult of+ Left err ->+ pure (Left (HandshakeTransportError err))+ Right () ->+ pure (Right ())+ where+ readInitialInfo = go mempty+ where+ go acc = do+ result <- readData (reader connectionApi) conn 4096+ case result of+ Left err ->+ pure (Left err)+ Right chunk ->+ if BS.null chunk+ then pure (Left "read returned empty chunk before INFO")+ else do+ let bytes = acc <> chunk+ case parse parserApi bytes of+ NeedMore ->+ go bytes+ DropPrefix n _ ->+ go (BS.drop n bytes)+ Reject reason ->+ pure (Left 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 ())+ go rest
+ internal/Policy/Publish.hs view
@@ -0,0 +1,8 @@+module Publish+ ( defaultPublishConfig+ ) where++import Publish.Config (PublishConfig)++defaultPublishConfig :: PublishConfig+defaultPublishConfig = (Nothing, Nothing, Nothing)
+ internal/Policy/Publish/Config.hs view
@@ -0,0 +1,10 @@+module Publish.Config+ ( Payload+ , Headers+ , PublishConfig+ ) where++import qualified Types.Msg as M+import Types.Msg (Headers, Payload)++type PublishConfig = (Maybe Payload, Maybe (Maybe M.Msg -> IO ()), Maybe Headers)
+ internal/Policy/Router/Nats.hs view
@@ -0,0 +1,66 @@+module Router.Nats+ ( RouteDirective (..)+ , routeMessage+ ) where++import Control.Monad.IO.Class (liftIO)+import Lib.Logger (LogLevel (..), MonadLogger (..))+import Parser.API+ ( ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong)+ )+import Queue.API (QueueItem (QueueItem))+import State.Store+ ( ClientState+ , enqueue+ , runClient+ , runNextPingAction+ , setServerInfo+ , updateLogContextFromInfo+ )+import State.Types (ClientExitReason (ExitServerError))+import Subscription.Store (SubscriptionStore, dispatchMessage)+import qualified Types.Err as Err+import qualified Types.Msg as Msg+import Types.Pong (Pong (..))++data RouteDirective = RouteContinue+ | RouteExit ClientExitReason+ deriving (Eq, Show)++routeMessage :: ClientState -> SubscriptionStore -> ParsedMessage -> IO RouteDirective+routeMessage state store parsed =+ runClient state $+ case parsed of+ ParsedMsg msg -> do+ logMessage Debug ("routing MSG: " ++ show msg)+ handled <- liftIO $ dispatchMessage store msg+ if handled+ then pure RouteContinue+ else do+ logMessage Error ("callback missing for SID: " ++ show (Msg.sid msg))+ pure RouteContinue+ ParsedInfo info -> do+ logMessage Debug ("routing INFO: " ++ show info)+ liftIO $ setServerInfo state info+ liftIO $ updateLogContextFromInfo state info+ pure RouteContinue+ ParsedPing _ -> do+ logMessage Debug "routing PING"+ liftIO $ enqueue state (QueueItem Pong)+ pure RouteContinue+ ParsedPong _ -> do+ logMessage Debug "routing PONG"+ liftIO $ runNextPingAction state+ pure RouteContinue+ ParsedOk okMsg -> do+ logMessage Debug ("routing OK: " ++ show okMsg)+ pure RouteContinue+ ParsedErr err -> do+ logMessage Debug ("routing ERR: " ++ show err)+ if Err.isFatal err+ then do+ logMessage Error ("fatal server error: " ++ show err)+ pure (RouteExit (ExitServerError err))+ else do+ logMessage Warn ("server error: " ++ show err)+ pure RouteContinue
+ internal/Policy/State/Store.hs view
@@ -0,0 +1,224 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}++module State.Store+ ( ClientState+ , newClientState+ , runClient+ , config+ , queue+ , connection+ , enqueue+ , openQueue+ , closeQueue+ , pushPingAction+ , runNextPingAction+ , nextSid+ , nextInbox+ , readServerInfo+ , setServerInfo+ , updateLogContextFromInfo+ , setConnectName+ , setEndpoint+ , readStatus+ , setClosing+ , markClosed+ , waitForClosed+ , waitForNotRunning+ , waitForServerInfo+ , markConnected+ , readAttemptIndex+ , incrementAttemptIndex+ ) where++import Control.Concurrent.STM+import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as BC+import Data.Maybe (fromMaybe)+import Lib.Logger+ ( LogLevel (..)+ , MonadLogger (..)+ , loggerApi+ )+import Lib.Logger.Types (AppM, LogContext (..))+import Lib.LoggerAPI (runWithLogger, updateLogContext)+import Network.ConnectionAPI (Conn)+import Nuid (Nuid, newNuidIO, nextNuid)+import Queue.API (Queue (Queue), QueueItem, close, open)+import Sid (SIDCounter, initialSIDCounter, nextSID)+import State.Types+import Types.Info (Info, client_id)+import Types.Msg (SID, Subject)++data ClientState = ClientState+ { clientConfig :: ClientConfig+ , clientQueue :: Queue+ , clientPings :: TQueue (IO ())+ , clientConnectedOnce :: TVar Bool+ , 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+newClientState cfg queue conn logContext = do+ pings <- newTQueueIO+ connectedOnce <- newTVarIO False+ sidCounter <- newTVarIO initialSIDCounter+ inboxNuid <- newTVarIO =<< newNuidIO+ serverInfo <- newTVarIO Nothing+ attemptIndex <- newTVarIO 0+ status <- newTVarIO Running+ pure ClientState+ { clientConfig = cfg+ , clientQueue = queue+ , clientPings = pings+ , clientConnectedOnce = connectedOnce+ , clientSidCounter = sidCounter+ , clientInboxNuid = inboxNuid+ , clientServerInfo = serverInfo+ , clientAttemptIndex = attemptIndex+ , clientStatus = status+ , clientConnection = conn+ , clientLogContext = logContext+ }++runClient :: ClientState -> AppM a -> IO a+runClient client =+ runWithLogger loggerApi (loggerConfig (clientConfig client)) (clientLogContext client)++config :: ClientState -> ClientConfig+config = clientConfig++queue :: ClientState -> Queue+queue = clientQueue++connection :: ClientState -> Conn+connection = clientConnection++enqueue :: ClientState -> QueueItem -> IO ()+enqueue client item = do+ let Queue queueEnqueue _ _ _ = clientQueue client+ result <- queueEnqueue item+ case result of+ Left err ->+ runClient client $+ logMessage Error ("enqueueing item failed: " ++ err)+ Right () ->+ pure ()++openQueue :: ClientState -> IO ()+openQueue = open . clientQueue++closeQueue :: ClientState -> IO ()+closeQueue = close . clientQueue++pushPingAction :: ClientState -> IO () -> IO ()+pushPingAction client action =+ atomically $ writeTQueue (clientPings client) action++runNextPingAction :: ClientState -> IO ()+runNextPingAction client = do+ action <- atomically $ tryReadTQueue (clientPings client)+ fromMaybe (pure ()) action++nextSid :: ClientState -> IO SID+nextSid client =+ atomically $ do+ counter <- readTVar (clientSidCounter client)+ let (sid, counter') = nextSID counter+ writeTVar (clientSidCounter client) counter'+ pure sid++nextInbox :: ClientState -> IO Subject+nextInbox client =+ atomically $ do+ nuid <- readTVar (clientInboxNuid client)+ let (token, nextNuidValue) = nextNuid nuid+ writeTVar (clientInboxNuid client) nextNuidValue+ pure ("_INBOX." <> token)++readServerInfo :: ClientState -> IO (Maybe Info)+readServerInfo = readTVarIO . clientServerInfo++setServerInfo :: ClientState -> Info -> IO ()+setServerInfo client info =+ atomically $ writeTVar (clientServerInfo client) (Just info)++updateLogContextFromInfo :: ClientState -> Info -> IO ()+updateLogContextFromInfo client info =+ updateLogContext loggerApi (clientLogContext client) $ \ctx ->+ ctx { lcClientId = client_id info }++setConnectName :: ClientState -> Maybe BS.ByteString -> IO ()+setConnectName client maybeName =+ updateLogContext loggerApi (clientLogContext client) $ \ctx ->+ ctx { lcConnectName = BC.unpack <$> maybeName }++setEndpoint :: ClientState -> (String, Int) -> IO ()+setEndpoint client (host, port) =+ updateLogContext loggerApi (clientLogContext client) $ \ctx ->+ ctx { lcServer = Just (host ++ ":" ++ show port) }++readStatus :: ClientState -> IO ClientStatus+readStatus = readTVarIO . clientStatus++setClosing :: ClientState -> ClientExitReason -> IO ()+setClosing client reason =+ atomically . modifyTVar' (clientStatus client) $ \case+ Closed result -> Closed result+ Closing result -> Closing result+ Running -> Closing reason++markClosed :: ClientState -> ClientExitReason -> IO (Maybe ClientExitReason)+markClosed client fallbackReason =+ atomically $ do+ status <- readTVar (clientStatus client)+ case status of+ Closed _ ->+ pure Nothing+ Closing reason -> do+ writeTVar (clientStatus client) (Closed reason)+ pure (Just reason)+ Running -> do+ writeTVar (clientStatus client) (Closed fallbackReason)+ pure (Just fallbackReason)++waitForClosed :: ClientState -> STM ()+waitForClosed client = do+ status <- readTVar (clientStatus client)+ case status of+ Closed _ -> pure ()+ _ -> retry++waitForNotRunning :: ClientState -> STM ()+waitForNotRunning client = do+ status <- readTVar (clientStatus client)+ case status of+ Running -> retry+ _ -> pure ()++waitForServerInfo :: ClientState -> STM ()+waitForServerInfo client = do+ info <- readTVar (clientServerInfo client)+ case info of+ Just _ -> pure ()+ Nothing -> retry++markConnected :: ClientState -> IO Bool+markConnected client =+ atomically $ do+ alreadyConnected <- readTVar (clientConnectedOnce client)+ writeTVar (clientConnectedOnce client) True+ pure alreadyConnected++readAttemptIndex :: ClientState -> IO Int+readAttemptIndex = readTVarIO . clientAttemptIndex++incrementAttemptIndex :: ClientState -> IO ()+incrementAttemptIndex client =+ atomically (modifyTVar' (clientAttemptIndex client) (+ 1))
+ internal/Policy/State/Types.hs view
@@ -0,0 +1,41 @@+{-# LANGUAGE RankNTypes #-}++module State.Types+ ( TLSPublicKey+ , TLSPrivateKey+ , TLSCertData+ , ClientConfig (..)+ , ClientExitReason (..)+ , ClientStatus (..)+ ) where++import qualified Data.ByteString as BS+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)++data ClientConfig = ClientConfig+ { connectionAttempts :: Int+ , callbackConcurrency :: Int+ , bufferLimit :: Int+ , connectConfig :: Connect+ , loggerConfig :: LoggerConfig+ , tlsCert :: Maybe TLSCertData+ , exitAction :: ClientExitReason -> IO ()+ , connectOptions :: [(String, Int)]+ }++data ClientExitReason = ExitClosedByUser+ | ExitRetriesExhausted (Maybe String)+ | ExitServerError Err.Err+ | ExitResetRequested+ deriving (Eq, Show)++data ClientStatus = Running+ | Closing ClientExitReason+ | Closed ClientExitReason+ deriving (Eq, Show)
+ internal/Policy/Subscription/Store.hs view
@@ -0,0 +1,183 @@+module Subscription.Store+ ( SubscriptionStore+ , newSubscriptionStore+ , register+ , unregister+ , dispatchMessage+ , active+ , startWorkers+ , awaitCallbackDrain+ , startExpiryWorker+ , awaitNoTrackedExpiries+ , hasTrackedExpiries+ ) where++import Control.Concurrent (forkIO, threadDelay)+import Control.Concurrent.STM+import Control.Exception (SomeException, finally)+import Control.Monad (unless, void)+import qualified Data.Heap as Heap+import qualified Data.Map as Map+import Data.Time.Clock+import Lib.WorkerPool (startWorkerPool)+import Subscription.Types+import qualified Types.Msg as M+import Types.Msg (SID)++data SubscriptionState = SubscriptionState+ { subscriptionCallbacks :: Map.Map SID (Maybe M.Msg -> IO ())+ , subscriptionExpiryHeap :: Heap.MinHeap (UTCTime, SID)+ , subscriptionTrackedExpiries :: Map.Map SID UTCTime+ , subscriptionMeta :: Map.Map SID SubscriptionMeta+ }++data SubscriptionStore = SubscriptionStore+ { storeState :: TVar SubscriptionState+ , storeCallbackQueue :: TQueue (IO ())+ , storeCallbackPending :: TVar Int+ }++newSubscriptionStore :: IO SubscriptionStore+newSubscriptionStore =+ SubscriptionStore+ <$> newTVarIO emptySubscriptionState+ <*> newTQueueIO+ <*> newTVarIO 0++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+ 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)+ , subscriptionMeta = Map.insert sid meta (subscriptionMeta state)+ }+ in case expiryAt of+ Just expiresAt | isReply meta ->+ trackSubscriptionExpiry sid expiresAt stateWithCallback+ _ ->+ stateWithCallback++unregister :: SubscriptionStore -> SID -> IO ()+unregister store sid =+ atomically $+ modifyTVar' (storeState store) (removeSubscriptionLocal sid)++dispatchMessage :: SubscriptionStore -> M.Msg -> IO Bool+dispatchMessage store msg =+ atomically $ do+ state <- readTVar (storeState store)+ let sid = M.sid msg+ case Map.lookup sid (subscriptionCallbacks state) of+ Nothing ->+ pure False+ Just callback -> do+ let shouldRemove =+ maybe False isReply (Map.lookup sid (subscriptionMeta state))+ nextState =+ if shouldRemove+ then removeSubscriptionLocal sid state+ else state+ writeTVar (storeState store) nextState+ enqueueCallbackSTM store (callback (Just msg))+ pure True++active :: SubscriptionStore -> IO [(SID, SubscriptionMeta)]+active store =+ Map.toList . subscriptionMeta <$> readTVarIO (storeState store)++startWorkers :: Int -> SubscriptionStore -> STM () -> (SomeException -> IO ()) -> IO ()+startWorkers concurrency store =+ startWorkerPool (max 1 concurrency) (storeCallbackQueue store)++awaitCallbackDrain :: SubscriptionStore -> STM ()+awaitCallbackDrain store = do+ pending <- readTVar (storeCallbackPending store)+ check (pending == 0)++startExpiryWorker :: SubscriptionStore -> IO Bool -> IO ()+startExpiryWorker store shouldStop =+ void . forkIO $ loop+ where+ loop = do+ stop <- shouldStop+ unless stop $ do+ now <- getCurrentTime+ expiredAny <- expireReadySubscriptions store now+ unless expiredAny (threadDelay 1000000)+ loop++awaitNoTrackedExpiries :: SubscriptionStore -> STM ()+awaitNoTrackedExpiries store = do+ tracked <- subscriptionTrackedExpiries <$> readTVar (storeState store)+ unless (Map.null tracked) retry++hasTrackedExpiries :: SubscriptionStore -> IO Bool+hasTrackedExpiries store =+ not . Map.null . subscriptionTrackedExpiries <$> readTVarIO (storeState store)++enqueueCallbackSTM :: SubscriptionStore -> IO () -> STM ()+enqueueCallbackSTM store action = do+ modifyTVar' (storeCallbackPending store) (+1)+ let wrapped =+ action `finally` atomically (modifyTVar' (storeCallbackPending store) (subtract 1))+ writeTQueue (storeCallbackQueue store) wrapped++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+ pure (not (null actions))++trackSubscriptionExpiry :: SID -> UTCTime -> SubscriptionState -> SubscriptionState+trackSubscriptionExpiry sid expiry state =+ state+ { subscriptionExpiryHeap = Heap.insert (expiry, sid) (subscriptionExpiryHeap state)+ , subscriptionTrackedExpiries = Map.insert sid expiry (subscriptionTrackedExpiries state)+ }++removeSubscriptionLocal :: SID -> SubscriptionState -> SubscriptionState+removeSubscriptionLocal sid state =+ state+ { subscriptionCallbacks = Map.delete sid (subscriptionCallbacks state)+ , subscriptionTrackedExpiries = Map.delete sid (subscriptionTrackedExpiries state)+ , subscriptionMeta = Map.delete sid (subscriptionMeta state)+ }++collectExpiredCallbacks :: UTCTime -> SubscriptionState -> ([IO ()], SubscriptionState)+collectExpiredCallbacks now = go []+ where+ go actions state =+ case Heap.viewHead (subscriptionExpiryHeap state) of+ Nothing ->+ (reverse actions, state)+ Just (expiry, sid) ->+ case Map.lookup sid (subscriptionTrackedExpiries state) of+ Nothing ->+ go actions (dropHeapHead state)+ Just trackedExpiry+ | trackedExpiry /= expiry ->+ go actions (dropHeapHead state)+ | now < trackedExpiry ->+ (reverse actions, state)+ | otherwise ->+ let callback = Map.lookup sid (subscriptionCallbacks state)+ state' = removeSubscriptionLocal sid (dropHeapHead state)+ actions' = maybe actions ((: actions) . ($ Nothing)) callback+ in go actions' state'++dropHeapHead :: SubscriptionState -> SubscriptionState+dropHeapHead state =+ case Heap.view (subscriptionExpiryHeap state) of+ Nothing -> state+ Just (_, heapTail) -> state { subscriptionExpiryHeap = heapTail }
+ internal/Policy/Subscription/Types.hs view
@@ -0,0 +1,16 @@+module Subscription.Types+ ( SubscribeConfig (..)+ , SubscriptionMeta (..)+ ) where++import Data.Time.Clock (NominalDiffTime)+import Types.Msg (Subject)++newtype SubscribeConfig = SubscribeConfig { expiry :: Maybe NominalDiffTime }++data SubscriptionMeta = SubscriptionMeta+ { subject :: Subject+ , queueGroup :: Maybe Subject+ , isReply :: Bool+ }+ deriving (Eq, Show)
+ internal/Types/Connect.hs view
@@ -0,0 +1,171 @@+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE OverloadedStrings #-}++module Types.Connect+ ( Connect (..)+ , validateAuthToken+ , validateUser+ , validatePass+ , validateName+ , validateLang+ , validateVersion+ , validateProtocol+ , validateSig+ , validateJwt+ , validateNKey+ ) where++import Data.Aeson+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import qualified Data.Text.Encoding as E+import GHC.Generics (Generic)+import Transformers.Transformers (Transformer (..))+import Validators.Validators++data Connect = Connect+ { verbose :: Bool+ , pedantic :: Bool+ , tls_required :: Bool+ , auth_token :: Maybe BS.ByteString+ , user :: Maybe BS.ByteString+ , pass :: Maybe BS.ByteString+ , name :: Maybe BS.ByteString+ , lang :: BS.ByteString+ , version :: BS.ByteString+ , protocol :: Maybe Int+ , echo :: Maybe Bool+ , sig :: Maybe BS.ByteString+ , jwt :: Maybe BS.ByteString+ , nkey :: Maybe BS.ByteString+ , no_responders :: Maybe Bool+ , headers :: Maybe Bool+ }+ deriving (Eq, Show)++instance ToJSON Connect where+ toJSON = toJSON . toConnectJSON++instance Transformer Connect where+ transform c =+ let encoded = encode c+ in LBS.fromChunks ["CONNECT ", LBS.toStrict encoded, "\r\n"]++newtype Utf8ByteString = Utf8ByteString { unUtf8ByteString :: BS.ByteString }+ deriving (Eq, Show)++instance ToJSON Utf8ByteString where+ toJSON = toJSON . E.decodeUtf8 . unUtf8ByteString++data ConnectJSON = ConnectJSON+ { connectJSON_verbose :: Bool+ , connectJSON_pedantic :: Bool+ , connectJSON_tls_required :: Bool+ , connectJSON_auth_token :: Maybe Utf8ByteString+ , connectJSON_user :: Maybe Utf8ByteString+ , connectJSON_pass :: Maybe Utf8ByteString+ , connectJSON_name :: Maybe Utf8ByteString+ , connectJSON_lang :: Utf8ByteString+ , connectJSON_version :: Utf8ByteString+ , connectJSON_protocol :: Maybe Int+ , connectJSON_echo :: Maybe Bool+ , connectJSON_sig :: Maybe Utf8ByteString+ , connectJSON_jwt :: Maybe Utf8ByteString+ , connectJSON_nkey :: Maybe Utf8ByteString+ , connectJSON_no_responders :: Maybe Bool+ , connectJSON_headers :: Maybe Bool+ }+ deriving (Eq, Generic, Show)++instance ToJSON ConnectJSON where+ toJSON = genericToJSON defaultOptions+ { fieldLabelModifier = drop connectJSONPrefixLength+ , omitNothingFields = True+ }++connectJSONPrefixLength :: Int+connectJSONPrefixLength = 12++toConnectJSON :: Connect -> ConnectJSON+toConnectJSON c =+ ConnectJSON+ { connectJSON_verbose = verbose c+ , connectJSON_pedantic = pedantic c+ , connectJSON_tls_required = tls_required c+ , connectJSON_auth_token = Utf8ByteString <$> auth_token c+ , connectJSON_user = Utf8ByteString <$> user c+ , connectJSON_pass = Utf8ByteString <$> pass c+ , connectJSON_name = Utf8ByteString <$> name c+ , connectJSON_lang = Utf8ByteString (lang c)+ , connectJSON_version = Utf8ByteString (version c)+ , connectJSON_protocol = protocol c+ , connectJSON_echo = echo c+ , connectJSON_sig = Utf8ByteString <$> sig c+ , connectJSON_jwt = Utf8ByteString <$> jwt c+ , connectJSON_nkey = Utf8ByteString <$> nkey c+ , connectJSON_no_responders = no_responders c+ , connectJSON_headers = headers c+ }++instance Validator Connect where+ validate c = do+ validateAuthToken c+ validateUser c+ validatePass c+ validateName c+ validateLang c+ validateVersion c+ validateProtocol c+ validateSig c+ validateJwt c+ validateNKey c++validateAuthToken :: Connect -> Either BS.ByteString ()+validateAuthToken c+ | auth_token c == Just "" = Left "explicit empty auth token"+ | otherwise = Right ()++validateUser :: Connect -> Either BS.ByteString ()+validateUser c+ | user c == Just "" = Left "explicit empty user"+ | otherwise = Right ()++validatePass :: Connect -> Either BS.ByteString ()+validatePass c+ | pass c == Just "" = Left "explicit empty pass"+ | otherwise = Right ()++validateName :: Connect -> Either BS.ByteString ()+validateName c+ | name c == Just "" = Left "explicit empty name"+ | otherwise = Right ()++validateLang :: Connect -> Either BS.ByteString ()+validateLang c+ | lang c == "" = Left "explicit empty lang"+ | otherwise = Right ()++validateVersion :: Connect -> Either BS.ByteString ()+validateVersion c+ | version c == "" = Left "explicit empty version"+ | otherwise = Right ()++validateProtocol :: Connect -> Either BS.ByteString ()+validateProtocol c+ | protocol c `notElem` [Nothing, Just 0, Just 1] = Left "invalid protocol"+ | otherwise = Right ()++validateSig :: Connect -> Either BS.ByteString ()+validateSig c+ | sig c == Just "" = Left "explicit empty sig"+ | otherwise = Right ()++validateJwt :: Connect -> Either BS.ByteString ()+validateJwt c+ | jwt c == Just "" = Left "explicit empty jwt"+ | otherwise = Right ()++validateNKey :: Connect -> Either BS.ByteString ()+validateNKey c+ | nkey c == Just "" = Left "explicit empty nkey"+ | otherwise = Right ()
+ internal/Types/Err.hs view
@@ -0,0 +1,30 @@+module Types.Err+ ( Reason+ , Err (..)+ , isFatal+ ) where++import Data.ByteString (ByteString)++type Reason = ByteString++data Err = ErrUnknownOp Reason+ | ErrRoutePortConn Reason+ | ErrAuthViolation Reason+ | ErrAuthTimeout Reason+ | ErrInvalidProtocol Reason+ | ErrMaxControlLineEx Reason+ | ErrErr Reason+ | ErrTlsRequired Reason+ | ErrStaleConn Reason+ | ErrMaxConnsEx Reason+ | ErrSlowConsumer Reason+ | ErrMaxPayload Reason+ | ErrInvalidSubject Reason+ | ErrPermViolation Reason+ deriving (Eq, Show)++isFatal :: Err -> Bool+isFatal (ErrInvalidSubject _) = False+isFatal (ErrPermViolation _) = False+isFatal _ = True
+ internal/Types/Info.hs view
@@ -0,0 +1,79 @@+{-# LANGUAGE DeriveGeneric #-}++module Types.Info+ ( Info (..)+ ) where++import Data.Aeson+import qualified Data.ByteString as BS+import qualified Data.Text.Encoding as E+import GHC.Generics (Generic)++data Info = Info+ { server_id :: BS.ByteString+ , version :: BS.ByteString+ , go :: BS.ByteString+ , host :: BS.ByteString+ , port :: Int+ , max_payload :: Int+ , proto :: Int+ , client_id :: Maybe Int+ , nonce :: Maybe BS.ByteString+ , auth_required :: Maybe Bool+ , tls_required :: Maybe Bool+ , connect_urls :: Maybe [BS.ByteString]+ , ldm :: Maybe Bool+ , headers :: Maybe Bool+ }+ deriving (Eq, Show)++instance FromJSON Info where+ parseJSON = fmap fromInfoJSON . genericParseJSON defaultOptions+ { fieldLabelModifier = drop infoJSONPrefixLength+ }++infoJSONPrefixLength :: Int+infoJSONPrefixLength = 9++newtype Utf8ByteString = Utf8ByteString { unUtf8ByteString :: BS.ByteString }+ deriving (Eq, Show)++instance FromJSON Utf8ByteString where+ parseJSON = fmap (Utf8ByteString . E.encodeUtf8) . parseJSON++data InfoJSON = InfoJSON+ { infoJSON_server_id :: Utf8ByteString+ , infoJSON_version :: Utf8ByteString+ , infoJSON_go :: Utf8ByteString+ , infoJSON_host :: Utf8ByteString+ , infoJSON_port :: Int+ , infoJSON_max_payload :: Int+ , infoJSON_proto :: Int+ , infoJSON_client_id :: Maybe Int+ , infoJSON_nonce :: Maybe Utf8ByteString+ , infoJSON_auth_required :: Maybe Bool+ , infoJSON_tls_required :: Maybe Bool+ , infoJSON_connect_urls :: Maybe [Utf8ByteString]+ , infoJSON_ldm :: Maybe Bool+ , infoJSON_headers :: Maybe Bool+ }+ deriving (Eq, Generic, Show)++fromInfoJSON :: InfoJSON -> Info+fromInfoJSON infoJson =+ Info+ { server_id = unUtf8ByteString (infoJSON_server_id infoJson)+ , version = unUtf8ByteString (infoJSON_version infoJson)+ , go = unUtf8ByteString (infoJSON_go infoJson)+ , host = unUtf8ByteString (infoJSON_host infoJson)+ , port = infoJSON_port infoJson+ , max_payload = infoJSON_max_payload infoJson+ , proto = infoJSON_proto infoJson+ , client_id = infoJSON_client_id infoJson+ , nonce = unUtf8ByteString <$> infoJSON_nonce infoJson+ , auth_required = infoJSON_auth_required infoJson+ , tls_required = infoJSON_tls_required infoJson+ , connect_urls = fmap (map unUtf8ByteString) (infoJSON_connect_urls infoJson)+ , ldm = infoJSON_ldm infoJson+ , headers = infoJSON_headers infoJson+ }
+ internal/Types/Msg.hs view
@@ -0,0 +1,23 @@+module Types.Msg+ ( Subject+ , SID+ , Payload+ , Headers+ , Msg (..)+ ) where++import Data.ByteString (ByteString)++type Subject = ByteString+type SID = ByteString+type Payload = ByteString+type Headers = [(ByteString, ByteString)]++data Msg = Msg+ { subject :: Subject+ , sid :: SID+ , replyTo :: Maybe Subject+ , payload :: Maybe Payload+ , headers :: Maybe Headers+ }+ deriving (Eq, Show)
+ internal/Types/Nuid.hs view
@@ -0,0 +1,95 @@+module Nuid+ ( Nuid (..)+ , newNuid+ , newNuidIO+ , nextNuid+ ) where++import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as BC+import Data.Word (Word64, Word8)+import System.Random (StdGen, newStdGen, randomR)++data Nuid = Nuid+ { nuidPrefix :: BS.ByteString+ , nuidSeq :: Word64+ , nuidInc :: Word64+ , nuidRng :: StdGen+ }++newNuid :: StdGen -> Nuid+newNuid rng0 = Nuid prefix seqVal incVal rng3+ where+ (prefix, rng1) = randomPrefix rng0+ (seqVal, rng2) = randomSeq rng1+ (incVal, rng3) = randomInc rng2++newNuidIO :: IO Nuid+newNuidIO = newNuid <$> newStdGen++nextNuid :: Nuid -> (BS.ByteString, Nuid)+nextNuid nuid =+ let seqVal = nuidSeq nuid + nuidInc nuid+ in if seqVal >= maxSeq+ then resetAndNext nuid+ else let nuid' = nuid { nuidSeq = seqVal }+ in (renderNuid (nuidPrefix nuid) seqVal, nuid')++resetAndNext :: Nuid -> (BS.ByteString, Nuid)+resetAndNext nuid = (renderNuid prefix seqVal, Nuid prefix seqVal incVal rng3)+ where+ (prefix, rng1) = randomPrefix (nuidRng nuid)+ (seqVal, rng2) = randomSeq rng1+ (incVal, rng3) = randomInc rng2++renderNuid :: BS.ByteString -> Word64 -> BS.ByteString+renderNuid prefix seqVal = BS.append prefix (encodeSeq seqVal)++randomPrefix :: StdGen -> (BS.ByteString, StdGen)+randomPrefix rng0 = go prefixLen rng0 []+ where+ go 0 rng acc = (BS.pack (reverse acc), rng)+ go n rng acc =+ let (idx, rng') = randomR (0, base62Len - 1) rng+ in go (n - 1) rng' (base62Index idx : acc)++encodeSeq :: Word64 -> BS.ByteString+encodeSeq value = BS.pack (map base62Index (buildDigits seqLen value []))+ where+ buildDigits 0 _ acc = acc+ buildDigits n v acc =+ let (q, r) = v `quotRem` base+ in buildDigits (n - 1) q (fromIntegral r : acc)++base62Index :: Int -> Word8+base62Index = BS.index base62Alphabet++base62Alphabet :: BS.ByteString+base62Alphabet = BC.pack "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"++base62Len :: Int+base62Len = BS.length base62Alphabet++prefixLen :: Int+prefixLen = 12++seqLen :: Int+seqLen = 10++base :: Word64+base = 62++maxSeq :: Word64+maxSeq = base ^ seqLen++minInc :: Word64+minInc = 33++maxInc :: Word64+maxInc = 333++randomSeq :: StdGen -> (Word64, StdGen)+randomSeq = randomR (0, maxSeq - 1)++randomInc :: StdGen -> (Word64, StdGen)+randomInc = randomR (minInc, maxInc - 1)
+ internal/Types/Ok.hs view
@@ -0,0 +1,4 @@+module Types.Ok where++data Ok = Ok+ deriving (Eq, Show)
+ internal/Types/Ping.hs view
@@ -0,0 +1,15 @@+module Types.Ping where++import qualified Data.ByteString.Char8 as BC+import qualified Data.ByteString.Lazy as LBS+import Transformers.Transformers (Transformer (..))+import Validators.Validators++data Ping = Ping+ deriving (Eq, Show)++instance Transformer Ping where+ transform _ = LBS.fromStrict (BC.pack "PING\r\n")++instance Validator Ping where+ validate _ = Right ()
+ internal/Types/Pong.hs view
@@ -0,0 +1,15 @@+module Types.Pong where++import qualified Data.ByteString.Char8 as BC+import qualified Data.ByteString.Lazy as LBS+import Transformers.Transformers (Transformer (..))+import Validators.Validators++data Pong = Pong+ deriving (Eq, Show)++instance Transformer Pong where+ transform _ = LBS.fromStrict (BC.pack "PONG\r\n")++instance Validator Pong where+ validate _ = Right ()
+ internal/Types/Pub.hs view
@@ -0,0 +1,85 @@+{-# LANGUAGE OverloadedStrings #-}++module Types.Pub where++import Data.ByteString hiding (foldr, map)+import qualified Data.ByteString.Char8 as BC+import qualified Data.ByteString.Lazy as LBS+import Data.Maybe+import Prelude hiding (concat, length, null)+import Transformers.Transformers (Transformer (..))+import Types.Msg (Headers, Payload, Subject)+import Validators.Validators++data Pub = Pub+ { subject :: Subject+ , replyTo :: Maybe Subject+ , headers :: Maybe Headers+ , payload :: Maybe Payload+ }+ deriving (Eq, Show)++instance Transformer Pub where+ transform pubMsg =+ case headers pubMsg of+ Nothing ->+ let payload' = payloadChunk (payload pubMsg)+ payloadLen = length payload'+ control = concat+ (["PUB ", subject pubMsg, " "] ++ replyChunks (replyTo pubMsg)+ ++ [packInt payloadLen, "\r\n"])+ in LBS.fromChunks [control, payload', "\r\n"]+ Just headerList ->+ let headers' = headerString headerList+ headerLength = length headers'+ payload' = payloadChunk (payload pubMsg)+ totalLen = headerLength + length payload'+ control = concat+ (["HPUB ", subject pubMsg, " "] ++ replyChunks (replyTo pubMsg)+ ++ [packInt headerLength, " ", packInt totalLen, "\r\n"])+ in LBS.fromChunks [control, headers', payload', "\r\n"]++instance Validator Pub where+ validate p = do+ validateSubject p+ validateReplyTo p+ validatePayload p+ validateHeaders p++validateSubject :: Pub -> Either ByteString ()+validateSubject p+ | subject p == "" = Left "explicit empty subject"+ | otherwise = Right ()++validateReplyTo :: Pub -> Either ByteString ()+validateReplyTo p+ | replyTo p == Just "" = Left "explicit empty replyTo"+ | otherwise = Right ()++validatePayload :: Pub -> Either ByteString ()+validatePayload p+ | payload p == Just "" = Left "explicit empty payload"+ | otherwise = Right ()++validateHeaders :: Pub -> Either ByteString ()+validateHeaders p+ | isNothing (headers p) = Right ()+ | headers p == Just [] = Left "explicit empty headers"+ | Prelude.any (\(k, _) -> k == "") (fromJust (headers p)) = Left "explicit empty header key"+ | Prelude.any (\(_, v) -> v == "") (fromJust (headers p)) = Left "explicit empty header value"+ | otherwise = Right ()++headerString :: Headers -> ByteString+headerString hs =+ concat ("NATS/1.0\r\n" : foldr appendHeader ["\r\n"] hs)+ where+ appendHeader (key, value) acc = key : ":" : value : "\r\n" : acc++payloadChunk :: Maybe Payload -> Payload+payloadChunk = fromMaybe empty++replyChunks :: Maybe Subject -> [ByteString]+replyChunks = maybe [] (\reply -> [reply, " "])++packInt :: Int -> ByteString+packInt = BC.pack . show
+ internal/Types/Sid.hs view
@@ -0,0 +1,18 @@+module Sid+ ( SIDCounter+ , initialSIDCounter+ , nextSID+ ) where++import qualified Data.ByteString.Char8 as BC+import Data.Word (Word64)+import Types.Msg (SID)++type SIDCounter = Word64++initialSIDCounter :: SIDCounter+initialSIDCounter = 0++nextSID :: SIDCounter -> (SID, SIDCounter)+nextSID counter = (BC.pack (show next), next)+ where next = counter + 1
+ internal/Types/Sub.hs view
@@ -0,0 +1,43 @@+{-# LANGUAGE OverloadedStrings #-}++module Types.Sub where++import Data.ByteString (ByteString)+import qualified Data.ByteString.Lazy as LBS+import Transformers.Transformers (Transformer (..))+import Types.Msg (SID, Subject)+import Validators.Validators++data Sub = Sub+ { subject :: Subject+ , queueGroup :: Maybe Subject+ , sid :: SID+ }+ deriving (Eq, Show)++instance Transformer Sub where+ transform subMsg =+ case queueGroup subMsg of+ Just queue -> LBS.fromChunks ["SUB ", subject subMsg, " ", queue, " ", sid subMsg, "\r\n"]+ Nothing -> LBS.fromChunks ["SUB ", subject subMsg, " ", sid subMsg, "\r\n"]++instance Validator Sub where+ validate s = do+ validateSubject s+ validateQueueGroup s+ validateSid s++validateSubject :: Sub -> Either ByteString ()+validateSubject s+ | subject s == "" = Left "explicit empty subject"+ | otherwise = Right ()++validateQueueGroup :: Sub -> Either ByteString ()+validateQueueGroup s+ | queueGroup s == Just "" = Left "explicit empty queue group"+ | otherwise = Right ()++validateSid :: Sub -> Either ByteString ()+validateSid s+ | sid s == "" = Left "explicit empty sid"+ | otherwise = Right ()
+ internal/Types/Unsub.hs view
@@ -0,0 +1,31 @@+{-# LANGUAGE OverloadedStrings #-}++module Types.Unsub where++import Data.ByteString (ByteString)+import qualified Data.ByteString.Char8 as BC+import qualified Data.ByteString.Lazy as LBS+import Transformers.Transformers (Transformer (..))+import Types.Msg (SID)+import Validators.Validators++data Unsub = Unsub+ { sid :: SID+ , maxMsg :: Maybe Int+ }+ deriving (Eq, Show)++instance Transformer Unsub where+ transform unsubMsg =+ case maxMsg unsubMsg of+ Just count -> LBS.fromChunks ["UNSUB ", sid unsubMsg, " ", BC.pack (show count), "\r\n"]+ Nothing -> LBS.fromChunks ["UNSUB ", sid unsubMsg, "\r\n"]++instance Validator Unsub where+ validate u = do+ validateSid u++validateSid :: Unsub -> Either ByteString ()+validateSid u+ | sid u == "" = Left "explicit empty sid"+ | otherwise = Right ()
+ natskell.cabal view
@@ -0,0 +1,241 @@+cabal-version: 3.0+name: natskell+version: 0.0.0.1+synopsis: A NATS client library written in Haskell+tested-with:+ GHC ==8.8,+ GHC ==8.10,+ GHC ==9.0,+ GHC ==9.2,+ GHC ==9.4,+ GHC ==9.6,+ GHC ==9.8,+ GHC ==9.10,+ GHC ==9.12,+ GHC ==9.14++-- A longer description of the package.+description: Please see the README on GitHub at <https://github.com/samisagit/natskell#readme>++-- homepage:++-- A URL where users can report bugs.+bug-reports: https://github.com/samisagit/natskell/issues++license: MIT+license-file: LICENSE+author: samisagit+maintainer: sam@whiteteam.co.uk++-- A copyright notice.+-- copyright:++category: Web++source-repository head+ type: git+ location: https://github.com/samisagit/natskell.git++common shared+ default-language: Haskell2010+ build-depends:+ base >= 4.13 && < 5,+ ghc-options:+ -haddock -W -Wunused-packages+ pkgconfig-depends: zlib++library+ import: shared+ build-depends:+ natskell-internal,+ stm >= 2.5 && < 2.6,+ bytestring >= 0.10 && < 0.13,+ heap >= 0.6 && < 1.1,+ time >= 1.9 && < 1.16++ exposed-modules:+ API+ Client++ hs-source-dirs:+ client++library natskell-internal+ import: shared+ build-depends:+ conduit >= 1.3 && < 1.4,+ attoparsec >= 0.14 && < 0.15,+ bytestring >= 0.10 && < 0.13,+ aeson >= 2.1 && < 2.3,+ stm >= 2.5 && < 2.6,+ containers >= 0.5 && < 0.9,+ heap >= 0.6 && < 1.1,+ time >= 1.9 && < 1.16,+ network >= 3.1 && < 3.3,+ mtl >= 2.2 && < 2.4,+ cryptonite >= 0.29 && < 0.31,+ memory >= 0.14 && < 0.19,+ random >= 1.1 && < 1.3,+ network-simple >= 0.4 && < 0.5,+ tls >= 1.8 && < 3,+ text (>= 1.2 && < 1.3) || (>= 2.0 && < 2.2)++ exposed-modules:+ Auth.Types+ Auth.None+ Auth.Token+ Auth.UserPass+ Auth.NKey+ Auth.Jwt+ Handshake.Nats+ Router.Nats+ Subscription.Types+ Subscription.Store+ Engine+ Publish.Config+ Publish+ State.Types+ State.Store+ Parser.API+ Lib.Logger+ Lib.Logger.Types+ Lib.LoggerAPI+ Lib.CallOption+ Lib.WorkerPool+ Parser.Attoparsec+ Transformers.Transformers+ Validators.Validators+ Types.Msg+ Types.Ping+ Types.Pong+ Types.Ok+ Types.Err+ Types.Info+ Types.Pub+ Types.Connect+ Types.Sub+ Types.Unsub+ Sid+ Nuid+ WaitGroup+ Pipeline.Broadcasting.Source+ Pipeline.Broadcasting.Transformer+ Pipeline.Broadcasting.Sink+ Pipeline.Broadcasting.API+ Pipeline.Broadcasting+ Pipeline.Streaming.API+ Pipeline.Streaming+ Pipeline.Streaming.Source+ Pipeline.Streaming.Parser+ Pipeline.Streaming.Sink+ Network.Connection+ Network.ConnectionAPI+ Network.Connection.Types+ Network.Connection.Core+ Network.Connection.Tcp+ Network.Connection.Tls+ Queue.API+ Queue.TransactionalQueue+ hs-source-dirs:+ internal+ internal/Types+ internal/Lib+ internal/Plumbing+ internal/Policy++library test-common+ import: shared+ build-depends:+ stm >= 2.5 && < 2.6,+ exposed-modules:+ WaitGroup+ hs-source-dirs:+ internal/Lib++test-suite unit-test+ import: shared+ type: exitcode-stdio-1.0+ hs-source-dirs: test/Unit+ main-is: Spec.hs+ build-depends: + natskell-internal,+ hspec >= 2.7 && < 2.12,+ bytestring,+ network,+ word8,+ aeson,+ stm,+ text,+ random++ other-modules:+ ParserSpec+ MsgSpec+ PingSpec+ PongSpec+ OkSpec+ ErrSpec+ InfoSpec+ PubSpec+ ConnectSpec+ SubSpec+ UnsubSpec+ NKeySpec+ ConnectionSpec+ Fixtures+ StreamingSpec+ SidSpec+ NuidSpec+ WorkerPoolSpec+ WaitGroupSpec++test-suite fuzz-test+ import: shared+ type: exitcode-stdio-1.0+ hs-source-dirs: test/Fuzz+ main-is: Spec.hs+ build-depends: + natskell-internal,+ hspec >= 2.7 && < 2.12,+ bytestring,+ QuickCheck >= 2.14 && < 2.17+ other-modules:+ ParserSpec+ ValidatorsSpec++benchmark parser-bench+ import: shared+ type: exitcode-stdio-1.0+ hs-source-dirs: bench+ main-is: Main.hs+ ghc-options:+ -rtsopts+ -with-rtsopts=-T+ build-depends:+ natskell-internal,+ bytestring++flag impure+ default: False+ manual: True++test-suite integration-tests+ import: shared+ hs-source-dirs: test/Integration+ build-depends:+ natskell,+ test-common,+ port-utils >= 0.1 && < 0.3,+ bytestring,+ word8,+ network,+ stm,+ hspec >= 2.7 && < 2.12+ other-modules:+ ClientSpec+ if flag (impure)+ type: exitcode-stdio-1.0+ main-is: Spec.hs+ else+ type: exitcode-stdio-1.0+ main-is: Skip.hs
+ test/Fuzz/ParserSpec.hs view
@@ -0,0 +1,370 @@+{-# LANGUAGE OverloadedStrings #-}++module ParserSpec (spec) where++import Data.Bifunctor (bimap)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as B8+import qualified Data.ByteString.Lazy as LBS+import Data.Maybe (fromMaybe)+import Data.Word (Word8)+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore)+ , ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Test.Hspec.QuickCheck (modifyMaxSuccess)+import Test.QuickCheck+import Transformers.Transformers (Transformer (transform))+import Types.Err+ ( Err (ErrAuthTimeout, ErrAuthViolation, ErrErr, ErrInvalidProtocol, ErrInvalidSubject, ErrMaxConnsEx, ErrMaxControlLineEx, ErrMaxPayload, ErrPermViolation, ErrRoutePortConn, ErrSlowConsumer, ErrStaleConn, ErrTlsRequired, ErrUnknownOp)+ )+import Types.Info+ ( Info (Info)+ , auth_required+ , client_id+ , connect_urls+ , go+ , headers+ , host+ , ldm+ , max_payload+ , nonce+ , port+ , proto+ , server_id+ , tls_required+ , version+ )+import qualified Types.Msg as Msg+import Types.Msg (Msg (Msg))+import Types.Ok (Ok (Ok))+import Types.Ping (Ping (Ping))+import Types.Pong (Pong (Pong))++spec :: Spec+spec = do+ describe "parser fuzz" $ do+ modifyMaxSuccess (const 5000) .+ it "parses generated valid frames" . property $+ propParsesValidFrames+ modifyMaxSuccess (const 5000) .+ it "requests more input for proper prefixes of valid frames" . property $+ propNeedsMoreForProperPrefixes+ modifyMaxSuccess (const 5000) .+ it "leaves the remaining bytes untouched after one parsed frame" . property $+ propLeavesRemainingBytes+ modifyMaxSuccess (const 5000) .+ it "recovers from generated noisy prefixes" . property $+ propRecoversFromNoisyPrefixes++data ValidFrame = ValidFrame+ { frameBytes :: BS.ByteString+ , parsedValue :: ParsedMessage+ }+ deriving (Eq, Show)++instance Arbitrary ValidFrame where+ arbitrary =+ oneof+ [ pure (ValidFrame (LBS.toStrict (transform Ping)) (ParsedPing Ping))+ , pure (ValidFrame (LBS.toStrict (transform Pong)) (ParsedPong Pong))+ , pure (ValidFrame "+OK\r\n" (ParsedOk Ok))+ , genErrFrame+ , genInfoFrame+ , genMsgFrame+ ]++propParsesValidFrames :: ValidFrame -> Property+propParsesValidFrames validFrame =+ parse parserApi (frameBytes validFrame)+ === Emit (parsedValue validFrame) ""++propNeedsMoreForProperPrefixes :: ValidFrame -> Property+propNeedsMoreForProperPrefixes validFrame =+ not (null candidatePrefixes)+ ==>+ forAll (elements candidatePrefixes) (\prefixBytes -> parse parserApi prefixBytes === NeedMore)+ where+ bytes = frameBytes validFrame+ candidatePrefixes = map (`BS.take` bytes) [1 .. BS.length bytes - 1]++propLeavesRemainingBytes :: ValidFrame -> ValidFrame -> Property+propLeavesRemainingBytes firstFrame secondFrame =+ parse parserApi (frameBytes firstFrame <> frameBytes secondFrame)+ === Emit (parsedValue firstFrame) (frameBytes secondFrame)++propRecoversFromNoisyPrefixes :: ValidFrame -> Property+propRecoversFromNoisyPrefixes validFrame =+ forAll genNoisePrefix $ \noisePrefix ->+ recoverParse (noisePrefix <> frameBytes validFrame)+ === Emit (parsedValue validFrame) ""++recoverParse :: BS.ByteString -> ParseStep ParsedMessage+recoverParse bytes =+ case parse parserApi bytes of+ DropPrefix dropped _+ | dropped <= 0 ->+ error "parser requested a non-positive prefix drop"+ DropPrefix dropped _ ->+ recoverParse (BS.drop dropped bytes)+ result ->+ result++genNoisePrefix :: Gen BS.ByteString+genNoisePrefix =+ fmap BS.pack (listOf1 (elements [33, 48, 63, 90, 95, 122]))++genErrFrame :: Gen ValidFrame+genErrFrame =+ oneof+ [ pure (errFrame "Unknown Protocol Operation" (ErrUnknownOp "Unknown Protocol Operation"))+ , pure (errFrame "Attempted To Connect To Route Port" (ErrRoutePortConn "Attempted To Connect To Route Port"))+ , pure (errFrame "Authorization Violation" (ErrAuthViolation "Authorization Violation"))+ , pure (errFrame "Authorization Timeout" (ErrAuthTimeout "Authorization Timeout"))+ , pure (errFrame "Invalid Client Protocol" (ErrInvalidProtocol "Invalid Client Protocol"))+ , pure (errFrame "Maximum Control Line Exceeded" (ErrMaxControlLineEx "Maximum Control Line Exceeded"))+ , pure (errFrame "Parser Error" (ErrErr "Parser Error"))+ , pure (errFrame "Secure Connection - TLS Required" (ErrTlsRequired "Secure Connection - TLS Required"))+ , pure (errFrame "Stale Connection" (ErrStaleConn "Stale Connection"))+ , pure (errFrame "Maximum Connections Exceeded" (ErrMaxConnsEx "Maximum Connections Exceeded"))+ , pure (errFrame "Slow Consumer" (ErrSlowConsumer "Slow Consumer"))+ , pure (errFrame "Maximum Payload Violation" (ErrMaxPayload "Maximum Payload Violation"))+ , pure (errFrame "Invalid Subject" (ErrInvalidSubject "Invalid Subject"))+ , genPermViolationFrame+ ]+ where+ errFrame reason err =+ ValidFrame ("-ERR '" <> reason <> "'\r\n") (ParsedErr err)++genPermViolationFrame :: Gen ValidFrame+genPermViolationFrame = do+ direction <- elements ["Subscription", "Publish"]+ subjectName <- genSubject+ let reason = "Permissions Violation For " <> direction <> " To " <> subjectName+ pure (ValidFrame ("-ERR '" <> reason <> "'\r\n") (ParsedErr (ErrPermViolation reason)))++genInfoFrame :: Gen ValidFrame+genInfoFrame = do+ infoValue <- genInfo+ pure (ValidFrame (renderInfoFrame infoValue) (ParsedInfo infoValue))++genMsgFrame :: Gen ValidFrame+genMsgFrame = do+ subjectName <- genSubject+ sidValue <- genAlphaNumBytes 1 12+ replySubject <- frequency [(2, pure Nothing), (3, Just <$> genSubject)]+ useHeaders <- arbitrary+ if useHeaders+ then do+ rawHeaders <- genHeaders+ payloadValue <- genPayload 1 128+ let renderedMsg = Msg subjectName sidValue replySubject (Just payloadValue) (Just rawHeaders)+ parsedMsg =+ Msg+ subjectName+ sidValue+ replySubject+ (Just payloadValue)+ (Just (normalizeHeaders rawHeaders))+ pure (ValidFrame (renderHMsgFrame renderedMsg) (ParsedMsg parsedMsg))+ else do+ maybePayload <- frequency [(1, pure Nothing), (4, Just <$> genPayload 1 128)]+ let parsedMsg = Msg subjectName sidValue replySubject maybePayload Nothing+ pure (ValidFrame (renderMsgFrame parsedMsg) (ParsedMsg parsedMsg))++genInfo :: Gen Info+genInfo =+ Info+ <$> genSafeText 1 16+ <*> genSafeText 1 12+ <*> genSafeText 1 8+ <*> genSafeText 1 16+ <*> chooseInt (1, 65535)+ <*> chooseInt (0, 1048576)+ <*> elements [0 .. 5]+ <*> frequency [(2, pure Nothing), (3, Just <$> chooseInt (1, 100000))]+ <*> frequency [(2, pure Nothing), (3, Just <$> genSafeText 1 16)]+ <*> genOptional arbitrary+ <*> genOptional arbitrary+ <*> genOptional (listOf1 (genSafeText 1 24))+ <*> genOptional arbitrary+ <*> genOptional arbitrary++genOptional :: Gen a -> Gen (Maybe a)+genOptional generator =+ frequency [(2, pure Nothing), (3, Just <$> generator)]++genHeaders :: Gen [(BS.ByteString, BS.ByteString)]+genHeaders =+ sized $ \size ->+ do+ pairCount <- chooseInt (1, max 1 (min 4 (size + 1)))+ vectorOf pairCount ((,) <$> genHeaderToken <*> genHeaderValue)++genHeaderToken :: Gen BS.ByteString+genHeaderToken =+ genBytesFromAlphabet 1 12 "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-"++genHeaderValue :: Gen BS.ByteString+genHeaderValue =+ genBytesFromAlphabet 1 32 "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_/.: "++genSubject :: Gen BS.ByteString+genSubject =+ frequency+ [ (1, pure ">")+ , (1, pure "*")+ , (6, genDottedSubject)+ ]++genDottedSubject :: Gen BS.ByteString+genDottedSubject = do+ segmentCount <- chooseInt (1, 4)+ segments <- vectorOf segmentCount genSubjectSegment+ useTailWildcard <- arbitrary+ let baseSubject = B8.intercalate "." segments+ pure+ ( if useTailWildcard+ then baseSubject <> ".>"+ else baseSubject+ )++genSubjectSegment :: Gen BS.ByteString+genSubjectSegment =+ frequency+ [ (1, pure "*")+ , (5, genBytesFromAlphabet 1 12 "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_$/" )+ ]++genSafeText :: Int -> Int -> Gen BS.ByteString+genSafeText minLength maxLength =+ genBytesFromAlphabet minLength maxLength "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_.:/"++genAlphaNumBytes :: Int -> Int -> Gen BS.ByteString+genAlphaNumBytes minLength maxLength =+ genBytesFromAlphabet minLength maxLength "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"++genPayload :: Int -> Int -> Gen BS.ByteString+genPayload minLength maxLength =+ genBytesFromAlphabet minLength maxLength payloadAlphabet+ where+ payloadAlphabet =+ "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789 !?._-/*+\r\n"++genBytesFromAlphabet :: Int -> Int -> BS.ByteString -> Gen BS.ByteString+genBytesFromAlphabet minLength maxLength alphabet = do+ targetLength <- chooseInt (minLength, maxLength)+ bytes <- vectorOf targetLength (elements (BS.unpack alphabet))+ pure (BS.pack bytes)++renderMsgFrame :: Msg -> BS.ByteString+renderMsgFrame msgValue =+ BS.concat+ [ "MSG "+ , Msg.subject msgValue+ , " "+ , Msg.sid msgValue+ , " "+ , maybe "" (<> " ") (Msg.replyTo msgValue)+ , decimalBytes payloadLength+ , "\r\n"+ , payloadBytes+ , "\r\n"+ ]+ where+ payloadBytes = fromMaybe "" (Msg.payload msgValue)+ payloadLength = BS.length payloadBytes++renderHMsgFrame :: Msg -> BS.ByteString+renderHMsgFrame msgValue =+ BS.concat+ [ "HMSG "+ , Msg.subject msgValue+ , " "+ , Msg.sid msgValue+ , " "+ , maybe "" (<> " ") (Msg.replyTo msgValue)+ , decimalBytes (BS.length headerBytes)+ , " "+ , decimalBytes (BS.length headerBytes + BS.length payloadBytes)+ , "\r\n"+ , headerBytes+ , payloadBytes+ , "\r\n"+ ]+ where+ headerBytes = renderHeaderBlock (fromMaybe [] (Msg.headers msgValue))+ payloadBytes = fromMaybe "" (Msg.payload msgValue)++renderHeaderBlock :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString+renderHeaderBlock headerPairs =+ BS.concat ("NATS/1.0\r\n" : foldr appendHeader ["\r\n"] headerPairs)+ where+ appendHeader (headerKey, headerValue) acc =+ [headerKey, ":", headerValue, "\r\n"] ++ acc++normalizeHeaders :: [(BS.ByteString, BS.ByteString)] -> [(BS.ByteString, BS.ByteString)]+normalizeHeaders =+ map (bimap trimHorizontalSpace trimHorizontalSpace)++trimHorizontalSpace :: BS.ByteString -> BS.ByteString+trimHorizontalSpace =+ BS.dropWhile isHorizontalSpaceByte . dropWhileEnd isHorizontalSpaceByte++dropWhileEnd :: (Word8 -> Bool) -> BS.ByteString -> BS.ByteString+dropWhileEnd predicate =+ BS.reverse . BS.dropWhile predicate . BS.reverse++isHorizontalSpaceByte :: Word8 -> Bool+isHorizontalSpaceByte byte =+ byte == 32 || byte == 9++renderInfoFrame :: Info -> BS.ByteString+renderInfoFrame infoValue =+ BS.concat+ [ "INFO {"+ , field "server_id" (quote (server_id infoValue))+ , ","+ , field "version" (quote (version infoValue))+ , ","+ , field "go" (quote (go infoValue))+ , ","+ , field "host" (quote (host infoValue))+ , ","+ , field "port" (decimalBytes (port infoValue))+ , ","+ , field "max_payload" (decimalBytes (max_payload infoValue))+ , ","+ , field "proto" (decimalBytes (proto infoValue))+ , maybeField "client_id" (decimalBytes <$> client_id infoValue)+ , maybeField "nonce" (quote <$> nonce infoValue)+ , maybeField "auth_required" (jsonBool <$> auth_required infoValue)+ , maybeField "tls_required" (jsonBool <$> tls_required infoValue)+ , maybeField "connect_urls" (jsonArray <$> connect_urls infoValue)+ , maybeField "ldm" (jsonBool <$> ldm infoValue)+ , maybeField "headers" (jsonBool <$> headers infoValue)+ , "}\r\n"+ ]++field :: BS.ByteString -> BS.ByteString -> BS.ByteString+field key value = quote key <> ":" <> value++maybeField :: BS.ByteString -> Maybe BS.ByteString -> BS.ByteString+maybeField _ Nothing = ""+maybeField key (Just value) = "," <> field key value++quote :: BS.ByteString -> BS.ByteString+quote value = "\"" <> value <> "\""++jsonBool :: Bool -> BS.ByteString+jsonBool True = "true"+jsonBool False = "false"++jsonArray :: [BS.ByteString] -> BS.ByteString+jsonArray values = "[" <> B8.intercalate "," (map quote values) <> "]"++decimalBytes :: Int -> BS.ByteString+decimalBytes = B8.pack . show
+ test/Fuzz/Spec.hs view
@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}
+ test/Fuzz/ValidatorsSpec.hs view
@@ -0,0 +1,131 @@+{-# LANGUAGE OverloadedStrings #-}++module ValidatorsSpec (spec) where++import qualified Data.ByteString as BS+import Data.Maybe+import Test.Hspec+import Test.Hspec.QuickCheck (modifyMaxSuccess)+import Test.QuickCheck+import Types.Connect+import Types.Pub+import Types.Sub+import Types.Unsub+import Validators.Validators++spec :: Spec+spec = do+ qc++qc = do+ describe "validator" $ do+ modifyMaxSuccess (const 100000) $ do+ it "passes quick check given connect" . property $+ propConnect+ modifyMaxSuccess (const 100000) $ do+ it "passes quick check given pub" . property $+ propPub+ modifyMaxSuccess (const 100000) $ do+ it "passes quick check given sub" . property $+ propSub+ modifyMaxSuccess (const 100000) $ do+ it "passes quick check given unsub" . property $+ propUnsub++propConnect :: Connect -> Bool+propConnect c =+ case validate c of+ Left _ -> not (connectRules c)+ Right _ -> connectRules c++connectRules :: Connect -> Bool+connectRules c+ | auth_token c == Just "" = False+ | user c == Just "" = False+ | pass c == Just "" = False+ | name c == Just "" = False+ | lang c == "" = False+ | version c == "" = False+ | protocol c `notElem` [Nothing, Just 0, Just 1] = False+ | sig c == Just "" = False+ | jwt c == Just "" = False+ | nkey c == Just "" = False+ | otherwise = True+++propPub :: Pub -> Bool+propPub p =+ case validate p of+ Left _ -> not (pubRules p)+ Right _ -> pubRules p++pubRules :: Pub -> Bool+pubRules p+ | Types.Pub.subject p == "" = False+ | Types.Pub.replyTo p == Just "" = False+ | Types.Pub.payload p == Just "" = False+ | headers == Just [] = False+ | isJust headers && Prelude.any (\(k, _) -> k == "") (fromJust headers) = False+ | isJust headers && Prelude.any (\(_, v) -> v == "") (fromJust headers) = False+ | otherwise = True+ where+ headers = Types.Pub.headers p++propSub :: Sub -> Bool+propSub s =+ case validate s of+ Left _ -> not (subRules s)+ Right _ -> subRules s++subRules :: Sub -> Bool+subRules s+ | Types.Sub.subject s == "" = False+ | Types.Sub.queueGroup s == Just "" = False+ | Types.Sub.sid s == "" = False+ | otherwise = True++propUnsub :: Unsub -> Bool+propUnsub u =+ case validate u of+ Left _ -> not (unsubRules u)+ Right _ -> unsubRules u++unsubRules :: Unsub -> Bool+unsubRules u+ | Types.Unsub.sid u == "" = False+ | otherwise = True++instance Arbitrary Connect where+ arbitrary = Connect <$> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary++instance Arbitrary Pub where+ arbitrary = Pub <$> arbitrary+ <*> arbitrary+ <*> arbitrary+ <*> arbitrary++instance Arbitrary Sub where+ arbitrary = Sub <$> arbitrary+ <*> arbitrary+ <*> arbitrary++instance Arbitrary Unsub where+ arbitrary = Unsub <$> arbitrary+ <*> arbitrary++instance Arbitrary BS.ByteString where arbitrary = BS.pack <$> arbitrary
+ test/Integration/ClientSpec.hs view
@@ -0,0 +1,365 @@+{-# LANGUAGE OverloadedStrings #-}++module ClientSpec where++import API+ ( Client (..)+ , MsgView (..)+ , withSubscriptionExpiry+ )+import Client+import Control.Concurrent+import Control.Concurrent.STM+import Control.Exception+import Control.Monad+import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as C+import Data.Foldable (for_)+import Data.IORef+ ( atomicModifyIORef'+ , newIORef+ , readIORef+ , writeIORef+ )+import Data.Word8+import Network.Socket (Socket, accept, listen)+import qualified Network.Socket+import Network.Socket.ByteString (sendAll)+import Network.Socket.Free+import System.Timeout (timeout)+import Test.Hspec+import WaitGroup++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"++tooLongMSG = "MSG a b 5000\r\n" <> BS.replicate 5000 _x <> "\r\n"++headerBlock :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString+headerBlock hs =+ BS.concat ("NATS/1.0\r\n" : foldr appendHeader ["\r\n"] hs)+ where+ appendHeader (key, value) acc = key : ":" : value : "\r\n" : acc++msgFrame :: BS.ByteString -> BS.ByteString -> BS.ByteString -> BS.ByteString+msgFrame subject sid payload =+ BS.concat+ [ "MSG "+ , subject+ , " "+ , sid+ , " "+ , C.pack (show (BS.length payload))+ , "\r\n"+ , payload+ , "\r\n"+ ]++hmsgFrame :: BS.ByteString -> BS.ByteString -> [(BS.ByteString, BS.ByteString)] -> BS.ByteString -> BS.ByteString+hmsgFrame subject sid headers payload =+ let headerBytes = headerBlock headers+ headerLen = BS.length headerBytes+ totalLen = headerLen + BS.length payload+ in BS.concat+ [ "HMSG "+ , subject+ , " "+ , sid+ , " "+ , C.pack (show headerLen)+ , " "+ , C.pack (show totalLen)+ , "\r\n"+ , headerBytes+ , payload+ , "\r\n"+ ]++startClientWith extraOptions = do+ (p, sock) <- openFreePort+ listen sock 1+ tva <- newEmptyTMVarIO+ void . forkIO $ do+ (serv, _) <- accept sock+ atomically $ putTMVar tva serv+ exited <- newEmptyTMVarIO+ tvb <- newEmptyTMVarIO+ void . forkIO $ do+ let configOptions =+ extraOptions+ ++ [ withExitAction (atomically . putTMVar exited)+ , withMinimumLogLevel Debug+ , withConnectionAttempts 1+ , withConnectName "test-client"+ ]+ c <- newClient [("127.0.0.1", p)] configOptions+ atomically $ putTMVar tvb c++ s <- atomically $ takeTMVar tva+ sendAll s defaultINFO+ c <- atomically $ takeTMVar tvb+ return (s, c, sock, exited)++startClient :: IO (Socket, Client, Socket, TMVar ClientExitReason)+startClient =+ startClientWith []++stopClient :: (Socket, Client, Socket, TMVar ClientExitReason) -> IO ()+stopClient (s, c, sock, _) = do+ 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 (Network.Socket.close sock) :: IO (Either SomeException ()))+ void (try (Network.Socket.close s) :: IO (Either SomeException ()))++withClient :: ((Socket, Client, Socket, TMVar ClientExitReason) -> IO()) -> IO ()+withClient action = do+ bracket startClient stopClient action++withClientWith configOptions =+ bracket (startClientWith configOptions) stopClient++spec :: Spec+spec = do+ describe "client integration" $ do+ around withClient $ do+ it "PING waits for PONG" $ \(serv, client, _, _) -> do+ wg <- newWaitGroup 1+ ping client $ done wg+ 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"+ sendAll serv "PONG\r\n"+ result <- timeout 1000000 (takeMVar doneVar)+ case result of+ Nothing -> expectationFailure "flush did not return after PONG"+ Just () -> pure ()+ 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+ result <- atomically $ readTMVar exited+ result `shouldBe` ExitClosedByUser+ it "fatal error results in disconnect" $ \(serv, _, _, exited) -> do+ sendAll serv "-ERR 'Unknown Protocol Operation'\r\n"+ result <- atomically $ readTMVar exited+ case result of+ ExitServerError _ -> return ()+ other -> expectationFailure $ "Unexpected exit reason: " ++ show other+ 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+ sendAll serv "PONG\r\n"+ wait wg+ it "garbled prefix bytes are ignored" $ \(serv, client, _, _) -> do+ wg <- newWaitGroup 1+ ping client $ done wg+ 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+ 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+ 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)+ sendAll serv "MSG SOAK."+ threadDelay 100000+ let headerTail =+ "SUBJECT "+ <> sid+ <> " "+ <> C.pack (show (BS.length payloadValue))+ <> "\r\n"+ chunk = headerTail <> payloadValue <> "\r\n"+ sendAll serv chunk+ result <- timeout 1000000 (takeMVar msgVar)+ case result of+ Nothing -> expectationFailure "message not received"+ Just Nothing -> expectationFailure "received empty message"+ Just (Just msg) -> do+ subject msg `shouldBe` topic+ payload msg `shouldBe` Just 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+ wg <- newWaitGroup 1+ ping client $ done wg+ 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 "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)+ 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)+ 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)+ case result of+ Nothing -> expectationFailure "close blocked after reply"+ Just () -> pure ()+ it "retries when the server disconnects before INFO" $ do+ (p, sock) <- openFreePort+ listen sock 2+ serverConn <- newEmptyTMVarIO+ clientVar <- newEmptyMVar+ void . forkIO $ do+ (firstConn, _) <- accept sock+ Network.Socket.close firstConn+ (secondConn, _) <- accept sock+ sendAll secondConn defaultINFO+ atomically $ putTMVar serverConn secondConn+ void . forkIO $ do+ let configOptions =+ [ withMinimumLogLevel Debug+ , withConnectionAttempts 2+ , withConnectName "test-client"+ ]+ c <- newClient [("127.0.0.1", p)] configOptions+ putMVar clientVar c+ clientResult <- timeout 1000000 (takeMVar clientVar)+ case clientResult of+ Nothing -> expectationFailure "client did not recover after disconnect before INFO"+ Just client -> do+ secondConn <- atomically $ takeTMVar serverConn+ wg <- newWaitGroup 1+ ping client $ done wg+ sendAll secondConn "PONG\r\n"+ wait wg+ close client+ Network.Socket.close secondConn+ Network.Socket.close sock+ around (withClientWith [withBufferLimit (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+ largePayloadSize = 16 * 1024+ smallCount = 100000+ largeCount = 2000+ headerValue = BS.replicate 128 _x+ headerPairs =+ [ ("X-Header-1", headerValue)+ , ("X-Header-2", headerValue)+ , ("X-Header-3", headerValue)+ , ("X-Header-4", headerValue)+ ]+ let smallPayload = BS.replicate smallPayloadSize _x+ let largePayload = BS.replicate largePayloadSize _x+ let expectedTotal = smallCount + largeCount+ let timeoutMicros = 120 * 1000000+ counter <- newIORef 0+ headerCounter <- newIORef 0+ headerChecked <- newIORef False+ errorRef <- newIORef Nothing+ done <- newEmptyMVar+ let recordError err =+ atomicModifyIORef' errorRef $ \current ->+ case current of+ Nothing -> (Just err, ())+ Just _ -> (current, ())+ let handleMsg msg = do+ let payloadLen = maybe 0 BS.length (payload msg)+ case headers msg of+ Just hs -> do+ atomicModifyIORef' headerCounter $ \count -> (count + 1, ())+ when (payloadLen /= largePayloadSize) $+ recordError "unexpected headers on small payload"+ checked <- readIORef headerChecked+ unless checked $ do+ when (hs /= headerPairs) $+ recordError "unexpected header contents"+ writeIORef headerChecked True+ Nothing ->+ when (payloadLen == largePayloadSize) $+ recordError "missing headers on large payload"+ count <- atomicModifyIORef' counter $ \count' ->+ let nextCount = count' + 1+ in (nextCount, nextCount)+ when (count == expectedTotal) $+ void (tryPutMVar done ())+ sid <- subscribe client subject [] (maybe (pure ()) handleMsg)+ let msg = msgFrame subject sid smallPayload+ let hmsg = hmsgFrame subject sid headerPairs largePayload+ let buffer = BS.concat (replicate smallCount msg ++ replicate largeCount hmsg)+ sendAll serv buffer+ result <- timeout timeoutMicros (takeMVar done)+ case result of+ Nothing -> expectationFailure "soak buffer timed out"+ Just () -> pure ()+ errors <- readIORef errorRef+ for_ errors expectationFailure+ headersSeen <- readIORef headerCounter+ headersSeen `shouldBe` largeCount
+ test/Integration/Skip.hs view
@@ -0,0 +1,4 @@+module Main where++main :: IO ()+main = putStrLn "Skipping all system test. Use the flag 'impure' to run them."
+ test/Unit/ConnectSpec.hs view
@@ -0,0 +1,75 @@+{-# LANGUAGE OverloadedStrings #-}++module ConnectSpec (spec) where++import Control.Monad+import Data.Aeson+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Test.Hspec+import Text.Printf+import Transformers.Transformers+import Types.Connect+import Validators.Validators++spec :: Spec+spec = do+ transformationCases+ validationCases++explicitTransformationCases :: [(Connect, BS.ByteString)]+explicitTransformationCases = [+ (+ Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing,+ "CONNECT {\"verbose\": false, \"pedantic\": false, \"tls_required\": false, \"lang\": \"Haskell\", \"version\": \"1\"}"+ ),+ (+ Connect True True True (Just "token") (Just "user") (Just "pass") (Just "name") "Haskell" "1" (Just 3) (Just True) (Just "sig") (Just "jwt") (Just "nkey") (Just True) (Just True),+ "CONNECT {\"verbose\": true, \"pedantic\": true, \"tls_required\": true, \"auth_token\": \"token\", \"user\": \"user\", \"pass\": \"pass\", \"name\": \"name\", \"lang\": \"Haskell\", \"version\": \"1\", \"protocol\": 3, \"echo\": true, \"sig\": \"sig\", \"jwt\": \"jwt\", \"nkey\": \"nkey\", \"no_responders\": true, \"headers\": true}"+ ),+ (+ Connect True True True (Just "token") (Just "user") (Just "pass") (Just "name") "Haskell" "1.0.0" (Just 3) (Just True) (Just "sig") (Just "jwt") (Just "nkey") (Just False) (Just False),+ "CONNECT {\"verbose\": true, \"pedantic\": true, \"tls_required\": true, \"auth_token\": \"token\", \"user\": \"user\", \"pass\": \"pass\", \"name\": \"name\", \"lang\": \"Haskell\", \"version\": \"1.0.0\", \"protocol\": 3, \"echo\": true, \"sig\": \"sig\", \"jwt\": \"jwt\", \"nkey\": \"nkey\", \"no_responders\": false, \"headers\": false}"+ )+ ]++transformationCases = parallel $ do+ describe "CONNECT transformer" $ do+ forM_ explicitTransformationCases $ \(input, want) -> do+ it (printf "correctly transforms %s" (show input)) $ do+ let wantProto = BS.take 8 want+ let wantJSONString = BS.drop 8 want++ let transformed = LBS.toStrict (transform input)+ let gotProto = BS.take 8 transformed+ let gotJSONString = BS.drop 8 transformed++ -- decode both to avoid field ordering issues+ let wantJSON = decode . LBS.fromStrict $ wantJSONString :: Maybe Value+ let gotJSON = decode . LBS.fromStrict $ gotJSONString :: Maybe Value++ gotProto `shouldBe` wantProto+ gotJSON `shouldBe` wantJSON++explicitValidationCases :: [(Connect, Either BS.ByteString ())]+explicitValidationCases = [+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Right ()),+ (Connect False False False (Just "") Nothing Nothing Nothing "Haskell" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Left "explicit empty auth token"),+ (Connect False False False Nothing (Just "") Nothing Nothing "Haskell" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Left "explicit empty user"),+ (Connect False False False Nothing Nothing (Just "") Nothing "Haskell" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Left "explicit empty pass"),+ (Connect False False False Nothing Nothing Nothing (Just "") "Haskell" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Left "explicit empty name"),+ (Connect False False False Nothing Nothing Nothing Nothing "" "1" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Left "explicit empty lang"),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "" Nothing Nothing Nothing Nothing Nothing Nothing Nothing, Left "explicit empty version"),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" (Just 0) Nothing Nothing Nothing Nothing Nothing Nothing, Right ()),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" (Just 1) Nothing Nothing Nothing Nothing Nothing Nothing, Right ()),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" (Just 2) Nothing Nothing Nothing Nothing Nothing Nothing, Left "invalid protocol"),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" Nothing Nothing (Just "") Nothing Nothing Nothing Nothing, Left "explicit empty sig"),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" Nothing Nothing Nothing (Just "") Nothing Nothing Nothing, Left "explicit empty jwt"),+ (Connect False False False Nothing Nothing Nothing Nothing "Haskell" "1" Nothing Nothing Nothing Nothing (Just "") Nothing Nothing, Left "explicit empty nkey")+ ]++validationCases = parallel $ do+ describe "CONNECT validater" $ do+ forM_ explicitValidationCases $ \(input, want) -> do+ it (printf "correctly validates %s" (show input)) $ do+ validate input `shouldBe` want
+ test/Unit/ConnectionSpec.hs view
@@ -0,0 +1,208 @@+{-# LANGUAGE OverloadedStrings #-}++module ConnectionSpec (spec) where++import Auth.None (auth)+import Control.Concurrent+import Control.Concurrent.STM+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Handshake.Nats (performHandshake)+import Lib.Logger+ ( LogLevel (Debug)+ , LoggerConfig (LoggerConfig)+ , newLogContext+ )+import Network.Connection (connectionApi)+import Network.Connection.Core (Transport (..), pointTransport)+import Network.ConnectionAPI+ ( closeReader+ , newConn+ , readData+ , reader+ )+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore, Reject)+ , ParsedMessage (ParsedInfo)+ , ParserAPI (ParserAPI)+ )+import Queue.TransactionalQueue (newQueue)+import State.Store (newClientState, readServerInfo)+import State.Types (ClientConfig (..))+import System.Timeout (timeout)+import Test.Hspec+import Transformers.Transformers (Transformer (transform))+import qualified Types.Connect as Connect+import Types.Info (Info (Info))++spec :: Spec+spec = do+ describe "Connection reader" $ do+ it "unblocks a blocking read when closeReader is called" $ do+ conn <- newConn connectionApi+ started <- newEmptyMVar+ blocker <- (newEmptyMVar :: IO (MVar BS.ByteString))+ let transport = Transport+ { transportRead = \_ -> putMVar started () >> takeMVar blocker+ , transportWrite = \_ -> pure ()+ , transportWriteLazy = \_ -> pure ()+ , transportFlush = pure ()+ , transportClose = pure ()+ , transportUpgrade = Nothing+ }+ pointTransport conn transport+ resultVar <- newEmptyMVar+ _ <- forkIO $ readData (reader connectionApi) conn 1 >>= putMVar resultVar+ _ <- takeMVar started+ closeReader (reader connectionApi) conn+ result <- timeout 1000000 (takeMVar resultVar)+ result `shouldBe` Just (Left "Read operation is blocked")+ describe "Handshake" $ do+ it "accepts an incremental parser backend for the initial INFO frame" $ do+ state <- newTestState+ conn <- newConn connectionApi+ writes <- newTVarIO []+ transport <-+ 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"+ ]+ writes+ pointTransport conn transport++ result <- performHandshake connectionApi incrementalInfoParser state auth conn "127.0.0.1"++ result `shouldBe` Right ()+ readServerInfo state `shouldReturn` Just testInfo+ readTVarIO writes `shouldReturn` [LBS.toStrict (transform testConnect)]++ it "accepts a parser backend that drops an invalid prefix before INFO" $ do+ state <- newTestState+ conn <- newConn connectionApi+ writes <- newTVarIO []+ transport <-+ 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"+ ]+ writes+ pointTransport conn transport++ result <- performHandshake connectionApi dropPrefixInfoParser state auth conn "127.0.0.1"++ result `shouldBe` Right ()+ readServerInfo state `shouldReturn` Just testInfo+ readTVarIO writes `shouldReturn` [LBS.toStrict (transform testConnect)]++newTestState = do+ queue <- newQueue+ ctx <- newLogContext+ conn <- newConn connectionApi+ logger <- newSilentLogger+ newClientState (testConfig logger) queue conn ctx++newSilentLogger :: IO LoggerConfig+newSilentLogger = do+ lock <- newTMVarIO ()+ pure (LoggerConfig Debug (\_ -> pure ()) lock)++testConfig :: LoggerConfig -> ClientConfig+testConfig logger =+ ClientConfig+ { connectionAttempts = 1+ , callbackConcurrency = 1+ , bufferLimit = 4096+ , connectConfig = testConnect+ , loggerConfig = logger+ , tlsCert = Nothing+ , exitAction = const (pure ())+ , connectOptions = []+ }++testConnect :: Connect.Connect+testConnect =+ 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+ }++testInfo :: Info+testInfo =+ Info+ "srv"+ "1.0.0"+ "go1"+ "127.0.0.1"+ 4222+ 1024+ 1+ Nothing+ Nothing+ Nothing+ Nothing+ Nothing+ Nothing+ Nothing++incrementalInfoParser :: ParserAPI ParsedMessage+incrementalInfoParser = ParserAPI parseInfo+ where+ parseInfo bytes+ | bytes == "INF" =+ NeedMore+ | bytes == infoFrame =+ Emit (ParsedInfo testInfo) ""+ | otherwise =+ Reject ("unexpected input: " ++ show bytes)++dropPrefixInfoParser :: ParserAPI ParsedMessage+dropPrefixInfoParser = ParserAPI parseInfo+ where+ parseInfo bytes+ | bytes == "XIN" =+ DropPrefix 1 "invalid prefix"+ | bytes == infoFrame =+ Emit (ParsedInfo testInfo) ""+ | otherwise =+ Reject ("unexpected input: " ++ show bytes)++infoFrame :: BS.ByteString+infoFrame =+ "INFO {\"server_id\":\"srv\",\"version\":\"1.0.0\",\"go\":\"go1\",\"host\":\"127.0.0.1\",\"port\":4222,\"max_payload\":1024,\"proto\":1}\r\n"++newScriptedTransport :: [BS.ByteString] -> TVar [BS.ByteString] -> IO Transport+newScriptedTransport chunks writes = do+ remainingChunks <- newTVarIO chunks+ pure $+ Transport+ { transportRead = \_ ->+ atomically $ do+ remaining <- readTVar remainingChunks+ case remaining of+ nextChunk:rest -> do+ writeTVar remainingChunks rest+ pure nextChunk+ [] ->+ pure BS.empty+ , transportWrite = \bytes ->+ atomically $ modifyTVar' writes (<> [bytes])+ , transportWriteLazy = \bytes ->+ atomically $ modifyTVar' writes (<> [LBS.toStrict bytes])+ , transportFlush = pure ()+ , transportClose = pure ()+ , transportUpgrade = Nothing+ }
+ test/Unit/ErrSpec.hs view
@@ -0,0 +1,54 @@+{-# LANGUAGE OverloadedStrings #-}++module ErrSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import Fixtures+import Parser.API+ ( ParseStep (Emit)+ , ParsedMessage (ParsedErr)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Text.Printf+import Types.Err++explicitCases :: [(BS.ByteString, Err)]+explicitCases = [+ ("-ERR 'Unknown Protocol Operation'\r\n", ErrUnknownOp "Unknown Protocol Operation"),+ ("-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 '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"),+ ("-ERR 'Stale Connection'\r\n", ErrStaleConn "Stale Connection"),+ ("-ERR 'Maximum Connections Exceeded'\r\n", ErrMaxConnsEx "Maximum Connections Exceeded"),+ ("-ERR 'Slow Consumer'\r\n", ErrSlowConsumer"Slow Consumer"),+ ("-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.")+ ]++generatedCases =+ zip+ ((\x n-> foldr BS.append "" ["-ERR 'Permissions Violation For ", x, " To ", n, "'\r\n"]) <$> ["Subscription", "Publish"] <*> invalidSubjectCases)+ ((\x n-> ErrPermViolation $ foldr BS.append "" ["Permissions Violation For ", x, " To ", n]) <$> ["Subscription", "Publish"] <*> invalidSubjectCases)++spec :: Spec+spec = do+ cases++cases = parallel $ do+ describe "generic parser" $ do+ forM_ explicitCases $ \(input, want) ->+ it (printf "correctly parses explicit case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedErr want) ""+ forM_ generatedCases $ \(input, want) ->+ it (printf "correctly parses generated case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedErr want) ""
+ test/Unit/Fixtures.hs view
@@ -0,0 +1,67 @@+{-# LANGUAGE OverloadedStrings #-}++module Fixtures where++import qualified Data.ByteString as BS++versionCases :: [BS.ByteString]+versionCases = ["0.0.0", "v1.0.1", "13.0.0+123"]++boolCases = [True, False]++userCases :: [BS.ByteString]+userCases = ["samisagit", "sam@google.com"]++passCases :: [BS.ByteString]+passCases = ["dsalkj09898(*)(UHJHI&*&*)(910"]++nameCases :: [BS.ByteString]+nameCases = ["natskell-client"]++sigCases :: [BS.ByteString]+sigCases = ["eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQSflKxwRJSMeKKF2QT4fwpMeJf36POk6yJVadQssw5c"]++jwtCases :: [BS.ByteString]+jwtCases = ["eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"]++nonceCases :: [BS.ByteString]+nonceCases = ["nonce-123"]++goVersionCases :: [BS.ByteString]+goVersionCases = ["1.17"]++hostCases :: [BS.ByteString]+hostCases = ["192.168.1.7"]++serverIDCases :: [BS.ByteString]+serverIDCases = ["123"]++portCases :: [Int]+portCases = [4222]+maxPayloadCases :: [Int]+maxPayloadCases = [0, 512, 1024]+protocolCases :: [Int]+protocolCases = [1,2,3]+clientIDCases :: [Int]+clientIDCases = [1, 100]++sidCases :: [BS.ByteString]+sidCases = ["1", "100", "abc", "abc123"]++connectStringCases :: [[BS.ByteString]]+connectStringCases = [[], ["127.0.0.1"], ["127.0.0.1:4222", "0.0.0.0:4222"]]++subjectCases :: [BS.ByteString]+subjectCases = ["FOO", "FOO.BAR", "FOO.BAR.>", ">", "foo.bar", "123.456", "FOO.*.BAR", "FOO.*.BAR.*.>"]++invalidSubjectCases :: [BS.ByteString]+invalidSubjectCases = [" FOO", "FOO>BAR", "FOO ", "F OO", "FOO.**"]++payloadCases :: [Maybe BS.ByteString]+payloadCases = [Just "some payload", Just "some\r\nmulti line payload", Just ".*some payload ** with specials chars > *", Nothing]++headerCases :: [[(BS.ByteString, BS.ByteString)]]+headerCases = [[("header", "abc")], [("header", "abc"), ("HEADER", "123")]]++maybeify :: [a] -> [Maybe a]+maybeify xs = Nothing : fmap Just xs
+ test/Unit/InfoSpec.hs view
@@ -0,0 +1,131 @@+{-# LANGUAGE OverloadedStrings #-}++module InfoSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import qualified Data.Text as T+import Data.Text.Encoding (encodeUtf8)+import Fixtures+import Parser.API+ ( ParseStep (Emit)+ , ParsedMessage (ParsedInfo)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Text.Printf+import Types.Info++spec :: Spec+spec = do+ cases++explicitCases :: [(BS.ByteString, Info)]+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 {\"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 {\"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)+ )+ ]++generatedCases :: [(BS.ByteString, Info)]+generatedCases = zip (map buildProtoInput infos) infos+ where+ infos = Info+ <$> serverIDCases+ <*> versionCases+ <*> goVersionCases+ <*> hostCases+ <*> portCases+ <*> maxPayloadCases+ <*> protocolCases+ <*> maybeify clientIDCases+ <*> maybeify nonceCases+ <*> maybeify boolCases+ <*> maybeify boolCases+ <*> maybeify connectStringCases+ <*> maybeify boolCases+ <*> maybeify boolCases++cases = parallel $ do+ describe "generic parser" $ do+ forM_ explicitCases $ \(input, want) ->+ it (printf "correctly parses explicit case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedInfo want) ""+ forM_ generatedCases $ \(input, want) ->+ it (printf "correctly parses generated case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedInfo want) ""++buildProtoInput :: Info -> BS.ByteString+buildProtoInput m = foldr BS.append "" [+ "INFO",+ " ",+ "{",+ newField "server_id" (Just . quote $ server_id m),+ ",",+ newField "version" (Just . quote $ version m),+ ",",+ newField "go" (Just . quote $ go m),+ ",",+ newField "host" (Just . quote $ host m),+ ",",+ newField "port" (Just . packStr' . show . port $ m),+ ",",+ newField "max_payload" (Just . packStr' . show . max_payload $ m),+ ",",+ newField "proto" (Just . packStr' . show . proto $ m),+ maybeComma (client_id m),+ newField "client_id" (fmap (packStr' . show) . client_id $ m),+ maybeComma (nonce m),+ newField "nonce" (fmap quote . nonce $ m),+ maybeComma (auth_required m),+ newField "auth_required" (fmap boolToJSON . auth_required $ m),+ maybeComma (tls_required m),+ newField "tls_required" (fmap boolToJSON . tls_required $ m),+ maybeComma (connect_urls m),+ newField "connect_urls" (fmap arrayToJSON . connect_urls $ m),+ maybeComma (ldm m),+ newField "ldm" (fmap boolToJSON . ldm $ m),+ maybeComma (headers m),+ newField "headers" (fmap boolToJSON . headers $ m),+ "}",+ "\r\n"+ ]++newField :: BS.ByteString -> Maybe BS.ByteString -> BS.ByteString+newField k v = case v of+ Nothing -> ""+ Just a -> foldr BS.append "" [quote k, ":", a]++quote :: BS.ByteString -> BS.ByteString+quote bs = foldr BS.append "" ["\"", bs, "\""]++packStr' :: String -> BS.ByteString+packStr' = encodeUtf8 . T.pack++boolToJSON :: Bool -> BS.ByteString+boolToJSON b = if b then "true" else "false"++commaSep :: [BS.ByteString] -> BS.ByteString+commaSep [] = ""+commaSep [x] = foldr BS.append "" [quote x, commaSep []]+commaSep (x:xs) = foldr BS.append "" [quote x, ",", commaSep xs]++arrayToJSON :: [BS.ByteString] -> BS.ByteString+arrayToJSON bs = foldr BS.append "" ["[", commaSep bs, "]"]++maybeComma :: Maybe a -> BS.ByteString+maybeComma m = case m of+ Nothing -> ""+ Just _ -> ","
+ test/Unit/MsgSpec.hs view
@@ -0,0 +1,107 @@+{-# LANGUAGE OverloadedStrings #-}++module MsgSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import Data.Maybe+import qualified Data.Text as T+import Data.Text.Encoding (encodeUtf8)+import Fixtures+import Parser.API+ ( ParseStep (Emit)+ , ParsedMessage (ParsedMsg)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Text.Printf+import Types.Msg++spec :: Spec+spec = do+ cases++explicitCases :: [(BS.ByteString, Msg)]+explicitCases = [+ ("MSG FOO 13 BAR 17\r\nsome payload bits\r\n", Msg "FOO" "13" (Just "BAR") (Just "some payload bits") Nothing),+ ("MSG FOO 13 BAR 0\r\n\r\n", Msg "FOO" "13" (Just "BAR") Nothing Nothing),+ ("MSG FOO 13 17\r\nsome payload bits\r\n", Msg "FOO" "13" Nothing (Just"some payload bits") Nothing),+ ("MSG FOO 13 0\r\n\r\n", Msg "FOO" "13" Nothing Nothing Nothing),+ ("MSG FOO 13 20\r\nmulti\r\nline\r\npayload\r\n", Msg "FOO" "13" Nothing (Just "multi\r\nline\r\npayload") Nothing),+ ("MSG FOO.BAR.BAZ 13 IN.*.BOX.> 0\r\n\r\n", Msg "FOO.BAR.BAZ" "13" (Just "IN.*.BOX.>") Nothing Nothing),+ ("HMSG FOO 13 BAR 24 41\r\nNATS/1.0\r\nKEY: VALUE\r\n\r\nsome payload bits\r\n", Msg "FOO" "13" (Just "BAR") (Just "some payload bits") (Just [("KEY", "VALUE")])),+ ("HMSG FOO 13 24 44\r\nNATS/1.0\r\nKEY: VALUE\r\n\r\nmulti\r\nline\r\npayload\r\n", Msg "FOO" "13" Nothing (Just "multi\r\nline\r\npayload") (Just [("KEY", "VALUE")]))+ ]++generatedCases :: [(BS.ByteString, Msg)]+generatedCases = zip (map buildProtoInput msgs) (map expectedMsg msgs)+ where+ msgs = (Msg+ <$> subjectCases+ <*> sidCases+ <*> maybeify subjectCases)+ <*> payloadCases+ <*> maybeify headerCases++expectedMsg :: Msg -> Msg+expectedMsg msg+ | isJust (headers msg) && isNothing (payload msg) = msg { payload = Just "" }+ | otherwise = msg++cases = parallel $ do+ describe "generic parser" $ do+ forM_ explicitCases $ \(input, want) -> do+ it (printf "correctly parses explicit case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedMsg want) ""+ forM_ generatedCases $ \(input, want) -> do+ it (printf "correctly parses generated case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedMsg want) ""++buildProtoInput :: Msg -> BS.ByteString+buildProtoInput m = do+ let headerProto = buildHeaderInput . headers $ m+ foldr BS.append "" [+ if isNothing . headers $ m then "" else "H",+ "MSG",+ " ",+ subject m,+ " ",+ sid m,+ " ",+ collapseNothing (replyTo m) " ",+ if isNothing . headers $ m then "" else packStr' (printf "%v" (BS.length headerProto)),+ if isNothing . headers $ m then "" else " ",+ packStr' (printf "%v" (BS.length headerProto + (payloadLength . payload $ m ))),+ if isNothing . headers $ m then "" else "\r\n",+ headerProto,+ if isJust . headers $ m then "" else "\r\n",+ payloadProto (payload m)+ ]++buildHeaderInput :: Maybe [(BS.ByteString, BS.ByteString)] -> BS.ByteString+buildHeaderInput Nothing = ""+buildHeaderInput (Just xs) = foldr BS.append "" ["NATS/1.0\r\n", headerPairProto xs, "\r\n"]++headerPairProto :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString+headerPairProto [] = ""+headerPairProto ((k, v):xs) = BS.append (foldr BS.append "" [k, ": ", v, "\r\n"]) (headerPairProto xs)++collapseNothing :: Maybe BS.ByteString -> BS.ByteString -> BS.ByteString+collapseNothing mbs suffix = case mbs of+ Just a -> BS.append a suffix+ Nothing -> ""++payloadProto :: Maybe BS.ByteString -> BS.ByteString+payloadProto mbs = case mbs of+ Just payloadValue -> BS.append payloadValue "\r\n"+ Nothing -> "\r\n"++payloadLength :: Maybe BS.ByteString -> Int+payloadLength Nothing = 0+payloadLength (Just xs) = BS.length xs++packStr' :: String -> BS.ByteString+packStr' = encodeUtf8 . T.pack
+ test/Unit/NKeySpec.hs view
@@ -0,0 +1,34 @@+{-# LANGUAGE OverloadedStrings #-}++module NKeySpec (spec) where++import qualified Auth.Jwt as Jwt+import qualified Auth.NKey as NKey+import qualified Data.ByteString as BS+import Test.Hspec++spec :: Spec+spec = do+ describe "parseJwtBundle" $ do+ it "extracts jwt and seed blocks" $ do+ let creds =+ "-----BEGIN NATS USER JWT-----\n\+ \jwt-token\n\+ \------END NATS USER JWT------\n\+ \-----BEGIN USER NKEY SEED-----\n\+ \seed-token\n\+ \------END USER NKEY SEED------"+ Jwt.parseJwtBundle creds `shouldBe` Just (Jwt.JwtBundle "jwt-token" "seed-token")++ it "returns Nothing when blocks are missing" $ do+ Jwt.parseJwtBundle "missing blocks" `shouldBe` Nothing++ describe "signNonceWithSeed" $ do+ it "derives the expected public key from a seed" $ do+ let seed = "SUAHR6JNS2HKJQEAQFHYPOXFXWE4JXBPKUWFX3IMYU72UHOGXT3ZMVFHXI"+ expectedPub = "UAB7EFDOTOBBMPOCK4SXFA62FVZOADQDZOU2W4IUDCGFKJXYVOK3LV7X"+ case NKey.signNonceWithSeed seed "nonce-123" of+ Left err -> expectationFailure err+ Right (pub, sig) -> do+ pub `shouldBe` expectedPub+ BS.length sig `shouldBe` 86
+ test/Unit/NuidSpec.hs view
@@ -0,0 +1,20 @@+{-# LANGUAGE OverloadedStrings #-}++module NuidSpec (spec) where++import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as BC+import Nuid+import System.Random (mkStdGen)+import Test.Hspec++spec :: Spec+spec = describe "NUID generation" $ do+ it "returns base62 tokens with the expected length" $ do+ let nuid0 = newNuid (mkStdGen 1)+ (token, _) = nextNuid nuid0+ BS.length token `shouldBe` 22+ BC.all (`BC.elem` base62Alphabet) token `shouldBe` True++base62Alphabet :: BC.ByteString+base62Alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+ test/Unit/OkSpec.hs view
@@ -0,0 +1,29 @@+{-# LANGUAGE OverloadedStrings #-}++module OkSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import Parser.API+ ( ParseStep (Emit)+ , ParsedMessage (ParsedOk)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Text.Printf+import Types.Ok++spec :: Spec+spec = do+ cases++explicitCases :: [(BS.ByteString, Ok)]+explicitCases = [("+OK\r\n", Ok)]++cases = parallel $ do+ describe "generic parser" $ do+ forM_ explicitCases $ \(input, want) ->+ it (printf "correctly parses explicit case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedOk want) ""
+ test/Unit/ParserSpec.hs view
@@ -0,0 +1,60 @@+{-# LANGUAGE OverloadedStrings #-}++module ParserSpec (spec) where++import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore)+ , ParsedMessage (ParsedMsg, ParsedPing)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import qualified Types.Msg as Msg+import Types.Ping (Ping (Ping))++spec :: Spec+spec = do+ describe "parser api" $ do+ it "parses a valid frame through the API" $ do+ parse parserApi "PING\r\n" `shouldBe` Emit (ParsedPing Ping) ""++ it "suggests pulling more data for truncated input" $ do+ parse parserApi "MSG FOO 1 5\r\nHEL" `shouldBe` NeedMore++ it "suggests dropping invalid prefix bytes" $ do+ case parse parserApi "LOL" of+ DropPrefix n _ ->+ n `shouldSatisfy` (> 0)+ other ->+ expectationFailure ("expected DropPrefix, got " ++ show other)++ describe "message parsing" $ do+ it "accepts tab-delimited fields and non-alphanumeric subjects" $ do+ let input = "MSG foo-bar\t1\t_INBOX.a_b\t5\r\nHELLO\r\n"+ case parse parserApi input of+ Emit parsed rest -> do+ rest `shouldBe` ""+ case parsed of+ ParsedMsg msg' -> do+ Msg.subject msg' `shouldBe` "foo-bar"+ Msg.replyTo msg' `shouldBe` Just "_INBOX.a_b"+ Msg.payload msg' `shouldBe` Just "HELLO"+ other ->+ expectationFailure ("unexpected parse result: " ++ show other)+ other ->+ expectationFailure ("unexpected parse result: " ++ show other)++ it "trims horizontal whitespace in HMSG headers" $ do+ let input =+ "HMSG FOO 13 30 35\r\nNATS/1.0\r\n KEY : VALUE \r\n\r\nHELLO\r\n"+ case parse parserApi input of+ Emit parsed rest -> do+ rest `shouldBe` ""+ case parsed of+ ParsedMsg msg' -> do+ Msg.headers msg' `shouldBe` Just [("KEY", "VALUE")]+ Msg.payload msg' `shouldBe` Just "HELLO"+ other ->+ expectationFailure ("unexpected parse result: " ++ show other)+ other ->+ expectationFailure ("unexpected parse result: " ++ show other)
+ test/Unit/PingSpec.hs view
@@ -0,0 +1,48 @@+{-# LANGUAGE OverloadedStrings #-}++module PingSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Parser.API+ ( ParseStep (Emit)+ , ParsedMessage (ParsedPing)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Text.Printf+import Transformers.Transformers+import Types.Ping+import Validators.Validators++spec :: Spec+spec = do+ parserCases+ transformerCases+ validateCase++explicitParserCases :: [(BS.ByteString, Ping)]+explicitParserCases = [("PING\r\n", Ping)]++explicitTransformerCases :: [(Ping, BS.ByteString)]+explicitTransformerCases = map (\(a,b) -> (b,a)) explicitParserCases++parserCases = parallel $ do+ describe "generic parser" $ do+ forM_ explicitParserCases $ \(input, want) -> do+ it (printf "correctly parses explicit case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedPing want) ""++transformerCases = parallel $ do+ describe "PING transformer" $ do+ forM_ explicitTransformerCases $ \(input, want) -> do+ it (printf"correctly transforms %s" (show input)) $ do+ LBS.toStrict (transform input) `shouldBe` want++validateCase = parallel $ do+ describe "PING validater" $ do+ it "correctly validates PING" $ do+ validate Ping `shouldBe` Right ()
+ test/Unit/PongSpec.hs view
@@ -0,0 +1,45 @@+{-# LANGUAGE OverloadedStrings #-}++module PongSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Parser.API+ ( ParseStep (Emit)+ , ParsedMessage (ParsedPong)+ , parse+ )+import Parser.Attoparsec (parserApi)+import Test.Hspec+import Text.Printf+import Transformers.Transformers+import Types.Pong+import Validators.Validators++spec :: Spec+spec = do+ cases+ validateCase++explicitParserCases :: [(BS.ByteString, Pong)]+explicitParserCases = [("PONG\r\n", Pong)]++explicitTransformerCases :: [(Pong, BS.ByteString)]+explicitTransformerCases = map (\(a,b) -> (b,a)) explicitParserCases++cases = parallel $ do+ describe "generic parser" $ do+ forM_ explicitParserCases $ \(input, want) -> do+ it (printf "correctly parses explicit case %s" (show input)) $ do+ let output = parse parserApi input+ output `shouldBe` Emit (ParsedPong want) ""+ describe "PONG transformer" $ do+ forM_ explicitTransformerCases $ \(input, want) -> do+ it (printf "correctly transforms %s" (show input)) $ do+ LBS.toStrict (transform input) `shouldBe` want++validateCase = parallel $ do+ describe "PONG validater" $ do+ it "correctly validates PONG" $ do+ validate Pong `shouldBe` Right ()
+ test/Unit/PubSpec.hs view
@@ -0,0 +1,56 @@+{-# LANGUAGE OverloadedStrings #-}++module PubSpec (spec) where++import Control.Monad+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Test.Hspec+import Text.Printf+import Transformers.Transformers+import Types.Pub+import Validators.Validators++spec :: Spec+spec = do+ transformerCases+ validateCases++explicitTransformerCases :: [(Pub, BS.ByteString)]+explicitTransformerCases = [+ (Pub "FOO.BAR" Nothing Nothing Nothing, "PUB FOO.BAR 0\r\n\r\n"),+ (Pub "FOO.BAR" Nothing Nothing (Just "Some payload bits"), "PUB FOO.BAR 17\r\nSome payload bits\r\n"),+ (Pub "FOO.BAR" (Just "FOO.BAR.REPLY") Nothing Nothing, "PUB FOO.BAR FOO.BAR.REPLY 0\r\n\r\n"),+ (Pub "FOO.BAR" (Just "FOO.BAR.REPLY") Nothing (Just "Some payload bits"), "PUB FOO.BAR FOO.BAR.REPLY 17\r\nSome payload bits\r\n"),+ (Pub "FOO.BAR" Nothing (Just [("key", "value")]) Nothing, "HPUB FOO.BAR 23 23\r\nNATS/1.0\r\nkey:value\r\n\r\n\r\n"),+ (Pub "FOO.BAR" Nothing (Just [("key", "value"), ("foo", "bar")]) Nothing, "HPUB FOO.BAR 32 32\r\nNATS/1.0\r\nkey:value\r\nfoo:bar\r\n\r\n\r\n"),+ (Pub "FOO.BAR" (Just "FOO.BAR.REPLY") (Just [("key", "value")]) Nothing, "HPUB FOO.BAR FOO.BAR.REPLY 23 23\r\nNATS/1.0\r\nkey:value\r\n\r\n\r\n"),+ (Pub "FOO.BAR" Nothing (Just [("key", "value"), ("foo", "bar")]) (Just "Some payload bits"), "HPUB FOO.BAR 32 49\r\nNATS/1.0\r\nkey:value\r\nfoo:bar\r\n\r\nSome payload bits\r\n"),+ (Pub "FOO.BAR" (Just "FOO.BAR.REPLY") (Just [("key", "value"), ("foo", "bar")]) (Just "Some payload bits"), "HPUB FOO.BAR FOO.BAR.REPLY 32 49\r\nNATS/1.0\r\nkey:value\r\nfoo:bar\r\n\r\nSome payload bits\r\n")+ ]++transformerCases = parallel $ do+ describe "PUB transformer" $ do+ forM_ explicitTransformerCases $ \(input, want) ->+ it (printf "correctly transforms %s" (show input)) $ do+ LBS.toStrict (transform input) `shouldBe` want++explicitValidatorCases :: [(Pub, Either BS.ByteString ())]+explicitValidatorCases = [+ (Pub ">" Nothing Nothing Nothing, Right ()),+ (Pub ">" (Just "SUB.ME") Nothing Nothing, Right ()),+ (Pub ">" Nothing Nothing (Just "payload"), Right ()),+ (Pub ">" (Just "SUB.ME") Nothing (Just "payload"), Right ()),+ (Pub "" Nothing Nothing Nothing, Left "explicit empty subject"),+ (Pub ">" (Just "") Nothing Nothing, Left "explicit empty replyTo"),+ (Pub ">" Nothing Nothing (Just ""), Left "explicit empty payload"),+ (Pub ">" Nothing (Just []) Nothing, Left "explicit empty headers"),+ (Pub ">" Nothing (Just [("", "value")]) Nothing, Left "explicit empty header key"),+ (Pub ">" Nothing (Just [("key", "")]) Nothing, Left "explicit empty header value")+ ]++validateCases = parallel $ do+ describe "PUB validater" $ do+ forM_ explicitValidatorCases $ \(input, want) ->+ it (printf "correctly validates %s" (show input)) $ do+ validate input `shouldBe` want
+ test/Unit/SidSpec.hs view
@@ -0,0 +1,15 @@+{-# LANGUAGE OverloadedStrings #-}++module SidSpec (spec) where++import Sid+import Test.Hspec++spec :: Spec+spec = describe "SID generation" $ do+ it "increments sequentially from the initial counter" $ do+ let (sid1, counter1) = nextSID initialSIDCounter+ (sid2, counter2) = nextSID counter1+ sid1 `shouldBe` "1"+ sid2 `shouldBe` "2"+ counter2 `shouldBe` 2
+ test/Unit/Spec.hs view
@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}
+ test/Unit/StreamingSpec.hs view
@@ -0,0 +1,212 @@+{-# LANGUAGE OverloadedStrings #-}++module StreamingSpec where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Exception+import qualified Control.Monad+import Data.ByteString+ ( ByteString+ , append+ , empty+ , hGetSome+ , hPut+ , head+ , isPrefixOf+ , pack+ , tail+ )+import qualified Data.ByteString.Lazy as LBS+import Data.Char (chr)+import Data.Word8 (isUpper)+import GHC.IO.Handle+ ( BufferMode (NoBuffering)+ , Handle+ , hClose+ , hSetBuffering+ )+import GHC.IO.IOMode+import Lib.Logger+import Network.ConnectionAPI (ReaderAPI (..), WriterAPI (..))+import Network.Socket hiding (Debug, close)+import Parser.API+ ( ParseStep (DropPrefix, Emit, NeedMore)+ , ParserAPI (ParserAPI)+ )+import Pipeline.Broadcasting (broadcastingApi)+import Pipeline.Broadcasting.API (BroadcastingAPI (BroadcastingAPI))+import Pipeline.Streaming (streamingApi)+import Pipeline.Streaming.API (StreamingAPI (StreamingAPI))+import Prelude hiding+ ( head+ , last+ , length+ , null+ , replicate+ , tail+ , take+ )+import Queue.API (QueueItem (..), close, enqueue)+import Queue.TransactionalQueue (newQueue)+import Test.Hspec+import qualified Types.Pub as Pub++defaultLogger' :: IO LoggerConfig+defaultLogger' = do+ lock <- newTMVarIO ()+ pure $ LoggerConfig Debug (putStrLn . renderLogEntry) lock++bufferLimit :: Int+bufferLimit = 4096++newtype CaptureWriter = CaptureWriter (TVar [ByteString])++captureWriterApi :: WriterAPI CaptureWriter+captureWriterApi =+ WriterAPI+ { writeData = \(CaptureWriter output) bytes -> do+ atomically $ modifyTVar' output (<> [bytes])+ pure (Right ())+ , writeDataLazy = \writer bytes ->+ writeData captureWriterApi writer (LBS.toStrict bytes)+ , closeWriter = \_ -> pure ()+ , openWriter = \_ -> pure ()+ }++handleReaderApi :: ReaderAPI Handle+handleReaderApi =+ ReaderAPI+ { readData = \h n -> do+ result <- try (hGetSome h n) :: IO (Either SomeException ByteString)+ case result of+ Left err -> return $ Left (show err)+ Right bytes -> return $ Right bytes+ , closeReader = hClose+ , openReader = \_ -> pure ()+ }++spec :: Spec+spec = do+ describe "Streaming" $ do+ it "reads from source and writes to sink" $ do+ (server, client) <- makeSocketPair+ result <- newTVarIO "" :: IO (TVar ByteString)+ let sink curr = atomically $ modifyTVar' result (`append` curr)+ dl <- defaultLogger'+ ctx <- newLogContext+ let StreamingAPI runStreaming = streamingApi+ (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client singleByteParser sink :: AppM ())+ hPut server "Hello, World"+ atomically $ assertTVarWithRetry result "Hello, World"+ hClose server+ hClose client+ it "continues reading from the handle after reaching end" $ do+ (server, client) <- makeSocketPair+ result <- newTVarIO "" :: IO (TVar ByteString)+ let sink curr = atomically $ modifyTVar' result (`append` curr)+ dl <- defaultLogger'+ ctx <- newLogContext+ let StreamingAPI runStreaming = streamingApi+ (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client singleByteParser sink :: AppM ())+ hPut server "Hello, World"+ atomically $ assertTVarWithRetry result "Hello, World"+ hPut server "Hello, again"+ atomically $ assertTVarWithRetry result "Hello, WorldHello, again"+ hClose server+ hClose client+ it "applies the parser healing" $ do+ (server, client) <- makeSocketPair+ result <- newTVarIO "" :: IO (TVar ByteString)+ let sink curr = atomically $ modifyTVar' result (`append` curr)+ dl <- defaultLogger'+ ctx <- newLogContext+ let StreamingAPI runStreaming = streamingApi+ (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client excludingUpperParser sink :: AppM ())+ hPut server "HELLO WORLD hello, world"+ atomically $ assertTVarWithRetry result " hello, world"+ hClose server+ hClose client+ it "waits for more data" $ do+ (server, client) <- makeSocketPair+ result <- newTVarIO "" :: IO (TVar ByteString)+ let sink curr = atomically $ modifyTVar' result (`append` curr)+ dl <- defaultLogger'+ ctx <- newLogContext+ let StreamingAPI runStreaming = streamingApi+ (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client explicitWordParser sink :: AppM ())+ hPut server "part1"+ threadDelay 100000+ atomically $ ensureTVarIsEmpty result+ hPut server "part2"+ atomically $ assertTVarWithRetry result "part1part2"+ hClose server+ hClose client+ describe "Broadcasting" $ do+ it "drops messages larger than the buffer limit" $ do+ dl <- defaultLogger'+ ctx <- newLogContext+ q <- newQueue+ output <- newTVarIO []+ let pub = Pub.Pub "FOO" Nothing Nothing (Just "0123456789")+ 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` []++ensureTVarIsEmpty :: TVar ByteString -> STM ()+ensureTVarIsEmpty tvar = do+ content <- readTVar tvar+ Control.Monad.when (content /= empty) retry++assertTVarWithRetry :: Eq a => TVar a -> a -> STM ()+assertTVarWithRetry tvar expected = do+ actual <- readTVar tvar+ Control.Monad.unless (actual == expected) retry++singleByteParser :: ParserAPI ByteString+singleByteParser = ParserAPI (\bs -> Emit (pack [head bs]) (tail bs))++excludingUpperParser :: ParserAPI ByteString+excludingUpperParser =+ ParserAPI $ \bs ->+ case isUpper (head bs) of+ False ->+ Emit (pack [head bs]) (tail bs)+ True ->+ DropPrefix 1 [chr . fromIntegral . head $ bs]++explicitWordParser :: ParserAPI ByteString+explicitWordParser = ParserAPI parseWord+ where+ parseWord bs+ | bs == "part1part2" =+ Emit bs empty+ | bs `isPrefixOf` "part1part2" =+ NeedMore+ | otherwise =+ DropPrefix 1 [chr . fromIntegral . head $ bs]++makeSocketPair :: IO (Handle, Handle)+makeSocketPair = do+ serverSock <- socket AF_INET Stream defaultProtocol+ setSocketOption serverSock ReuseAddr 1+ bind serverSock (SockAddrInet 0 (tupleToHostAddress (127,0,0,1)))+ listen serverSock 1+ SockAddrInet port _ <- getSocketName serverSock++ -- initiate client connection+ clientSock <- socket AF_INET Stream defaultProtocol+ connect clientSock (SockAddrInet port (tupleToHostAddress (127,0,0,1)))++ -- accept connection from server side+ (serverConn, _) <- accept serverSock++ -- convert both sides to handles+ serverHandle <- socketToHandle serverConn ReadWriteMode+ clientHandle <- socketToHandle clientSock ReadWriteMode+ hSetBuffering serverHandle NoBuffering+ hSetBuffering clientHandle NoBuffering++ return (serverHandle, clientHandle)
+ test/Unit/SubSpec.hs view
@@ -0,0 +1,44 @@+{-# LANGUAGE OverloadedStrings #-}++module SubSpec (spec) where++import Control.Monad+import Data.ByteString+import qualified Data.ByteString.Lazy as LBS+import Test.Hspec+import Text.Printf+import Transformers.Transformers+import Types.Sub+import Validators.Validators++spec :: Spec+spec = do+ transformerCases+ validateCases++explicitTransformerCases :: [(Sub, ByteString)]+explicitTransformerCases = [+ (Sub "SOME.SUBJ" Nothing "uuid", "SUB SOME.SUBJ uuid\r\n"),+ (Sub "SOME.SUBJ" (Just "QUEUE.GROUP.A") "uuid", "SUB SOME.SUBJ QUEUE.GROUP.A uuid\r\n")+ ]++transformerCases = parallel $ do+ describe "SUB transformer" $ do+ forM_ explicitTransformerCases $ \(input, want) ->+ it (printf "correctly transforms %s" (show input)) $ do+ LBS.toStrict (transform input) `shouldBe` want++explicitValidaterCases :: [(Sub, Either ByteString ())]+explicitValidaterCases = [+ (Sub ">" Nothing "a", Right ()),+ (Sub ">" (Just "1") "a", Right ()),+ (Sub "" Nothing "a", Left "explicit empty subject"),+ (Sub ">" (Just "") "a", Left "explicit empty queue group"),+ (Sub ">" Nothing "", Left "explicit empty sid")+ ]++validateCases = parallel $ do+ describe "SUB validater" $ do+ forM_ explicitValidaterCases $ \(input, want) ->+ it (printf "correctly validates %s" (show input)) $ do+ validate input `shouldBe` want
+ test/Unit/UnsubSpec.hs view
@@ -0,0 +1,43 @@+{-# LANGUAGE OverloadedStrings #-}++module UnsubSpec (spec) where++import Control.Monad+import Data.ByteString+import qualified Data.ByteString.Lazy as LBS+import Test.Hspec+import Text.Printf+import Transformers.Transformers+import Types.Unsub+import Validators.Validators++spec :: Spec+spec = do+ transformerCases+ validateCases++explicitTransformerCases :: [(Unsub, ByteString)]+explicitTransformerCases = [+ (Unsub "uuid" Nothing, "UNSUB uuid\r\n"),+ (Unsub "uuid" (Just 10), "UNSUB uuid 10\r\n")+ ]++transformerCases = parallel $ do+ describe "UNSUB transformer" $ do+ forM_ explicitTransformerCases $ \(input, want) ->+ it (printf "correctly transforms %s" (show input)) $ do+ LBS.toStrict (transform input) `shouldBe` want++explicitValidaterCases :: [(Unsub, Either ByteString ())]+explicitValidaterCases = [+ (Unsub "a" Nothing, Right ()),+ (Unsub "a" (Just 1), Right ()),+ (Unsub "a" (Just 0), Right ()),+ (Unsub "" Nothing, Left "explicit empty sid")+ ]++validateCases = parallel $ do+ describe "UNSUB validater" $ do+ forM_ explicitValidaterCases $ \(input, want) ->+ it (printf "correctly validates %s" (show input)) $ do+ validate input `shouldBe` want
+ test/Unit/WaitGroupSpec.hs view
@@ -0,0 +1,14 @@+module WaitGroupSpec (spec) where++import Control.Exception (throwIO)+import System.Timeout (timeout)+import Test.Hspec+import WaitGroup++spec :: Spec+spec = describe "WaitGroup" $ do+ it "unblocks wait when a forked action throws" $ do+ wg <- newWaitGroup 1+ _ <- forkWaitGroup wg (throwIO (userError "boom"))+ result <- timeout 1000000 (wait wg)+ result `shouldBe` Just ()
+ test/Unit/WorkerPoolSpec.hs view
@@ -0,0 +1,59 @@+module WorkerPoolSpec (spec) where++import Control.Concurrent.STM+ ( atomically+ , check+ , modifyTVar'+ , newEmptyTMVarIO+ , newTQueueIO+ , newTVarIO+ , putTMVar+ , readTMVar+ , readTQueue+ , readTVar+ , readTVarIO+ , writeTQueue+ , writeTVar+ )+import Control.Monad (replicateM_, void, when)+import Lib.WorkerPool (startWorkerPool)+import System.Timeout (timeout)+import Test.Hspec++spec :: Spec+spec = describe "WorkerPool" $ do+ it "limits concurrent jobs to the worker count" $ do+ queue <- newTQueueIO+ gate <- newTQueueIO+ stopVar <- newEmptyTMVarIO+ runningVar <- newTVarIO 0+ maxRunningVar <- newTVarIO 0++ let job = do+ atomically $ do+ running <- readTVar runningVar+ let running' = running + 1+ writeTVar runningVar running'+ maxRunning <- readTVar maxRunningVar+ when (running' > maxRunning) $+ writeTVar maxRunningVar running'+ atomically $ readTQueue gate+ atomically $ modifyTVar' runningVar (subtract 1)++ startWorkerPool 2 queue (void (readTMVar stopVar)) (const (pure ()))+ replicateM_ 4 (atomically $ writeTQueue queue job)++ started <- timeout 2000000 (atomically $ do+ maxRunning <- readTVar maxRunningVar+ check (maxRunning >= 2))+ started `shouldBe` Just ()++ replicateM_ 4 (atomically $ writeTQueue gate ())+ done <- timeout 2000000 (atomically $ do+ running <- readTVar runningVar+ check (running == 0))+ done `shouldBe` Just ()++ atomically $ putTMVar stopVar ()+ maxRunning <- readTVarIO maxRunningVar+ maxRunning `shouldBe` 2